From 4e1f0e2968c77ad391397d785b55af67a3cd23a5 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 2 Sep 2026 12:40:20 +0200 Subject: [PATCH 01/16] feat: rebuild the Atlas reader on the single-file format Atlas 0.16 replaced the directory of per-array files with one write-once container, `data.atlas`, holding every dataset and a footer that describes them all. The reader here was written against the old layout, every call it made is gone, and the crate had been excluded from the workspace since the morsel scan landed. `STORED AS ATLAS` and `read_atlas` failed. The crate is rewritten on the new format and registered again. A `LOCATION` names the container rather than a marker beside it. A collection written before 0.16 is not read at all: its registry is not a marker, so a listing passes over it. There is no compatibility path. One dataset is one unit of work. The format lists a collection at plan time and emits one entry per dataset, those entries go into the shared morsel queue, and the scan sits under the nd spine like netCDF and Zarr. A worker takes the next dataset when it is free and helps drain an open one when none is left, so a collection of a million small datasets and one of four large ones both divide over every core. Level two of the queue follows the chunk shape the writer chose, so one pop reads one stored chunk. A predicate skips whole datasets. The footer records the minimum, the maximum and the null count of every array, so the first scan of a collection pivots those into one index -- one row per dataset, one typed Arrow column per column the predicate names -- and judges every dataset in a single vectorised pass. A million datasets cost one pass rather than a million decisions. A dataset-level attribute is exact in the footer, so a predicate on one prunes too. Every path fails open, and the filter above the scan still decides each row. Three behaviour changes come with the rebuild. A dataset attribute is a column under a leading dot, matching netCDF and Zarr rather than the bare key. A column two datasets type in two families refuses the merge by name instead of silently becoming text; `keep_first` settles it the other way. And collections are crawlable now, because a collection is one file whose extension is its format. Atlas applies no CF decoding: it has a native timestamp type and `atlas create` applies scale, offset and time units before the write. Verified against a collection built by the real `atlas create`: its per-dataset statistics are present, so pruning works on collections built the normal way, and xarray's NaN fills and marker attributes read as documented. --- CHANGELOG.md | 28 + Cargo.lock | 72 + Cargo.toml | 7 +- .../src/components/external-table-dialog.tsx | 2 +- .../beacon-web/src/pages/crawlers.tsx | 1 + beacon-db/beacon-core/Cargo.toml | 4 + .../beacon-core/src/crawler/discovery.rs | 24 +- beacon-db/beacon-core/src/runtime_builder.rs | 12 + beacon-db/beacon-core/tests/atlas.rs | 221 +++ .../beacon-core/tests/schema_functions.rs | 1 + .../python/beacondb/_beacondb.pyi | 1 + beacon-db/beacon-db-py/src/connection.rs | 2 + .../beacon-arrow-atlas/Cargo.toml | 7 +- .../beacon-arrow-atlas/REBUILD_PLAN.md | 906 ++++++++++ .../beacon-arrow-atlas/src/backend.rs | 429 +++-- .../beacon-arrow-atlas/src/compat.rs | 474 +++-- .../beacon-arrow-atlas/src/config.rs | 48 + .../src/datafusion/cache.rs | 151 -- .../src/datafusion/metrics.rs | 52 +- .../beacon-arrow-atlas/src/datafusion/mod.rs | 1553 +++++++++-------- .../src/datafusion/options.rs | 19 +- .../src/datafusion/pruning.rs | 1065 +++++++---- .../src/datafusion/source.rs | 1019 +++++------ .../src/datafusion/statistics.rs | 292 ++++ .../src/datafusion/table_function.rs | 106 +- .../beacon-arrow-atlas/src/lib.rs | 78 +- .../beacon-arrow-atlas/src/reader.rs | 988 ++++++----- .../beacon-arrow-atlas/src/store.rs | 405 +++++ .../beacon-arrow-atlas/src/test_support.rs | 353 ++++ .../beacon-arrow-atlas/src/util.rs | 164 -- .../beacon-nd-array/src/arrow/file_read.rs | 37 + beacon-db/beacon-functions/Cargo.toml | 1 + .../beacon-functions/src/file_formats/mod.rs | 4 + beacon-server/beacon-server-config/Cargo.toml | 1 + beacon-server/beacon-server-config/src/lib.rs | 33 + beacon-server/beacon-server/src/main.rs | 2 + .../beacon-server/src/server/catalog.rs | 1 + beacon-server/beacon-server/src/server/mod.rs | 1 + docs/docs/2.0.0-rc5/cf-decoding.md | 16 +- .../2.0.0-rc5/data-sources/external-tables.md | 12 +- docs/docs/2.0.0-rc5/data-sources/index.md | 4 +- docs/docs/2.0.0-rc5/faq.md | 4 +- docs/docs/2.0.0-rc5/formats/atlas.md | 189 +- docs/docs/2.0.0-rc5/formats/index.md | 14 +- .../docs/2.0.0-rc5/guides/speed-up-queries.md | 2 +- docs/docs/2.0.0-rc5/how-it-works.md | 5 +- docs/docs/2.0.0-rc5/server/configuration.md | 6 +- docs/docs/2.0.0-rc5/server/crawlers.md | 13 +- docs/docs/2.0.0-rc5/server/datasets.md | 2 +- .../2.0.0-rc5/server/performance-tuning.md | 26 +- .../2.0.0-rc5/sql/create-external-table.md | 9 +- docs/docs/2.0.0-rc5/sql/table-functions.md | 18 +- integration-tests/formats/README.md | 1 + integration-tests/formats/test_atlas.py | 241 +++ integration-tests/requirements.txt | 2 + 55 files changed, 6237 insertions(+), 2891 deletions(-) create mode 100644 beacon-db/beacon-core/tests/atlas.rs create mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/REBUILD_PLAN.md create mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs delete mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/cache.rs create mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs create mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/store.rs create mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs delete mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/util.rs create mode 100644 integration-tests/formats/test_atlas.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b16b8e53..ce4dfd17 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,34 @@ tag. Releases before 2.0.0 are recorded in the ### Added +- **Atlas is readable again, on the single-file format.** Atlas 0.16 replaced the directory of + per-array files with one write-once container, `data.atlas`, holding every dataset and a footer + that describes them all. Beacon's reader was written against the old layout and had been + excluded from the build, so `STORED AS ATLAS` and `read_atlas` failed. It is rebuilt on the new + format and registered again. A `LOCATION` now names the container rather than a marker beside + it: `LOCATION 'obs/data.atlas'`, or a glob such as `'obs/**/data.atlas'`. A collection written + before 0.16 is not read at all — rewrite it with `atlas create`. Three behaviour changes come + with the rebuild. A dataset-level attribute is a column under a leading dot, `".platform"`, + matching NetCDF and Zarr instead of the bare key. A scan reads through the shared nd pipeline, + so one dataset is one unit of work and a collection divides over every core. And a column two + datasets type in two families now refuses the merge by name rather than silently becoming text; + `BEACON_TYPE_WIDENING_ON_CONFLICT=keep_first` settles it the other way. Atlas collections are + also crawlable now, because a collection is one file whose extension is its format. + `BEACON_ATLAS_USE_READER_CACHE`, `BEACON_ATLAS_READER_CACHE_SIZE`, `BEACON_ATLAS_USE_PRUNING` + and `BEACON_ATLAS_ENABLE_STATISTICS` configure it, and the same keys work per table through + `OPTIONS`. + +- **A predicate over an Atlas collection skips whole datasets.** The footer records the minimum, + the maximum and the null count of every array, so a collection can be judged before it is read. + The first scan of a collection pivots those statistics into one index — one row per dataset, + one typed Arrow column per column the predicate names — and evaluates the predicate over all of + them in a single vectorised pass. A collection of a million datasets therefore costs one pass + rather than a million decisions, and a dataset ruled out is never opened. A dataset-level + attribute is exact in the footer, so `WHERE ".platform" = 'p3'` prunes on it too. Pruning only + ever removes datasets that hold no matching row: every path falls back to reading everything, + and the filter above the scan still decides each row. `EXPLAIN ANALYZE` reports it as + `atlas_datasets_pruned`, `atlas_index_builds` and `atlas_index_rows`. + - **`BEACON_TYPE_WIDENING_ON_CONFLICT` settles a column that no type holds.** A collection can type one column as a number in one file and as a string in another. No type holds both, so the schema merge refused the whole table and the table answered no query: `Incompatible types for diff --git a/Cargo.lock b/Cargo.lock index b48d5802..4c119418 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -287,6 +287,27 @@ dependencies = [ "password-hash", ] +[[package]] +name = "array-format" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7379c3303a9b9f0693ccd040d85da31e302a601efeafae8a6cb69341e42de41" +dependencies = [ + "bytes", + "futures", + "indexmap 2.14.0", + "lz4_flex 0.11.6", + "moka", + "ndarray 0.17.2", + "object_store 0.13.2", + "rkyv 0.8.10", + "tempfile", + "thiserror 2.0.20", + "tokio", + "zerocopy", + "zstd", +] + [[package]] name = "array-init" version = "2.1.0" @@ -1169,6 +1190,30 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "atlas-rust" +version = "0.16.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f7ba85f9c7eee0f991be08eefbbe9eef18fc9daf0992fde1a4215b6befd2fec" +dependencies = [ + "array-format", + "async-trait", + "bytes", + "chrono", + "futures", + "indexmap 2.14.0", + "ndarray 0.17.2", + "object_store 0.13.2", + "parking_lot", + "rmp-serde", + "serde", + "tempfile", + "thiserror 2.0.20", + "tokio", + "tracing", + "zstd", +] + [[package]] name = "atoi" version = "2.0.0" @@ -1392,6 +1437,29 @@ dependencies = [ "tokio-postgres", ] +[[package]] +name = "beacon-arrow-atlas" +version = "2.0.0-rc.5" +dependencies = [ + "anyhow", + "arrow 58.4.0", + "async-trait", + "atlas-rust", + "beacon-common", + "beacon-datafusion-ext", + "beacon-nd-array", + "chrono", + "datafusion 53.1.0", + "futures", + "indexmap 2.14.0", + "moka", + "ndarray 0.17.2", + "object_store 0.13.2", + "tempfile", + "tokio", + "tracing", +] + [[package]] name = "beacon-arrow-bbf" version = "2.0.0-rc.5" @@ -1723,6 +1791,7 @@ dependencies = [ "async-stream", "async-trait", "base64", + "beacon-arrow-atlas", "beacon-arrow-bbf", "beacon-arrow-csv", "beacon-arrow-geoparquet", @@ -1756,6 +1825,7 @@ dependencies = [ "glob", "iceberg", "iceberg-datafusion", + "ndarray 0.17.2", "num_cpus", "object_store 0.13.2", "parking_lot", @@ -1880,6 +1950,7 @@ name = "beacon-functions" version = "2.0.0-rc.5" dependencies = [ "arrow 58.4.0", + "beacon-arrow-atlas", "beacon-arrow-bbf", "beacon-arrow-csv", "beacon-arrow-geoparquet", @@ -2119,6 +2190,7 @@ name = "beacon-server-config" version = "2.0.0-rc.5" dependencies = [ "base64", + "beacon-arrow-atlas", "beacon-arrow-bbf", "beacon-arrow-hdf5", "beacon-arrow-netcdf", diff --git a/Cargo.toml b/Cargo.toml index f39e5465..72d16def 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,6 @@ [workspace] resolver = "3" -exclude = ["beacon-db/beacon-file-formats/beacon-arrow-atlas"] -members = ["beacon-server/beacon-server", "beacon-server/beacon-server-config", "beacon-server/beacon-mcp", "beacon-db/beacon-auth", "beacon-db/beacon-common", "beacon-db/beacon-core", "beacon-db/beacon-datafusion-ext", "beacon-db/beacon-db-py", "beacon-db/beacon-file-formats/beacon-arrow-bbf", "beacon-db/beacon-file-formats/beacon-arrow-csv", "beacon-db/beacon-file-formats/beacon-arrow-geoparquet", "beacon-db/beacon-file-formats/beacon-arrow-hdf5", "beacon-db/beacon-file-formats/beacon-arrow-ipc", "beacon-db/beacon-file-formats/beacon-arrow-netcdf", "beacon-db/beacon-file-formats/beacon-arrow-odv", "beacon-db/beacon-file-formats/beacon-arrow-parquet", "beacon-db/beacon-file-formats/beacon-arrow-tiff", "beacon-db/beacon-file-formats/beacon-arrow-zarr", "beacon-db/beacon-file-formats/beacon-binary-format", "beacon-db/beacon-file-formats/beacon-binary-format-toolbox", "beacon-db/beacon-file-formats/beacon-delta", "beacon-db/beacon-file-formats/beacon-iceberg", "beacon-db/beacon-file-formats/beacon-icechunk", "beacon-db/beacon-file-formats/beacon-lance", "beacon-db/beacon-file-formats/beacon-nd-array", "beacon-db/beacon-file-formats/beacon-nd-arrow", "beacon-db/beacon-file-stats", "beacon-db/beacon-functions", "beacon-db/beacon-redb-store", "beacon-db/beacon-sql-databases"] +members = ["beacon-server/beacon-server", "beacon-server/beacon-server-config", "beacon-server/beacon-mcp", "beacon-db/beacon-auth", "beacon-db/beacon-common", "beacon-db/beacon-core", "beacon-db/beacon-datafusion-ext", "beacon-db/beacon-db-py", "beacon-db/beacon-file-formats/beacon-arrow-atlas", "beacon-db/beacon-file-formats/beacon-arrow-bbf", "beacon-db/beacon-file-formats/beacon-arrow-csv", "beacon-db/beacon-file-formats/beacon-arrow-geoparquet", "beacon-db/beacon-file-formats/beacon-arrow-hdf5", "beacon-db/beacon-file-formats/beacon-arrow-ipc", "beacon-db/beacon-file-formats/beacon-arrow-netcdf", "beacon-db/beacon-file-formats/beacon-arrow-odv", "beacon-db/beacon-file-formats/beacon-arrow-parquet", "beacon-db/beacon-file-formats/beacon-arrow-tiff", "beacon-db/beacon-file-formats/beacon-arrow-zarr", "beacon-db/beacon-file-formats/beacon-binary-format", "beacon-db/beacon-file-formats/beacon-binary-format-toolbox", "beacon-db/beacon-file-formats/beacon-delta", "beacon-db/beacon-file-formats/beacon-iceberg", "beacon-db/beacon-file-formats/beacon-icechunk", "beacon-db/beacon-file-formats/beacon-lance", "beacon-db/beacon-file-formats/beacon-nd-array", "beacon-db/beacon-file-formats/beacon-nd-arrow", "beacon-db/beacon-file-stats", "beacon-db/beacon-functions", "beacon-db/beacon-redb-store", "beacon-db/beacon-sql-databases"] [workspace.package] @@ -82,6 +81,10 @@ datafusion-spatial = { git = "https://github.com/robinskil/datafusion-spatial.gi datafusion-spatial-kernels = { git = "https://github.com/robinskil/datafusion-spatial.git", rev = "92245ff8e2eaad939425993cff7fa9180d471ec1" } object_store = { version = "0.13.1", features = ["aws", "gcp", "azure", "http"] } +# The atlas container embeds array-format files verbatim, so a change to the +# crate's bytes is a change to the format Beacon reads. Pin it exactly. +atlas-rust = "=0.16.4" + oxcdf = { git = "https://github.com/robinskil/oxcdf.git", version = "0.4.0", features = ["async", "object-store", "ndarray"] } oxcdf-hdf5 = { git = "https://github.com/robinskil/oxcdf.git", version = "0.4.0", features = ["async", "object-store"] } diff --git a/beacon-clients/beacon-web/src/components/external-table-dialog.tsx b/beacon-clients/beacon-web/src/components/external-table-dialog.tsx index 53f8e156..e9a6e6df 100644 --- a/beacon-clients/beacon-web/src/components/external-table-dialog.tsx +++ b/beacon-clients/beacon-web/src/components/external-table-dialog.tsx @@ -55,7 +55,7 @@ const TYPE_HINTS: Record = { NC: "Datasets-store path or glob to netCDF files, e.g. argo/**/*.nc.", HDF5: "Datasets-store path or glob, e.g. data/**/*.h5. NetCDF-4 files are HDF5.", ZARR: "Datasets-store path to a Zarr v3 store (the zarr.json marker).", - ATLAS: "Datasets-store path to an Atlas store (the atlas.json marker).", + ATLAS: "Datasets-store path to an Atlas collection (the data.atlas file).", TIFF: "Datasets-store path or glob to GeoTIFF/COG files.", BBF: "Datasets-store path or glob to Beacon Binary Format files.", ODV: "Datasets-store path or glob to ODV files.", diff --git a/beacon-clients/beacon-web/src/pages/crawlers.tsx b/beacon-clients/beacon-web/src/pages/crawlers.tsx index a5d3d8ca..7233f69a 100644 --- a/beacon-clients/beacon-web/src/pages/crawlers.tsx +++ b/beacon-clients/beacon-web/src/pages/crawlers.tsx @@ -34,6 +34,7 @@ const FORMATS = [ { value: "nc", label: "NetCDF" }, { value: "csv", label: "CSV" }, { value: "zarr", label: "Zarr" }, + { value: "atlas", label: "Atlas" }, { value: "arrow", label: "Arrow" }, { value: "odv", label: "ODV" }, { value: "tiff", label: "GeoTIFF" }, diff --git a/beacon-db/beacon-core/Cargo.toml b/beacon-db/beacon-core/Cargo.toml index 5e43c654..6d6537f8 100644 --- a/beacon-db/beacon-core/Cargo.toml +++ b/beacon-db/beacon-core/Cargo.toml @@ -72,6 +72,7 @@ beacon-arrow-parquet = { path = "../beacon-file-formats/beacon-arrow-parquet" } beacon-arrow-geoparquet = { path = "../beacon-file-formats/beacon-arrow-geoparquet" } beacon-arrow-bbf = { path = "../beacon-file-formats/beacon-arrow-bbf" } beacon-arrow-zarr = { path = "../beacon-file-formats/beacon-arrow-zarr" } +beacon-arrow-atlas = { path = "../beacon-file-formats/beacon-arrow-atlas" } beacon-datafusion-ext = { path = "../beacon-datafusion-ext" } beacon-lance = { path = "../beacon-file-formats/beacon-lance" } beacon-delta = { path = "../beacon-file-formats/beacon-delta" } @@ -82,6 +83,9 @@ beacon-redb-store = { path = "../beacon-redb-store" } beacon-file-stats = { path = "../beacon-file-stats", features = ["datafusion"] } [dev-dependencies] +# Writes the Atlas fixtures the read tests read; the atlas writer takes ndarray +# views, and Beacon itself never writes a collection. +ndarray = { workspace = true } # The GeoJSON filter of the JSON query renders `ST_Within`; its test registers the set. datafusion-spatial = { workspace = true } deltalake = { workspace = true } diff --git a/beacon-db/beacon-core/src/crawler/discovery.rs b/beacon-db/beacon-core/src/crawler/discovery.rs index 55105489..31c436b2 100644 --- a/beacon-db/beacon-core/src/crawler/discovery.rs +++ b/beacon-db/beacon-core/src/crawler/discovery.rs @@ -310,15 +310,33 @@ mod tests { #[test] fn skips_marker_and_overlapping_formats() { - // zarr marker (.json != zarr), atlas marker, geoparquet (.parquet != geoparquet) + // zarr marker (.json != zarr), geoparquet (.parquet != geoparquet) let datasets = vec![ ds("d/foo.zarr/zarr.json", "zarr"), - ds("d/atlas.json", "atlas"), ds("d/g.parquet", "geoparquet"), ]; let (cands, skipped) = group_into_tables(&datasets, &def()); assert!(cands.is_empty()); - assert_eq!(skipped.len(), 3); + assert_eq!(skipped.len(), 2); + } + + /// An Atlas collection is one file, `data.atlas`, so its extension is its + /// format and the crawler can build a table over it. That is the difference + /// from Zarr, whose store is a directory behind a `zarr.json`. + /// + /// Each collection lives in its own directory, and tables group by + /// directory, so each becomes its own table. Several collections in one + /// table is what an external table over a glob is for. + #[test] + fn an_atlas_collection_is_crawlable() { + let datasets = vec![ + ds("d/january/data.atlas", "atlas"), + ds("d/february/data.atlas", "atlas"), + ]; + let (cands, skipped) = group_into_tables(&datasets, &def()); + assert!(skipped.is_empty(), "{skipped:?}"); + assert_eq!(cands.len(), 2, "one per directory"); + assert!(cands.iter().all(|table| table.format == "atlas")); } #[test] diff --git a/beacon-db/beacon-core/src/runtime_builder.rs b/beacon-db/beacon-core/src/runtime_builder.rs index d3e10f38..56e6326e 100644 --- a/beacon-db/beacon-core/src/runtime_builder.rs +++ b/beacon-db/beacon-core/src/runtime_builder.rs @@ -6,6 +6,7 @@ use std::{ use crate::crawler::{new_crawler_manager_handle, CrawlerConfig, CrawlerManager}; use crate::schema_persistence::{init_tables, PersistentSchemaProvider}; +use beacon_arrow_atlas::{AtlasConfig, AtlasFormatFactory, AtlasOptions}; use beacon_arrow_bbf::datafusion::BBFFormatFactory; use beacon_arrow_csv::datafusion::CsvFormatFactory; use beacon_arrow_geoparquet::datafusion::GeoParquetFormatFactory; @@ -98,6 +99,7 @@ pub struct RuntimeBuilder { pub netcdf: NetcdfConfig, pub hdf5: Hdf5Config, pub zarr: ZarrConfig, + pub atlas: AtlasConfig, pub auth_provider: Option>, pub secrets_encryption_key: Option<[u8; 32]>, @@ -236,6 +238,12 @@ impl RuntimeBuilder { self } + /// Replaces the whole Atlas reader configuration. + pub fn with_atlas_config(mut self, atlas: AtlasConfig) -> Self { + self.atlas = atlas; + self + } + pub fn with_auth_provider(mut self, provider: Arc) -> Self { self.auth_provider = Some(provider); self @@ -791,6 +799,10 @@ fn register_file_formats( Arc::new(ArrowFormatFactory), Arc::new(TiffFormatFactory::new(Default::default())), Arc::new(ZarrFormatFactory::new(builder.zarr.clone())), + Arc::new(AtlasFormatFactory::new( + AtlasOptions::default(), + builder.atlas.clone(), + )), Arc::new(BBFFormatFactory::new(Default::default())), Arc::new(GeoParquetFormatFactory::default()), Arc::new(NetCDFFormatFactory::new( diff --git a/beacon-db/beacon-core/tests/atlas.rs b/beacon-db/beacon-core/tests/atlas.rs new file mode 100644 index 00000000..11b95505 --- /dev/null +++ b/beacon-db/beacon-core/tests/atlas.rs @@ -0,0 +1,221 @@ +//! Atlas collections through an assembled runtime. +//! +//! The format crate tests the reader and the scan against its own fixtures. +//! What this covers is the wiring: that a runtime registers the format and its +//! table function, that `STORED AS ATLAS` resolves and survives a restart, and +//! that dataset pruning does not change an answer. +//! +//! An Atlas collection is one write-once file, `data.atlas`, holding many +//! datasets. These tests write real ones with the real writer. + +mod common; + +use std::path::Path; + +use beacon_arrow_atlas::atlas::{AtlasWriter, Attr, WriterConfig}; +use common::{scalar_i64, total_rows, TestRuntime}; +use ndarray::arr1; + +/// Write a collection of `n` datasets at `dir`, named `d0..d{n-1}`. +/// +/// Dataset `i` holds `temperature: Float32[4]` over the range `[10i, 10i + 3]` +/// and the attribute `platform = "p{i}"`, so a threshold predicate has an +/// answer that can be written down. +async fn write_collection(dir: &Path, n: usize) { + std::fs::create_dir_all(dir).expect("create the collection directory"); + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + + for i in 0..n { + let mut dataset = writer + .add_dataset(&format!("d{i}")) + .await + .expect("add a dataset"); + dataset + .define_array::("temperature", vec!["obs".into()], vec![4], None, None) + .await + .expect("define temperature"); + let base = (10 * i) as f32; + dataset + .write_array( + "temperature", + vec![0], + arr1(&[base, base + 1.0, base + 2.0, base + 3.0]) + .into_dyn() + .view(), + ) + .await + .expect("write temperature"); + dataset.set_attribute("platform", Attr::String(format!("p{i}"))); + dataset.finish().await.expect("finish a dataset"); + } + + writer.finish().await.expect("finish the collection"); +} + +/// Every value of `temperature`, sorted, as the query returned them. +async fn temperatures(rt: &TestRuntime, sql: &str) -> Vec { + use arrow::array::Float32Array; + + let batches = rt.sql(sql).await; + let mut values = Vec::new(); + for batch in &batches { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("temperature is f32"); + values.extend(column.iter().flatten()); + } + values +} + +// ── the table function ────────────────────────────────────────────────── + +#[tokio::test(flavor = "multi_thread")] +async fn read_atlas_reads_every_dataset() { + let rt = common::runtime("atlas-read").await; + write_collection(&rt.datasets_dir().join("obs"), 5).await; + + let rows = total_rows( + &rt.sql("SELECT temperature FROM read_atlas('obs/data.atlas')") + .await, + ); + assert_eq!(rows, 20, "five datasets of four rows"); +} + +/// A glob covers several collections in one call. +#[tokio::test(flavor = "multi_thread")] +async fn read_atlas_covers_a_glob_of_collections() { + let rt = common::runtime("atlas-glob").await; + write_collection(&rt.datasets_dir().join("obs/january"), 2).await; + write_collection(&rt.datasets_dir().join("obs/february"), 3).await; + + let rows = total_rows( + &rt.sql("SELECT temperature FROM read_atlas('obs/**/data.atlas')") + .await, + ); + assert_eq!(rows, 20, "five datasets across two collections"); +} + +/// The schema counterpart is registered, and reports the columns without a +/// scan. A dataset attribute is a column under a leading dot. +#[tokio::test(flavor = "multi_thread")] +async fn read_atlas_schema_reports_the_columns() { + let rt = common::runtime("atlas-schema").await; + write_collection(&rt.datasets_dir().join("obs"), 2).await; + + let columns = common::column_strings( + &rt.sql("SELECT column_name FROM read_atlas_schema('obs/data.atlas') ORDER BY column_name") + .await, + 0, + ); + assert_eq!(columns, vec![".platform", "temperature"]); +} + +// ── the external table ────────────────────────────────────────────────── + +#[tokio::test(flavor = "multi_thread")] +async fn an_external_table_reads_a_collection() { + let rt = common::runtime("atlas-external").await; + write_collection(&rt.datasets_dir().join("obs"), 5).await; + + rt.sql("CREATE EXTERNAL TABLE obs STORED AS ATLAS LOCATION 'obs/data.atlas'") + .await; + + assert_eq!(scalar_i64(&rt.sql("SELECT count(*) FROM obs").await), 20); +} + +/// A table in a Beacon-native format has to rebuild at startup. Its definition +/// names the format, and recovery resolves that through the session's registry +/// — so a format registered too late leaves the table silently missing. +#[tokio::test(flavor = "multi_thread")] +async fn an_external_table_survives_a_restart() { + let rt = common::restartable_runtime("atlas-restart", |builder| builder).await; + write_collection(&rt.datasets_dir().join("obs"), 4).await; + + rt.sql("CREATE EXTERNAL TABLE obs STORED AS ATLAS LOCATION 'obs/data.atlas'") + .await; + assert_eq!(scalar_i64(&rt.sql("SELECT count(*) FROM obs").await), 16); + + let rt = rt.restart().await; + assert_eq!( + scalar_i64(&rt.sql("SELECT count(*) FROM obs").await), + 16, + "the table must come back after a restart" + ); +} + +/// The dimensions argument narrows what a read returns, as it does for the +/// other nd formats. +#[tokio::test(flavor = "multi_thread")] +async fn read_atlas_takes_a_dimension_list() { + let rt = common::runtime("atlas-dimensions").await; + write_collection(&rt.datasets_dir().join("obs"), 2).await; + + let rows = total_rows( + &rt.sql("SELECT temperature FROM read_atlas(['obs/data.atlas'], ['obs'])") + .await, + ); + assert_eq!(rows, 8, "`temperature` lives on `obs`, so it survives"); +} + +// ── pruning, through the assembled runtime ────────────────────────────── + +/// A predicate returns the same rows whether or not whole datasets were +/// skipped to find them. +#[tokio::test(flavor = "multi_thread")] +async fn pruning_does_not_change_the_answer() { + let rt = common::runtime_with("atlas-pruning-on", |builder| { + builder.with_atlas_config(beacon_arrow_atlas::AtlasConfig { + use_pruning: true, + ..Default::default() + }) + }) + .await; + let unpruned = common::runtime_with("atlas-pruning-off", |builder| { + builder.with_atlas_config(beacon_arrow_atlas::AtlasConfig { + use_pruning: false, + ..Default::default() + }) + }) + .await; + + for rt in [&rt, &unpruned] { + write_collection(&rt.datasets_dir().join("obs"), 10).await; + } + + for predicate in [ + "temperature > 45", + "temperature < 25", + "temperature > 1000", + "temperature > 45 AND temperature < 75", + ] { + let sql = format!( + "SELECT temperature FROM read_atlas('obs/data.atlas') \ + WHERE {predicate} ORDER BY temperature" + ); + assert_eq!( + temperatures(&rt, &sql).await, + temperatures(&unpruned, &sql).await, + "pruning changed the answer for `{predicate}`" + ); + } +} + +/// An attribute is exact in the footer, so a predicate on one reaches the right +/// dataset without reading the others. +#[tokio::test(flavor = "multi_thread")] +async fn an_attribute_predicate_selects_one_dataset() { + let rt = common::runtime("atlas-attribute").await; + write_collection(&rt.datasets_dir().join("obs"), 6).await; + + let values = temperatures( + &rt, + r#"SELECT temperature FROM read_atlas('obs/data.atlas') + WHERE ".platform" = 'p3' ORDER BY temperature"#, + ) + .await; + assert_eq!(values, vec![30.0, 31.0, 32.0, 33.0]); +} diff --git a/beacon-db/beacon-core/tests/schema_functions.rs b/beacon-db/beacon-core/tests/schema_functions.rs index 371e2bdb..31b5bc0c 100644 --- a/beacon-db/beacon-core/tests/schema_functions.rs +++ b/beacon-db/beacon-core/tests/schema_functions.rs @@ -38,6 +38,7 @@ async fn schema_functions_are_registered_for_each_reader() { "read_parquet_schema", "read_csv_schema", "read_netcdf_schema", + "read_atlas_schema", "read_arrow_schema", "read_delta_schema", "read_iceberg_schema", diff --git a/beacon-db/beacon-db-py/python/beacondb/_beacondb.pyi b/beacon-db/beacon-db-py/python/beacondb/_beacondb.pyi index 380210da..05a891fb 100644 --- a/beacon-db/beacon-db-py/python/beacondb/_beacondb.pyi +++ b/beacon-db/beacon-db-py/python/beacondb/_beacondb.pyi @@ -139,6 +139,7 @@ class Connection: def read_netcdf(self, *args: Any, **kwargs: Any) -> Relation: ... def read_hdf5(self, *args: Any, **kwargs: Any) -> Relation: ... def read_zarr(self, *args: Any, **kwargs: Any) -> Relation: ... + def read_atlas(self, *args: Any, **kwargs: Any) -> Relation: ... def read_delta(self, *args: Any, **kwargs: Any) -> Relation: ... def read_iceberg(self, *args: Any, **kwargs: Any) -> Relation: ... def read_geoparquet(self, *args: Any, **kwargs: Any) -> Relation: ... diff --git a/beacon-db/beacon-db-py/src/connection.rs b/beacon-db/beacon-db-py/src/connection.rs index 83fecd32..c3aed4d8 100644 --- a/beacon-db/beacon-db-py/src/connection.rs +++ b/beacon-db/beacon-db-py/src/connection.rs @@ -756,6 +756,8 @@ const TABLE_FUNCTIONS: &[(&str, &[&str])] = &[ ("read_tiff_schema", &["glob_paths"]), ("read_zarr", &["glob_paths"]), ("read_zarr_schema", &["glob_paths"]), + ("read_atlas", &["glob_paths", "dimensions"]), + ("read_atlas_schema", &["glob_paths", "dimensions"]), ("view_dataset_statistics", &["path"]), ("view_external_table_statistics", &["table_name"]), ("view_statistics_cache", &[]), diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/Cargo.toml b/beacon-db/beacon-file-formats/beacon-arrow-atlas/Cargo.toml index 9125a8d1..b514f673 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/Cargo.toml +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/Cargo.toml @@ -4,7 +4,12 @@ version.workspace = true edition = "2024" rust-version.workspace = true + [dependencies] +# The reader. Pinned exactly in the workspace: the container embeds +# `array-format` files verbatim, so a change to those bytes is a format change. +atlas-rust = { workspace = true } + object_store = { workspace = true } datafusion = { workspace = true } arrow = { workspace = true } @@ -16,9 +21,7 @@ tracing = { workspace = true } tokio = { workspace = true } ndarray = { workspace = true } chrono = { workspace = true } -crossbeam = { workspace = true } moka = { workspace = true } -atlas-rust = "0.14.0" beacon-datafusion-ext = { path = "../../beacon-datafusion-ext" } beacon-nd-array = { path = "../beacon-nd-array" } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/REBUILD_PLAN.md b/beacon-db/beacon-file-formats/beacon-arrow-atlas/REBUILD_PLAN.md new file mode 100644 index 00000000..ea536617 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/REBUILD_PLAN.md @@ -0,0 +1,906 @@ +# Rebuild `beacon-arrow-atlas` on atlas-rust 0.16 + +Status: complete, phases 0 to 5. Date: 2026-09-02. +Scope: this crate, its wiring in core, server, clients, tests and docs. + +| Phase | State | +| --- | --- | +| 0 — prerequisites | done: the crate is a workspace member and pins `atlas-rust = "=0.16.4"` | +| 1 — the reader | done: `config`, `store`, `compat`, `backend`, `reader`, `test_support` | +| 2 — the scan | done: `datafusion/{mod,source,options,metrics,table_function}`, `FileRead::skipped`; 74 atlas + 243 nd-array tests pass, clippy clean | +| 3 — pruning and statistics | done: `pruning`, `statistics`; 105 tests pass, clippy clean | +| 4 — wiring | done: core, functions, server, config, clients, changelog; 8 end-to-end tests in `beacon-core/tests/atlas.rs` | +| 5 — integration and docs | done: `formats/test_atlas.py`, the format page rewritten and 13 other pages updated; the docs site builds | + +Sections 11 to 15 record what each phase settled that this plan had wrong. + +## 1. Why a rebuild + +- This crate targets atlas-rust 0.14. Version 0.14 stores a collection as a + directory: an `atlas.json` marker plus one `.af` file per array. +- atlas-rust 0.16.4 stores a collection as one immutable file. The file is + `/data.atlas`. An optional `/deleted.mask` sits beside it. +- Every atlas call this crate makes is gone in 0.16: `Atlas::open_dataset`, + `Atlas::merged_schema`, `Atlas::pruning_index`, `StoreConfig`, + `MergedSchema`, `ColumnKey`, `StatVal`, `PruningIndex`, `ArraySchema::codec`, + `DatasetSchema::array_attrs` and `DatasetSchema::global_attrs`. +- The workspace excludes the crate since #371 (commit `cebecfda`). Core does + not register it. `STORED AS ATLAS` and `read_atlas` fail today. +- The old crate predates the morsel scan. It dealt dataset names round-robin at + plan time and streamed flat batches outside the nd spine. NetCDF, HDF5, Zarr + and TIFF now read through `MorselSource`, `FileRead`, `NdSourceExec` and + `NdBroadcastExec`. + +Decision: rewrite the crate from scratch. Keep the crate name and path. Keep +three ideas from the old crate: the reader cache, the `atlas_*` metrics and +the per-table `OPTIONS`. + +## 2. The new format, checked against the 0.16.4 source + +Source: `~/.cargo/registry/src/*/atlas-rust-0.16.4/`. + +| Fact | Where | +| --- | --- | +| Container: `ATLS` header, segments back to back, `zstd(msgpack(footer))`, 16-byte trailer. | `src/format/mod.rs` | +| One dataset is one segment. A segment is a complete array-format 0.12.0 file. `SegmentStore` presents it as `seg{ordinal}.af`. | `src/format/segment_store.rs` | +| An open costs one `HEAD`, one 64 KiB tail read, one more read when the footer is larger, and one `GET` of `deleted.mask`. A missing mask is fine. | `src/reader/mod.rs:57-158` | +| The footer holds every dataset name, segment range, schema, attribute value and per-array statistic. Metadata costs no I/O after the open. | `src/format/footer.rs` | +| A schema holds `dtype`, `shape`, `chunk_shape`, `dimension_names` and `fill_value` per array. Datasets with equal schemas share one pool entry. | `src/schema/array.rs` | +| Statistics per array: `min`, `max`, `null_count`, `row_count`. `null_count` counts the elements equal to the fill value. A never-written array has no entry. Lists have no `min` or `max`. | `src/format/footer.rs:305-320` | +| Attributes sit at dataset scope and at array scope. A value is a scalar, a list, or `TimestampNanoseconds`. | `src/schema/attr.rs` | +| Dtypes: `Bool`, `Int8..Int64`, `UInt8..UInt64`, `Float32`, `Float64`, `String`, `Binary`, `TimestampNs`, `List`, `FixedSizeList`. | `array-format/src/dtype.rs` | +| `read_array::` needs `T: ArrayElement`. That is the numeric types, `String`, `Vec` and `TimestampNs`. There is no `bool` and no list element. | `array-format/src/array.rs` | +| `read_array` checks the dtype and the shape against the footer, opens the segment once per handle, and fetches only the chunks the region overlaps. Unwritten cells come from the fill value. | `src/reader/mod.rs:521-591` | +| Atlas merges no schema. Two datasets can declare one array name with two dtypes. | `README.md`, "Types" | +| The collection is immutable. A delete writes ordinals to `deleted.mask`. `list_datasets()` hides them. `dataset(name)` refuses them. | `src/format/mask.rs` | +| `Atlas::dataset(name)` scans the footer linearly. There is no lookup by ordinal. | `src/reader/mod.rs:394-402` | +| Each `Atlas` handle owns one `DeltaCache`: a 256 MiB block budget and a 64 MiB I/O budget. | `src/reader/mod.rs:87-90` | +| `Atlas` and `DatasetView` are `Send + Sync` and implement `Debug`. | `src/lib.rs`, `src/reader/mod.rs` | + +The Python ingest (`atlas create`) shapes the data Beacon meets in practice: + +- One dataset per NetCDF file. The dataset name is the file name. +- xarray opens the file with `mask_and_scale=True`. Scale and offset are + applied. Time is `datetime64[ns]`, so it lands as `TimestampNs`. +- `_FillValue` becomes the fill value. The defaults are `NaN` for floats, + `NaT` (`i64::MIN`) for timestamps, `""` for strings and none for integers. +- Attributes such as `units` and `calendar` stay as plain attributes. +- Python adds `_pyatlas_coords` (a JSON string), `_pyatlas_timedelta` and + `json:`-prefixed string attributes. Beacon shows them as strings. + +Dependency facts: + +- atlas-rust 0.16.4 was published on 2026-09-01. That is today. The repo rule + says: do not take a package that is less than one week old. This is the + user's own crate, so the rule is a prompt to confirm, not a block. +- The local checkout `~/git/atlas` is at 0.15.0. It lacks the 0.16 commits. + Fetch before any upstream work. +- Every transitive dependency fits the workspace: `object_store 0.13`, + `ndarray 0.17`, `rkyv 0.8` (the lock pins 0.8.10), `moka 0.12`, + `lz4_flex 0.11`, `zerocopy 0.8`, `zstd 0.13`, `rmp-serde 1`, `tempfile 3`. + MSRV 1.85 sits under the workspace floor of 1.94. + +## 3. Design + +### 3.1 One dataset is one morsel + +```text +CREATE EXTERNAL TABLE t STORED AS ATLAS LOCATION 'obs/**/data.atlas' + +listing obs/a/data.atlas obs/b/data.atlas markers + │ │ +create_physical_plan │ open (cached) │ list_datasets() + ▼ ▼ +entries [a#d0 a#d1 ... a#dN] [b#d0 ... b#dM] one PartitionedFile per dataset + │ +repartitioned MorselSource queue ──▶ one standing entry per partition + │ +OpenFile::open entry ─▶ Atlas (cached) ─▶ first open of a collection: build its pruning index + │ ─▶ index row of the entry: kept, or skipped + │ ─▶ DatasetView ─▶ AnyDataset (lazy backends) + ▼ ─▶ FileRead::plan (chunk grid, predicate masks) +workers pop chunk ─▶ view.read_array(start, shape) ─▶ nd-encoded batch + │ +plan DataSourceExec ─▶ NdSourceExec ─▶ NdBroadcastExec +``` + +Level 1 of the morsel scan is a dataset. Level 2 is a stored chunk of it. The +backend reports the array's `chunk_shape`, so `FileRead` cuts the dataset on +the chunk grid the writer chose. One pop then reads one stored chunk per +projected array. array-format packs chunks into 8 MiB blocks and caches the +decompressed block, so arrays that share a block cost one fetch. + +### 3.2 Discovery + +- The marker is `data.atlas`. The collection prefix is its parent directory. +- The factory `get_ext()` returns `atlas`. The format `get_ext()` returns + `atlas` too. The listing filter matches by suffix, so `data.atlas` matches + and `deleted.mask` does not. `STORED AS ATLAS` already falls back to the + glob `*.atlas` (see `listing_table_factory_ext.rs`). +- `is_atlas_marker(obj)`: the location is `data.atlas` or ends in + `/data.atlas`. Another `*.atlas` name is skipped with a debug log, because + `Atlas::open` hardcodes the object name. A rename adapter is out of scope. +- `top_level_atlas_markers` keeps one marker per directory. +- `discover_datasets` emits one `DatasetMetadata` per marker with format + `atlas`. The crawler rule "extension equals format" now admits + `data.atlas`, so a crawler can build an atlas table. Update its test. +- `schema_units` uses `units_over_stores`. `schema_options_fingerprint` + returns `SchemaOptions::new("atlas").finish()` when no `read_dimensions` is + set, and `None` otherwise. That matches Zarr. +- Only the 0.16 format is read. A pre-0.16 collection is a directory of `.af` + files with an `atlas.json` registry; nothing here recognizes one, so its + marker is simply not a marker and a listing passes over it. There is no + compatibility path and no migration hint. + +### 3.3 Open and cache + +- `open_collection(store, marker) -> Atlas` calls `Atlas::open(store, prefix)`. +- `AtlasReaderCache` is a `moka::future::Cache>` sized by + `reader_cache_size`. The key is the marker path, its `last_modified`, its + `size`, and a mask stamp. +- The mask stamp is one `HEAD` of `/deleted.mask`: `None` when absent, + else `(last_modified, size, e_tag)`. The container never changes, so this + stamp is the only thing that can retire a cached handle. One `HEAD` per + open-through-cache is the cost. A query pays one or two. +- `get_or_open_atlas(cache: Option<&AtlasReaderCache>, store, marker)` opens + directly when the cache is `None`. +- Memory bound: each handle owns 320 MiB of cache budget. A cache of 32 + handles can hold 10 GiB. Keep the default at 32 and document the bound. See + section 10 for the upstream fix. + +### 3.4 Dataset to `AnyDataset` + +Build one `Arc` per dataset and share it across its backends. +That avoids one linear name lookup per array read. + +Column model. Use the netCDF and Zarr convention: + +| Atlas | Column | +| --- | --- | +| array `a` | `a` | +| array attribute `k` of `a` | `a.k` | +| dataset attribute `k` | `.k` | + +The old crate used the bare key for a dataset attribute. The new convention +avoids a collision between an attribute and an array of one name. Ragged +detection reads `a.sample_dimension`, which is unchanged. + +Type mapping. Put it in `compat.rs` and pin it with tests: + +| Atlas dtype | `NdArrayDataType` | Note | +| --- | --- | --- | +| `Int8..Int64`, `UInt8..UInt64` | the same width | | +| `Float32`, `Float64` | `F32`, `F64` | | +| `String` | `String` | | +| `Binary` | `Binary` | | +| `TimestampNs` | `Timestamp` | `TimestampNanosecond` wraps the same `i64` | +| `Bool` array | skipped | array-format reads no `bool` | +| `List`, `FixedSizeList` | skipped | the nd model has no list | +| `Attr` scalar | rank-0 array | `Bool` attributes are kept | +| `Attr` list | skipped | | + +Log a skip at `debug`, not `warn`. A collection of a million datasets would +flood the log. + +Fill values. `view.array_fill_value(a)` gives a `FillValue`. Convert it with +`::fill_element` and report it from the backend. The nd +engine nulls the cells equal to it. Three consequences to document: + +- A `NaT` timestamp fill reads as null. +- A `""` string fill reads as null. An empty string in the data reads as null + too. That mirrors the Python ingest, which cannot store a null string. +- A `NaN` float fill leaves `NaN` cells as `NaN`, because `NaN != NaN`. That + is the engine's rule for every format. + +No CF decoding. The Python ingest applies scale and offset and decodes time +before the write. The format has a native timestamp type. Beacon therefore +reads every atlas array as stored. Update `cf-decoding.md`, which says the +opposite today. A Rust-written collection with packed integers and CF `units` +reads as integers. Document that. + +Projection. `dataset_from_view(view, projected: Option<&[String]>)` builds a +backend only for a projected column. A wide dataset then pays nothing for the +columns a query never names. + +Dimensions. Apply `resolve_read_dimensions` and `DatasetProjection` as Zarr's +`project_read_dimensions` does. Every array has one dimension name per axis, +so nothing is invented. + +### 3.5 Array backend + +```rust +pub struct AtlasArrayBackend { + view: Arc, + array: String, + shape: Vec, + dimensions: Vec, + chunk_shape: Vec, + fill_value: Option, +} +``` + +- `read_subset(subset)` calls `view.read_array::(array, start, shape)` and + converts the `ArcArray` with `into_owned()`. +- `AtlasElement` bridges `NdArrayType` and `ArrayElement`. Numeric types, + `String` and `Vec` pass through. `TimestampNanosecond` maps element-wise + from `TimestampNs`. +- `chunk_shape()` returns the stored chunk shape. This is what aligns level 2 + of the morsel scan with the file. +- `AttributeBackend` holds one value at rank 0. Copy it from the Zarr crate. + +### 3.6 Schema inference + +Input: the markers of the listing. + +1. Open each collection through the cache. +2. For each live dataset, compute a dedupe key: the address of + `view.schema()` (the interned pool entry) plus the sorted list of + `(attribute key, Attr::dtype())` for the dataset and its arrays. +3. For each new key, build the lazy `AnyDataset`, apply the dimension + narrowing, and derive its Arrow schema with `any_dataset_to_arrow_schema`. + Label it with `marker#dataset`. +4. Merge every labeled schema with `session_widening(state).merge_schemas`. +5. Return an empty schema for a collection with no live dataset. + +A fleet of similar datasets then costs one schema per distinct shape, not one +per dataset. The pass over the footer stays O(datasets) but does no I/O. + +### 3.7 Physical plan + +`AtlasFormat::create_physical_plan`: + +1. `reject_partition_columns("Atlas", &conf)`. A dataset has no path, so no + partition value can come from one. +2. Collect the markers from the file groups. Reduce them to top-level markers. +3. For each marker, open the collection and list its datasets. Build one + `PartitionedFile` per dataset, in `list_datasets()` order. Keep the + marker's `ObjectMeta` verbatim, so the file-statistics pruner still + recognises the store. Put the dataset name and its position in that + listing in `extensions` as `AtlasEntry { dataset: String, position: usize }`. + The position is the row of the dataset in the pruning index (section 3.8). +4. Encode the file schema with `beacon_datafusion_ext::nd::encoded_schema`. +5. Build `AtlasSource` with the read dimensions, the pushed projection, the + cache and the pruning switch. Rebuild the config with the new groups. +6. Wrap the scan: `DataSourceExec` under `NdSourceExec` under + `NdBroadcastExec`. Copy Zarr's `nd_scan_plan`. + +`AtlasSource` implements `FileSource` as `ZarrSource` does: + +- `repartitioned` puts every entry in one `MorselSource` through + `morsel_scan`, unless the scan is ordered or has one partition. +- `try_pushdown_filters` folds the filters into one predicate. +- `try_pushdown_projection` merges projections. +- `create_file_opener` builds an `AtlasOpener` that holds the queue and an + `Arc`. + +`AtlasDatasets` implements `OpenFile`: + +1. Read the `AtlasEntry`. A missing entry is an internal error. +2. Open the collection through the cache. +3. Get the collection's `CandidateFilter` from the `PruneCache`. The first + open of a collection in a scan builds the pruning index (section 3.8). + Every later open reads it. A pruned entry returns `FileRead::skipped()` + and counts in `atlas_datasets_pruned`. +4. Build the `DatasetView`. +5. Build the projected `AnyDataset`, narrow the dimensions, and call + `FileRead::plan` with the projected schema, the batch size, the predicate, + `FilePartitions::none()` and the read metrics. + +`FileRead::skipped()` does not exist yet. Add it to `beacon-nd-array` as a +public constructor for "a file the scan decided not to read": no queue and +`Output::Nothing`. It is four lines. + +The single-partition path reads each entry through the same `open`, then +streams it, as Zarr's opener does. + +The EXPLAIN size of a partition is the marker size times the dataset count. +That is cosmetic. Document it. + +### 3.8 Pruning with a collection index + +A collection can hold millions of datasets. A predicate evaluation per +dataset would cost millions of evaluations. Build one index per collection +instead, and evaluate the predicate once, over every dataset in one +vectorised pass. + +**When.** The first `OpenFile::open` for an entry of a collection builds the +index. Every later open of that collection reads the result. A `PruneCache` +on `AtlasSource` memoises the result per marker path for the life of the +scan. Its `moka::future::Cache::get_with` coalesces the partitions: the first +one builds, the rest await the same future. Each partition's opener holds a +clone of the cache, and the clones share one store. + +**What.** `PruningIndex` holds, in `list_datasets()` order: + +- `names: Vec`. The row order, and the guard of section "How an + entry uses it". +- One `StatColumn` per referenced column: `min: ArrayRef`, `max: ArrayRef`, + `null_count: UInt64Array`, `row_count: UInt64Array`. Every array has length + N, one row per live dataset. + +**How it is built.** + +1. Take the pushed predicate and the logical projected schema. Derive that + schema from the encoded projected schema through + `nd::encoding::nd_value_type`, because the scan schema is nd-encoded. + Build one `PruningPredicate`. A predicate the engine refuses gives + `CandidateFilter::KeepAll`. +2. `collect_columns` names the referenced columns. Resolve each one: + - An array name. Call `atlas.array_stats_by_dataset(name)`. That is one + linear pass over the footer per column, with no view and no name lookup. + Align its `(dataset, stats)` pairs to the rows with a + `HashMap<&str, usize>` built once from `names`. A dataset without an + entry stays unknown in that row. + - An attribute, `.k` or `a.k`. The value is exact, so `min = max = value` + and `null_count = 0`. This prunes `WHERE ".platform" = 'X'` from the + footer alone. Atlas has no bulk attribute accessor. A value needs one + `DatasetView`, and `Atlas::dataset(name)` is a linear scan, so the pass + is quadratic. Build an attribute column only while N is at most + `ATTRIBUTE_INDEX_LIMIT`, 100 000. Above it, leave the column unknown + until the upstream lookup of section 10 lands. + - Anything else. Unknown. The column gets no `StatColumn`, and the + predicate cannot prune on it. +3. Pack each column into typed Arrow arrays in the table type. Use a typed + builder per target type: `Float64Builder`, `Int64Builder`, + `TimestampNanosecondBuilder`, `StringBuilder` and so on. Keep a + `ScalarValue::cast_to` fallback for a type the fast path lacks. Rules per + value: + - Cast the dataset's native dtype to the table type. A cast failure is + null. + - A `NaN` bound is null. `total_cmp` sorts `NaN` last, so a `NaN` max says + nothing about the other values. + - `Bytes` becomes `Utf8` for a `Utf8` column when it is valid UTF-8, and + `Binary` for a `Binary` column. Otherwise null. + - `TimestampNs` becomes `TimestampNanosecond`. + - A missing entry is null for `min` and `max`, and `None` for both counts. + Run the pack on `spawn_blocking`. A million rows is CPU work, not I/O. +4. `AtlasPruningStatistics` implements `PruningStatistics` over the index. + `num_containers` is N. Each accessor returns the column's array, or `None` + for a column the index lacks. +5. `pruning_predicate.prune(&stats)` returns one `bool` per row. Keep it as + `CandidateFilter::Rows { kept: BooleanBuffer, names }`. + +**How an entry uses it.** `AtlasEntry` carries `position`, the dataset's row +in the plan-time listing. `CandidateFilter::keeps(position, name)` reads the +bit at `position` when `names[position] == name`. A mismatch means the +listing changed between plan and open, so the entry is kept. `KeepAll` keeps +every entry. No string is hashed per entry. + +**Counts.** `row_count` is the element count of one array, not the row count +of the broadcast. The predicate uses the counts to decide "every value is +null" and "no value is null". Both hold per array before and after a +broadcast, so the counts are exact for that purpose. + +**Cost.** One footer pass per referenced column, one typed pack, one +vectorised evaluation. For a million datasets and one `Float64` column the +index holds 32 MB and builds in well under a second. A pruned entry then +costs one cache lookup and one empty `FileRead` at its pop. + +**Fail open.** Any error in the build gives `KeepAll`. A row the index cannot +judge stays in. The predicate runs again above the scan, so a kept dataset +that matches nothing costs a read and never a wrong row. + +The pushed predicate also reaches `FileRead::plan`, which prunes chunks on the +coordinate arrays. The two levels compose. + +**Later options, out of scope.** + +- Keep the packed `StatColumn`s on the reader-cache entry, keyed by column + name and table type. The collection is immutable, so a second query pays no + pack. +- When `repartitioned` already holds the predicate, build the index there and + queue only the candidates. That saves one pop and one prefetch task per + pruned dataset. +- Implement `contained()` for attribute columns. Their values are exact, so an + `IN` list prunes too. + +### 3.9 Statistics for the analyzer + +`FileFormat::infer_stats` folds the footer per marker, as Zarr's +`StoreRanges` does: + +- An array column: the lowest `min` and the highest `max` over the live + datasets that hold statistics for it, cast to the table type. Unknown when + any dataset's bound is missing, `NaN`, or fails to cast. +- An attribute column: the same fold over the values. +- A dataset without the column adds nothing. + +It costs no I/O beyond the open. Gate it like Zarr all the same: only +`create_for_analysis` enables it, and `enable_statistics` decides. + +### 3.10 Configuration and `OPTIONS` + +```rust +pub struct AtlasConfig { + pub use_reader_cache: bool, // true + pub reader_cache_size: u64, // 32 + pub use_pruning: bool, // true + pub enable_statistics: bool, // true +} +``` + +| `OPTIONS` key | Env | Effect | +| --- | --- | --- | +| `read_dimensions` | | The dimensions the table reads | +| `use_reader_cache` | `BEACON_ATLAS_USE_READER_CACHE` | Consult the reader cache | +| `use_pruning` | `BEACON_ATLAS_USE_PRUNING` | Prune datasets on footer statistics | +| `enable_statistics` | `BEACON_ATLAS_ENABLE_STATISTICS` | Let the analyzer measure a collection | + +Read the keys with `format_option()`. A key arrives `format.`-prefixed and +lowercased. Reject a bad boolean at `CREATE EXTERNAL TABLE` time. + +### 3.11 Table functions + +`read_atlas(glob_paths)` and `read_atlas(glob_paths, dimensions)`. A path +names a `data.atlas` or a glob such as `**/data.atlas`. The function builds +the format from the session factory with `read_dimensions`, then a +`FastObjectTable`. The `read_atlas_schema` wrapper comes for free. + +### 3.12 Metrics + +Keep `atlas_open_time`, `atlas_prune_time`, `atlas_dataset_build_time`, +`atlas_datasets_scanned` and `atlas_datasets_pruned`. Add +`atlas_index_builds`, the number of pruning indexes a scan built, and +`atlas_index_rows`, their row total. The partition that builds an index +records its build time in `atlas_prune_time`. Register one `ReadMetrics` per +partition, as the other nd formats do. + +### 3.13 Behaviour changes versus the old crate + +1. The marker is `data.atlas`, not `atlas.json`. Every `LOCATION` changes. +2. A dataset attribute column is `.k`, not `k`. +3. A dataset that lacks every projected column contributes no rows. The old + crate null-filled its rows. NetCDF and Zarr already behave this way. +4. The scan goes through the nd spine. `NdBroadcastExec` sits above it, and + the nd projection pushdown rule applies. +5. A partitioned atlas table is refused with a clear error. +6. The pruning index is built from the footer at the first open of a scan. + Nothing is persisted, and nothing is read from disk to build it. +7. A pre-0.16 collection is not read at all. Its `atlas.json` is not a marker, + so a listing passes over it rather than failing a query. +8. A column two datasets type in two families — `String` in one and `Int64` in + another — fails schema inference, and the error names both datasets. The old + crate took atlas's own merge, which made every such column text. Beacon now + settles it the way it settles two files of any other format, and + `BEACON_TYPE_WIDENING_ON_CONFLICT=keep_first` takes the first dataset's type + instead. See section 11. + +## 4. Crate layout + +```text +beacon-arrow-atlas/ + Cargo.toml atlas-rust = "=0.16.4" via the workspace + src/lib.rs crate docs, module list, re-export of `atlas` + src/config.rs AtlasConfig + src/store.rs markers, prefix, open, AtlasReaderCache + src/compat.rs dtype, Attr and FillValue mapping; column names + src/backend.rs AtlasArrayBackend, AttributeBackend, AtlasElement + src/reader.rs dataset_from_view, collection_schema, project_read_dimensions + src/datafusion/mod.rs AtlasFormatFactory, AtlasFormat, nd_scan_plan + src/datafusion/source.rs AtlasSource, AtlasOpener, AtlasDatasets, AtlasEntry + src/datafusion/pruning.rs PruningIndex, StatColumn, CandidateFilter, PruneCache + src/datafusion/statistics.rs the infer_stats fold + src/datafusion/options.rs AtlasOptions + src/datafusion/metrics.rs AtlasScanMetrics + src/datafusion/table_function.rs ReadAtlasFunc + src/test_support.rs #[cfg(test)] fixtures built with AtlasWriter +``` + +Every test lives beside the code it covers, including the end-to-end ones in +`datafusion/mod.rs`. The fixtures are `#[cfg(test)]`, so an integration target +under `tests/` could not reach them. + +Delete every file of the old `src/` first. Nothing in it compiles against +0.16. + +## 5. Work plan + +Environment for every step: + +```bash +export PATH="$HOME/.cargo/bin:$PATH" +source ~/.config/beacon/build-env.sh +``` + +The active toolchain is stable 1.98. CI also builds at 1.94. Use no feature +newer than 1.94. `cargo fmt --check` is not clean repo-wide, so format the +new crate alone. + +### Phase 0: prerequisites + +1. Confirm the dependency rule for atlas-rust 0.16.4 (published today). +2. Add `atlas-rust = "=0.16.4"` to `[workspace.dependencies]`. +3. Remove the `exclude` line from the workspace `Cargo.toml`. Add the crate + to `members`. +4. Run `cargo tree -p beacon-arrow-atlas -e normal | head` after step 5 of + phase 1 to confirm one `rkyv`, one `object_store` and one `ndarray`. + +### Phase 1: the reader + +1. Write `config.rs`, `store.rs`, `compat.rs`, `backend.rs`, `reader.rs`. +2. Write `test_support.rs`. Build collections in a `tempdir` with + `AtlasWriter`: two datasets with attributes and a fill; a widening pair + (`Int16` and `Float32`); an incompatible pair (`String` and `Int64`); a + ranged fleet of `n` datasets; a chunked 2-D grid; a dataset with a list + attribute; an empty collection. A `Bool` or list *array* cannot be a + fixture: no Rust writer can produce one, so that mapping is unit-tested + alone. A deleted dataset is made in the test that wants one, by calling + `delete_dataset` on the open collection. +3. Unit tests: marker recognition, prefix, cache hit and miss on the mask + stamp, every dtype mapping, fill conversion, column names, a full read, a + window read that spans chunks, the timestamp path, the skips. +4. `cargo test -p beacon-arrow-atlas`. + +### Phase 2: the scan + +1. Add `FileRead::skipped()` to `beacon-nd-array`. +2. Write `datafusion/mod.rs`, `source.rs`, `options.rs`, `metrics.rs`, + `table_function.rs`. +3. End-to-end tests in `datafusion/mod.rs`, through `ListingTable` and + `FastObjectTable`: every row once at 1, 4 and 8 partitions; `COUNT(*)`; + projection; the widening cast; the null fill of a missing column; the + incompatible pair as `Utf8`; the deleted dataset absent; `read_dimensions` + narrows the schema; the plan shape `NdBroadcastExec` over `NdSourceExec` + over `DataSourceExec`; a chunk-pruned scan reads fewer encoded batches; + `EXPLAIN` does not open a segment. +4. `cargo test -p beacon-arrow-atlas` and + `cargo test -p beacon-nd-array --lib`. + +### Phase 3: pruning and statistics + +1. Write `pruning.rs` and `statistics.rs`. +2. Tests: the index over the ranged fleet has one row per live dataset in + listing order; `> 45` keeps `d5..d9`; an impossible predicate prunes + everything; a permissive one keeps everything; the mixed-dtype pair casts + before it compares; an attribute predicate prunes from the footer; a `NaN` + bound is null and fails open; an unknown column fails open; a deleted + dataset has no row; a position whose name differs is kept; eight + partitions build the index once (`atlas_index_builds` is 1); results match + with pruning on and off; the metrics report the counts; a synthetic index + of 200 000 rows builds and prunes in one test without a timeout; + `infer_stats` folds a fleet and goes unknown on a mixed dtype. +3. `cargo test -p beacon-arrow-atlas`. + +### Phase 4: wiring + +Section 6 lists the files. Then: + +```bash +cargo clippy --workspace --lib --bins --tests +cargo test --workspace --no-fail-fast --lib --bins --tests +cargo fmt -p beacon-arrow-atlas +``` + +`beacon-datafusion-ext` does not test standalone. Use the workspace run. + +### Phase 5: integration and docs + +1. `integration-tests/formats/test_atlas.py`. It needs `atlas-python`. Add it + to `requirements-optional.txt` and skip when absent. Build a collection + from `test_file.nc` with `atlas.create`, query it, create an external + table, restart, check the table survives. +2. Update the docs of section 8. +3. Add a CHANGELOG entry. + +## 6. Wiring outside the crate + +| File | Change | +| --- | --- | +| `Cargo.toml` (workspace) | Drop the `exclude`. Add the member. Add `atlas-rust = "=0.16.4"`. | +| `beacon-db/beacon-core/Cargo.toml` | Add the crate. | +| `beacon-db/beacon-core/src/runtime_builder.rs` | `pub atlas: AtlasConfig`, `with_atlas_config`, and `AtlasFormatFactory::new(AtlasOptions::default(), builder.atlas.clone())` in `register_file_formats`. | +| `beacon-db/beacon-core/src/crawler/discovery.rs` | Update the marker test: `d/x/data.atlas` is crawlable. | +| `beacon-db/beacon-core/tests/schema_functions.rs` | Add `read_atlas_schema`. | +| `beacon-db/beacon-functions/Cargo.toml`, `src/file_formats/mod.rs` | Register `ReadAtlasFunc`. | +| `beacon-db/beacon-file-formats/beacon-nd-array/src/arrow/file_read.rs` | Add `FileRead::skipped()`. | +| `beacon-db/beacon-db-py/src/connection.rs`, `python/beacondb/_beacondb.pyi` | Add `read_atlas` and `read_atlas_schema`. | +| `beacon-server/beacon-server-config/src/lib.rs` | Re-export `AtlasConfig`. Add the four `BEACON_ATLAS_*` fields. Fill `atlas`. | +| `beacon-server/beacon-server/src/server/mod.rs` | `.with_atlas_config(config.atlas.clone())`. | +| `beacon-server/beacon-server/src/server/catalog.rs` | `"atlas" => "read_atlas"`. | +| `beacon-server/beacon-server/src/main.rs` | Add `atlas` and `array_format` to the quiet log list. | +| `beacon-clients/beacon-web/src/components/external-table-dialog.tsx` | Hint: "the `data.atlas` file". | +| `beacon-clients/beacon-web/src/pages/crawlers.tsx` | Add `{ value: "atlas", label: "Atlas" }`. | +| `integration-tests/formats/test_atlas.py`, `requirements-optional.txt` | New suite. | +| `CHANGELOG.md` | Entry. | + +`beacon-file-stats` and `fast_object` mention Atlas in comments only. They +need no code change. + +## 7. Tests to keep from the old crate + +Port these assertions. Rewrite the fixtures with `AtlasWriter`. + +- `reads_all_datasets_through_datafusion` and the `FastObjectTable` twin. +- `widened_array_dtype_is_cast_from_each_dataset`. +- `missing_column_is_null_filled_per_dataset`. +- The incompatible-dtype case, with its answer corrected: the merge is + refused and the error names both datasets, and `keep_first` resolves it to the + first dataset's type. See section 11. +- `pruning_matches_unpruned_results`, `pruning_on_mixed_dtype_column_end_to_end`, + `pruning_across_many_partitions_is_correct`, + `scan_metrics_report_pruned_and_scanned_counts`, and the `pack_column` + tests of the old `pruning.rs`. The old index came from atlas; the new one is + built here, but the pack and the `PruningStatistics` adapter are the same + shape. +- `partitioned_scan_reads_every_dataset_row`. +- `cache_returns_same_arc_for_identical_marker` and + `cache_reopens_when_last_modified_changes`. Add a mask-change case. +- `discover_datasets_emits_one_entry_per_store`. + +## 8. Docs + +Rewrite `docs/docs/2.0.0-rc5/formats/atlas.md`: the `data.atlas` file, the +mask, the column model, the `OPTIONS` table, footer pruning, the fill rules, +the skips, and "not readable: 0.14 collections". Then update every page that +names `atlas.json`: + +- `formats/index.md` (two rows and the marker note) +- `data-sources/external-tables.md#atlas` +- `server/datasets.md`, `server/configuration.md`, `server/performance-tuning.md` +- `server/crawlers.md` (atlas is crawlable now) +- `cf-decoding.md` ("Zarr and Atlas": Atlas decodes nothing) +- `sql/table-functions.md`, `sql/table-functions-utility.md` +- `guides/speed-up-queries.md`, `guides/query-a-collection.md`, `guides/query-s3.md` +- `faq.md`, `how-it-works.md` + +## 9. Risks and open points + +1. Freshness. atlas-rust 0.16.4 is one day old. A 0.16.5 with an API change + would land on this crate first. Pin exactly and accept. +2. Name lookup. `Atlas::dataset(name)` is O(datasets). A full scan of a + million datasets makes a million lookups. That is quadratic. The pruning + index avoids it for array columns, because `array_stats_by_dataset` is one + linear pass. A dataset that survives pruning still pays one lookup at its + open, and an attribute column in the index pays one per dataset. The + rebuild therefore works as is for a selective query over a large + collection, and for any query over a collection up to the tens of + thousands. Above that the upstream lookup of section 10 is required. +3. Memory. Every cached handle owns 320 MiB of cache budget. Document the + bound. The upstream shared cache removes it. +4. Mask freshness. The reader cache pays one `HEAD` per open. The schema + cache keys on the listed objects, and a `.atlas` listing omits the mask. A + delete can therefore leave a stale merged schema until the next listing + change. The stale schema can only hold an extra column, which reads as + null. Accept. +5. Chunk reads. `read_array` walks every chunk coordinate of the array per + call to find the overlap. A long array read chunk by chunk pays + O(chunks²) comparisons. Acceptable for now. See section 10. +6. Pruning schema. `PruningPredicate` needs the logical column types, and the + scan's projected schema is nd-encoded. Derive the logical schema through + `nd_value_type` and pin it with a test that a pushed `>` prunes. +7. Bool arrays and list values are invisible. The Python ingest refuses bool + by default too. Document. +8. `EXPLAIN` sizes. See section 3.7. + +## 10. Requests to atlas-rust + +Not blockers. Each one lifts a limit above. + +1. `Atlas::dataset_at(ordinal)` and a `HashMap` name index built + at open. Removes the quadratic scan, and lets attribute columns join the + pruning index at any size. +2. `Atlas::attribute_by_dataset(key)` and `Atlas::array_attribute_by_dataset(array, key)`, + the attribute twins of `array_stats_by_dataset`. One footer pass per + attribute column, with no view at all. +3. `Atlas::open_with_cache(store, prefix, Arc)`. Lets Beacon share + one block budget across every open collection. +4. Expose the schema pool, or `Atlas::array_dtypes()`: per array name, the + set of dtypes the live datasets declare. Makes schema inference + O(pool + attributes). +5. In array-format `assemble_nd`, iterate the chunk coordinates that overlap + the slice, not every coordinate of the array. + +## 11. What building it settled + +Phases 0 and 1 answered five questions this plan had guessed at. Each is now +pinned by a test. + +**1. Two families of type are refused, not stringified.** The old crate took +atlas's own merge, where `String` absorbed everything, so a collection whose +datasets typed one array as `String` and `Int64` read back as text. The +schema of a collection now merges through the session's +`ArrowTypeWidening`, exactly as the files of every other format do, and its +default refuses that pair: + +```text +Incompatible types for field 'value': Utf8 in 'sensor#a' vs Int64 in 'sensor#b' +``` + +The label is `{collection}#{dataset}`, so the offending dataset is named +rather than searched for. `BEACON_TYPE_WIDENING_ON_CONFLICT=keep_first` takes +the first dataset's type instead and casts the rest to it. This is a +behaviour change for a collection that holds such a column, and section 3.13 +records it. + +**2. An integer beside a `Float32` widens to `Float64`.** Not to `Float32`, as +the old crate's test asserted. A `Float32` holds no `Int32`, so a rule that +kept it would make the merged type depend on which dataset came first. That is +issue #377, and the session rule is what settles it. + +**3. `list_datasets()` reports write order.** Not sorted order. The pruning +index of section 3.8 keys its rows on that order, and `AtlasEntry::position` +indexes into it, so both are read from the same call at plan time. + +**4. A `Bool` or list array cannot be a fixture.** `array-format` implements +no element type for either, so no Rust writer can produce one and no +collection can hold one. The refusal is unit-tested at the mapping instead. A +list *attribute* is writable, and the reader drops it as designed. + +**5. The marker list must sort by directory, not by path.** A path sort puts +`a/b/data.atlas` before `a/data.atlas`, because `b` sorts under `d`, so a +nested collection would be kept and its parent dropped. The old crate sorted +by path and was correct only by accident: its marker was `atlas.json`, and +`a/atlas.json` sorts before `a/b/atlas.json`. Renaming the marker broke the +assumption. `top_level_atlas_markers` now sorts on the directory. + +Two Phase 0 results worth keeping: + +- The dependency graph is clean. `atlas-rust 0.16.4` brings + `array-format 0.12.0`, and the subtree resolves to one `rkyv 0.8.10` (shared + with `beacon-file-stats`), one `object_store 0.13.2`, one `ndarray 0.17.2` + and one `zstd 0.13.3`. Nothing was duplicated by adding the crate. +- The crate compiles clean at the workspace toolchain with no warnings of its + own, and `cargo check --workspace --lib` still passes. + +## 12. What the scan settled + +Phase 2 answered four more questions, all of them about the DataFusion +surface rather than about atlas. + +**1. The schema-adapter hooks are gone.** DataFusion 53 deprecates +`SchemaAdapterFactory` and gives `FileSource::with_schema_adapter_factory` +and `schema_adapter_factory` defaults that refuse and return `None`. The old +crate implemented both. The nd read adapts its own batches through +`batch_adapter_factory` inside `FileRead::plan`, so the source leaves the +deprecated pair alone and carries no adapter field. + +**2. A `COUNT(*)` still has to read one array.** The projection reaches the +dataset build, so projecting nothing would build nothing, and the read would +have no grid to take a row count from — `FileRead` would then index into an +empty dataset and fail. The opener picks the widest readable array out of the +footer instead and builds that one column. A dataset with no readable array +at all falls back to building its attributes, and contributes the single row +its scalars define. + +**3. A predicate column is built even when the projection omits it.** The +filter that stays above the scan forces its columns into the projection +today, so this changes nothing now. But the chunk pruning inside the read +matches columns by name, and it would stop pruning silently if that ever +stopped holding. + +**4. Partition columns reach a scan through the source's `TableSchema`.** +`FileScanConfigBuilder` has no `with_table_partition_cols`; the config reads +them off the source. Only the refusal test needed to know this — an Atlas +dataset lives inside a container rather than at a path, so +`reject_partition_columns` turns such a table away before a scan is built. + +The legacy format is gone from the crate entirely, by request: there is no +`atlas.json` constant, no detection, and no migration hint. A pre-0.16 +collection's marker is simply not a marker, so a listing passes over it. + +## 13. What pruning settled + +**1. A scan is offered its filters even when it has none.** DataFusion calls +`try_pushdown_filters` with an empty list, and `conjunction` over nothing is +the literal `true`. A source that stores that becomes a source with a +predicate, and this one would then build a pruning index for every scan, +including `SELECT *`. The fix is to treat a predicate that names no column as +no predicate. The other nd formats store the same literal; it costs them +nothing, because their chunk pruning finds no column ranges in it and stops. + +**2. Statistics need a count guard, and pruning does not.** A dataset that +declares an array and never writes it has no footer entry, and its cells read +back as the array's fill — or as zeros, when it declares none. In the pruning +index that dataset gets null bounds, and DataFusion keeps a container whose +statistics are null, so the answer stays right. In the statistics *fold* the +same dataset would simply be skipped, and the range would then exclude values +the collection holds. So `infer_stats` claims a bound only when every live +dataset reported one. A uniform collection — what `atlas create` writes — is +unaffected; a heterogeneous one reports unknown and is read in full. + +**3. The index is what makes attribute pruning affordable at all.** An array +column costs one `array_stats_by_dataset` pass. An attribute has no bulk +accessor, so it costs one `DatasetView` per dataset and each is a linear +scan. That is capped by `ATTRIBUTE_INDEX_LIMIT`, and it is the strongest case +for `Atlas::attribute_by_dataset` upstream (section 10, item 2). + +**4. `infer_stats` measures nothing during a query.** It follows Zarr: a +format built by `create()` reports unknown whatever the option says, and only +`create_for_analysis` turns the fold on. Folding a footer is cheap, but a +listing of thousands of collections would still open every one of them while +planning. + +Two costs are deliberate and worth revisiting if a profile asks: + +- The pack goes through `ScalarValue` per value rather than a typed builder. + It runs on a blocking thread, once per collection per scan. A typed builder + per target type is the optimization, and the packing is one function. +- The index is not kept between queries. The collection is immutable, so the + packed columns could live on the reader-cache entry, keyed by column and + target type. + +## 14. What the wiring settled + +**1. An Atlas collection is crawlable, and Zarr is not.** The crawler's rule is +that a file's extension must equal its format name. A Zarr store is a +directory behind a `zarr.json`, so `json != zarr` and it is skipped. A +collection is one file named `data.atlas`, so `atlas == atlas` and a crawler +builds a table over it. Tables group by directory and each collection has its +own, so a crawl of many collections makes one table each; several in one table +is what an external table over a glob is for. The discovery test that asserted +atlas was skipped is now the test that asserts it is not. + +**2. `beacon-core` needed `ndarray` as a dev-dependency.** Its integration +tests write real collections, and the atlas writer takes `ndarray` views. +Beacon itself never writes a collection, so the dependency belongs in +`dev-dependencies` alone. + +**3. The four settings reach a table two ways.** `BEACON_ATLAS_USE_READER_CACHE`, +`BEACON_ATLAS_READER_CACHE_SIZE`, `BEACON_ATLAS_USE_PRUNING` and +`BEACON_ATLAS_ENABLE_STATISTICS` set the runtime defaults, and every one but +the cache size is overridable per table through `OPTIONS`. An embedded caller +sets them with `RuntimeBuilder::with_atlas_config`. + +### What phase 4 verified, and what it did not + +Verified: `beacon-arrow-atlas` (105 tests), the whole of `beacon-core` +(45 suites, including 8 new end-to-end tests over a real runtime — the table +function, its `_schema` counterpart, a glob, `STORED AS ATLAS`, recovery +across a restart, the dimensions argument, and pruning against an unpruned +control). `cargo check` passes for `beacon-functions`, `beacon-server-config` +and `beacon-server` (lib and bins). Clippy is clean in every touched crate's +own sources. + +Not run: the test suites of `beacon-server-config` and `beacon-server`, and +any build of the web client. The changes there are a config struct, three +one-line wirings and two string edits, but they are unproven. + +### A note on `cargo fmt` + +Do not run `cargo fmt -p ` on the crates this touches. The repository is +not format-clean, so formatting a whole package rewrites files the change never +went near — one pass here reformatted 62 unrelated files across `beacon-core` +and `beacon-functions`, and they had to be reverted one by one. Format the new +crate, whose every file is new, and leave the rest alone. + +## 15. What the ingest settled + +The one thing no Rust test could answer: does the reader handle what `atlas +create` actually writes? It was checked directly — a collection built by +atlas-python 0.16.4 from five netCDF files, opened with this crate. + +**1. Real collections carry the statistics pruning needs.** Each dataset +reported its own `min`, `max`, `null_count` and `row_count` for every array, +with the ranges the source files held. Dataset pruning therefore works on +collections built the normal way, which was the open question behind the whole +design. + +**2. `atlas create` writes `{destination}/data.atlas`.** One dataset per source +file, named after the file *with its suffix* — `d0.nc`, not `d0`. A `LOCATION` +names the container inside that directory. + +**3. xarray's conventions arrive intact and are handled.** A float array gets a +`NaN` fill, because xarray reads with `mask_and_scale=True`; the reader carries +it through, and `NaN` never equals itself, so such a cell reads as `NaN` rather +than as null. Python's own marker attribute lands as the column +`._pyatlas_coords`, beside `.platform` and `temperature.units`. The collection +schema came out as `._pyatlas_coords`, `.platform`, `depth`, `temperature`, +`temperature.units` — exactly the column model this plan describes. + +**4. The reference writers live in `requirements.txt`.** Not +`requirements-optional.txt`, which this plan said: that file is for the Flight +SQL driver alone. The `formats/` convention is one pinned writer per format in +the main file, with each test skipping itself when its writer is absent. +`atlas-python==0.16.4` is pinned there. + +### What phase 5 verified, and what it did not + +Verified: the fixture half of `formats/test_atlas.py` runs against real +atlas-python; the reader handles its output; the test file imports, collects +and skips cleanly when `beacondb` is absent; the documentation site builds, +which is also a dead-link check because VitePress fails a build on one. + +Not run: the body of `formats/test_atlas.py`, which needs the `beacondb` +extension built with maturin. Its SQL and its use of the embedded API follow +`formats/test_zarr.py`, and the same ground is covered in Rust by +`beacon-core/tests/atlas.rs`. diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/backend.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/backend.rs index f8d9efde..53c507c2 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/backend.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/backend.rs @@ -1,131 +1,121 @@ -//! Array backend implementations used by the Atlas reader. +//! The lazy array backends the Atlas reader hands to `beacon-nd-array`. //! -//! Two flavors of lazy [`ArrayBackend`] are provided: -//! - [`AtlasArrayBackend`] — reads an atlas array on demand as its native dtype -//! `T`, reopening a cheap in-memory [`DatasetView`](atlas::DatasetView) per -//! read. -//! - [`AttributeBackend`] — surfaces a scalar attribute value as a rank-0 array. +//! [`AtlasArrayBackend`] reads a region of one atlas array on demand. +//! [`AttributeBackend`] holds one attribute value as a rank-0 array. use std::sync::Arc; +use atlas::{DatasetView, FillValue}; use beacon_nd_array::{ array::{backend::ArrayBackend, subset::ArraySubset}, datatypes::{NdArrayType, TimestampNanosecond}, }; use ndarray::ArrayD; -/// Trait implemented for `T: NdArrayType` values that can be read from an atlas -/// [`DatasetView`](atlas::DatasetView) as a typed `ArrayD`. +/// A Beacon element type that can be read out of an atlas array. /// -/// Atlas's `read_array::` is generic over `atlas::ArrayElement`. Most -/// `NdArrayType` impls are also `ArrayElement` (numeric primitives, `String`, -/// `Vec`), but [`TimestampNanosecond`] is a layout-compatible newtype over -/// `i64` that needs a thin element-wise conversion through -/// [`atlas::TimestampNs`]. This trait hides the difference behind a single -/// `read` entry point so [`AtlasArrayBackend`] stays generic. +/// Atlas reads through [`atlas::ArrayElement`], and Beacon's ND model through +/// [`NdArrayType`]. The two agree on the numeric types, `String` and +/// `Vec`, but Beacon's [`TimestampNanosecond`] is its own newtype over +/// `i64` and needs a conversion. This trait hides that difference behind one +/// entry point, so [`AtlasArrayBackend`] stays generic. #[async_trait::async_trait] -pub trait AtlasReadable: NdArrayType { +pub trait AtlasElement: NdArrayType { + /// Read `shape` elements of `array` from `start`. async fn read( - view: &atlas::DatasetView, - array_name: &str, + view: &DatasetView, + array: &str, start: Vec, shape: Vec, ) -> anyhow::Result>; - /// Convert an atlas [`FillValue`](atlas::FillValue) into this type's - /// per-element fill, using the same widening/sentinel rules `array_format` - /// applies when materializing missing chunks. - fn fill_element(fill: Option<&atlas::FillValue>) -> Self; + /// This type's form of an array's fill value. + /// + /// The engine nulls every element equal to it, so it has to be the value + /// the read actually returns for a cell nobody wrote. Deferring to + /// `array-format`'s own conversion is what guarantees that. + fn fill_element(fill: Option<&FillValue>) -> Self; } -macro_rules! impl_atlas_readable_passthrough { +macro_rules! passthrough { ($ty:ty) => { #[async_trait::async_trait] - impl AtlasReadable for $ty { + impl AtlasElement for $ty { async fn read( - view: &atlas::DatasetView, - array_name: &str, + view: &DatasetView, + array: &str, start: Vec, shape: Vec, ) -> anyhow::Result> { - let arr = view - .read_array::<$ty>(array_name, start, shape) + let values = view + .read_array::<$ty>(array, start, shape) .await .map_err(|e| { - anyhow::anyhow!("Failed to read atlas array '{}': {}", array_name, e) - })? - .ok_or_else(|| { anyhow::anyhow!( - "Atlas array '{}' not found in dataset '{}'", - array_name, + "Failed to read atlas array '{array}' of dataset '{}': {e}", view.name() ) })?; - Ok(arr.to_owned()) + Ok(values.into_owned()) } - fn fill_element(fill: Option<&atlas::FillValue>) -> Self { + fn fill_element(fill: Option<&FillValue>) -> Self { <$ty as atlas::ArrayElement>::fill_element(fill) } } }; } -impl_atlas_readable_passthrough!(i8); -impl_atlas_readable_passthrough!(i16); -impl_atlas_readable_passthrough!(i32); -impl_atlas_readable_passthrough!(i64); -impl_atlas_readable_passthrough!(u8); -impl_atlas_readable_passthrough!(u16); -impl_atlas_readable_passthrough!(u32); -impl_atlas_readable_passthrough!(u64); -impl_atlas_readable_passthrough!(f32); -impl_atlas_readable_passthrough!(f64); -impl_atlas_readable_passthrough!(String); -impl_atlas_readable_passthrough!(Vec); - +passthrough!(i8); +passthrough!(i16); +passthrough!(i32); +passthrough!(i64); +passthrough!(u8); +passthrough!(u16); +passthrough!(u32); +passthrough!(u64); +passthrough!(f32); +passthrough!(f64); +passthrough!(String); +passthrough!(Vec); + +/// Both types are `#[repr(transparent)]` over `i64`, so the conversion is a +/// rename. It is still done element by element, because the two are distinct +/// types and a transmute of a whole array would rest on layout rather than on +/// the type system. #[async_trait::async_trait] -impl AtlasReadable for TimestampNanosecond { +impl AtlasElement for TimestampNanosecond { async fn read( - view: &atlas::DatasetView, - array_name: &str, + view: &DatasetView, + array: &str, start: Vec, shape: Vec, ) -> anyhow::Result> { - let arr = view - .read_array::(array_name, start, shape) + let values = view + .read_array::(array, start, shape) .await .map_err(|e| { - anyhow::anyhow!("Failed to read atlas timestamp array '{}': {}", array_name, e) - })? - .ok_or_else(|| { anyhow::anyhow!( - "Atlas array '{}' not found in dataset '{}'", - array_name, + "Failed to read atlas timestamp array '{array}' of dataset '{}': {e}", view.name() ) })?; - // Map element-wise: TimestampNs(i64) -> TimestampNanosecond(i64). - // Both are #[repr(transparent)] over i64. - Ok(arr.to_owned().mapv(|ts| TimestampNanosecond(ts.0))) + Ok(values.into_owned().mapv(|ts| TimestampNanosecond(ts.0))) } - fn fill_element(fill: Option<&atlas::FillValue>) -> Self { - let ts = ::fill_element(fill); - TimestampNanosecond(ts.0) + fn fill_element(fill: Option<&FillValue>) -> Self { + TimestampNanosecond(::fill_element(fill).0) } } -/// Backend that reads atlas array data lazily. +/// Reads one atlas array lazily, one requested region at a time. /// -/// Holds an [`Arc`](atlas::Atlas) rather than a -/// [`DatasetView`](atlas::DatasetView): views borrow the store's shared -/// in-memory metadata and are cheap to reopen, so each read reopens the view -/// and issues the subset read against the store's shared, cached array files. +/// The backend holds the [`DatasetView`] rather than the collection and a +/// name. A view is resolved once, when the dataset is built; resolving it per +/// read would cost a linear scan of the collection footer every time. pub struct AtlasArrayBackend { - atlas: Arc, - dataset_name: String, - array_name: String, + view: Arc, + array: String, shape: Vec, dimensions: Vec, chunk_shape: Vec, @@ -135,29 +125,27 @@ pub struct AtlasArrayBackend { impl std::fmt::Debug for AtlasArrayBackend { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("AtlasArrayBackend") - .field("dataset_name", &self.dataset_name) - .field("array_name", &self.array_name) + .field("dataset", &self.view.name()) + .field("array", &self.array) .field("shape", &self.shape) .field("dimensions", &self.dimensions) .field("chunk_shape", &self.chunk_shape) - .finish() + .finish_non_exhaustive() } } impl AtlasArrayBackend { pub fn new( - atlas: Arc, - dataset_name: String, - array_name: String, + view: Arc, + array: String, shape: Vec, dimensions: Vec, chunk_shape: Vec, fill_value: Option, ) -> Self { Self { - atlas, - dataset_name, - array_name, + view, + array, shape, dimensions, chunk_shape, @@ -167,7 +155,7 @@ impl AtlasArrayBackend { } #[async_trait::async_trait] -impl ArrayBackend for AtlasArrayBackend { +impl ArrayBackend for AtlasArrayBackend { fn len(&self) -> usize { self.shape.iter().product() } @@ -180,6 +168,10 @@ impl ArrayBackend for AtlasArrayBackend { self.dimensions.clone() } + /// The chunk shape the writer chose. + /// + /// The scan cuts a dataset on this grid, so one unit of work is one stored + /// chunk and a read fetches no block it does not need. fn chunk_shape(&self) -> Vec { self.chunk_shape.clone() } @@ -189,14 +181,14 @@ impl ArrayBackend for AtlasArrayBackend { } async fn read_subset(&self, subset: ArraySubset) -> anyhow::Result> { - let view = self.atlas.open_dataset(&self.dataset_name).await.map_err(|e| { - anyhow::anyhow!("Failed to open atlas dataset '{}': {}", self.dataset_name, e) - })?; - T::read(&view, &self.array_name, subset.start, subset.shape).await + T::read(&self.view, &self.array, subset.start, subset.shape).await } } -/// Backend for scalar attribute values surfaced as rank-0 arrays. +/// Holds one attribute value as a rank-0 array. +/// +/// The value came from the collection footer, which the open already read, so +/// nothing here touches the store. #[derive(Debug)] pub struct AttributeBackend { value: T, @@ -234,133 +226,212 @@ impl ArrayBackend for AttributeBackend { #[cfg(test)] mod tests { use super::*; - use crate::reader::test_support::build_two_dataset_store; - use atlas::Atlas; - - // ── AttributeBackend ─────────────────────────────────────────────── + use crate::test_support; - #[tokio::test] - async fn attribute_backend_is_rank_zero() { - let backend = AttributeBackend::new("hello".to_string()); - assert_eq!(backend.len(), 1); - assert!(ArrayBackend::::shape(&backend).is_empty()); - assert!(ArrayBackend::::dimensions(&backend).is_empty()); + /// Open one dataset of a fixture collection. + async fn view(dir: &std::path::Path, dataset: &str) -> Arc { + let atlas = atlas::Atlas::open_path(dir).await.expect("open"); + Arc::new(atlas.dataset(dataset).expect("dataset")) } - #[tokio::test] - async fn attribute_backend_read_subset_returns_value() { - let backend = AttributeBackend::new(42i32); - let arr = backend - .read_subset(ArraySubset { - start: vec![], - shape: vec![], - }) - .await - .expect("read"); - assert_eq!(arr.ndim(), 0); - let raw = arr.into_raw_vec_and_offset().0; - assert_eq!(raw, vec![42i32]); - } + // ── AtlasArrayBackend ─────────────────────────────────────────────── - // ── AtlasReadable::fill_element ──────────────────────────────────── + #[tokio::test] + async fn the_backend_reports_what_the_footer_holds() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; - #[test] - fn fill_element_passthrough_numeric() { - use atlas::FillValue; + let backend = AtlasArrayBackend::::new( + view(tmp.path(), "winter").await, + "cycle".to_string(), + vec![4], + vec!["obs".to_string()], + vec![4], + Some(-1), + ); + assert_eq!(ArrayBackend::::shape(&backend), vec![4]); assert_eq!( - ::fill_element(Some(&FillValue::Int(-7))), - -7i32 + ArrayBackend::::dimensions(&backend), + vec!["obs".to_string()] ); - let nan = ::fill_element(Some(&FillValue::Float(f64::NAN))); - assert!(nan.is_nan(), "NaN fill must round-trip as NaN"); - assert_eq!(::fill_element(None), 0i32); + assert_eq!(ArrayBackend::::chunk_shape(&backend), vec![4]); + assert_eq!(ArrayBackend::::fill_value(&backend), Some(-1)); + assert_eq!(backend.len(), 4); } - // ── AtlasArrayBackend ────────────────────────────────────────────── - #[tokio::test] - async fn atlas_array_backend_reports_metadata() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; - let atlas = Atlas::open_path(tmp.path()).await.expect("open atlas"); + async fn a_full_read_returns_every_value() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; let backend = AtlasArrayBackend::::new( - Arc::new(atlas), - "winter".into(), - "temperature".into(), + view(tmp.path(), "winter").await, + "temperature".to_string(), vec![4], - vec!["obs".into()], + vec!["obs".to_string()], vec![4], - Some(-1.0f32), + None, ); - assert_eq!( - as ArrayBackend>::shape(&backend), - vec![4] + let values = backend + .read_subset(ArraySubset::new(vec![0], vec![4])) + .await + .unwrap(); + assert_eq!(values.into_raw_vec_and_offset().0, vec![1.0, 2.0, 3.0, 4.0]); + } + + #[tokio::test] + async fn a_window_returns_only_its_own_values() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let backend = AtlasArrayBackend::::new( + view(tmp.path(), "winter").await, + "cycle".to_string(), + vec![4], + vec!["obs".to_string()], + vec![4], + None, ); - assert_eq!( - as ArrayBackend>::dimensions(&backend), - vec!["obs".to_string()] + let values = backend + .read_subset(ArraySubset::new(vec![1], vec![2])) + .await + .unwrap(); + assert_eq!(values.into_raw_vec_and_offset().0, vec![20, 30]); + } + + /// A window that spans two stored chunks assembles across them, and lands + /// in row-major order. + #[tokio::test] + async fn a_window_across_chunks_is_assembled_in_order() { + let tmp = tempfile::tempdir().unwrap(); + test_support::chunked_grid(tmp.path()).await; + + let backend = AtlasArrayBackend::::new( + view(tmp.path(), "grid").await, + "temperature".to_string(), + vec![4, 6], + vec!["lat".to_string(), "lon".to_string()], + vec![2, 3], + None, ); + // Rows 1..3, columns 2..4 of a 4x6 grid whose value is row * 6 + col. + // That window straddles all four chunk columns and both chunk rows. + let values = backend + .read_subset(ArraySubset::new(vec![1, 2], vec![2, 2])) + .await + .unwrap(); + assert_eq!(values.shape(), &[2, 2]); assert_eq!( - as ArrayBackend>::chunk_shape(&backend), - vec![4] + values.into_raw_vec_and_offset().0, + vec![8.0, 9.0, 14.0, 15.0] ); - assert_eq!( - as ArrayBackend>::fill_value(&backend), - Some(-1.0f32) + } + + /// A region nobody wrote reads back as the fill value, and costs no bytes. + #[tokio::test] + async fn an_unwritten_region_reads_as_the_fill() { + let tmp = tempfile::tempdir().unwrap(); + test_support::chunked_grid(tmp.path()).await; + + let backend = AtlasArrayBackend::::new( + view(tmp.path(), "grid").await, + "sparse".to_string(), + vec![4, 6], + vec!["lat".to_string(), "lon".to_string()], + vec![2, 3], + Some(-999.0), ); - assert_eq!(backend.len(), 4); + let values = backend + .read_subset(ArraySubset::new(vec![2, 0], vec![1, 3])) + .await + .unwrap(); + assert_eq!(values.into_raw_vec_and_offset().0, vec![-999.0; 3]); } #[tokio::test] - async fn atlas_array_backend_read_subset_full_range() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; - let atlas = Atlas::open_path(tmp.path()).await.expect("open atlas"); + async fn a_timestamp_array_reads_as_beacons_own_newtype() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; - let backend = AtlasArrayBackend::::new( - Arc::new(atlas), - "winter".into(), - "temperature".into(), + let backend = AtlasArrayBackend::::new( + view(tmp.path(), "winter").await, + "time".to_string(), vec![4], - vec!["obs".into()], + vec!["obs".to_string()], vec![4], None, ); - let arr = backend - .read_subset(ArraySubset { - start: vec![0], - shape: vec![4], - }) + let values = backend + .read_subset(ArraySubset::new(vec![0], vec![2])) .await - .expect("read full"); - let raw = arr.into_raw_vec_and_offset().0; - assert_eq!(raw, vec![1.0f32, 2.0, 3.0, 4.0]); + .unwrap(); + assert_eq!( + values.into_raw_vec_and_offset().0, + vec![ + TimestampNanosecond(test_support::EPOCH_NANOS), + TimestampNanosecond(test_support::EPOCH_NANOS + 86_400_000_000_000), + ] + ); } #[tokio::test] - async fn atlas_array_backend_read_subset_partial_range() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; - let atlas = Atlas::open_path(tmp.path()).await.expect("open atlas"); - - let backend = AtlasArrayBackend::::new( - Arc::new(atlas), - "winter".into(), - "cycle".into(), - vec![4], - vec!["obs".into()], - vec![4], + async fn a_string_array_reads_its_values() { + let tmp = tempfile::tempdir().unwrap(); + test_support::incompatible(tmp.path()).await; + + let backend = AtlasArrayBackend::::new( + view(tmp.path(), "a").await, + "value".to_string(), + vec![2], + vec!["obs".to_string()], + vec![2], None, ); - let arr = backend - .read_subset(ArraySubset { - start: vec![1], - shape: vec![2], - }) + let values = backend + .read_subset(ArraySubset::new(vec![0], vec![2])) + .await + .unwrap(); + assert_eq!( + values.into_raw_vec_and_offset().0, + vec!["x".to_string(), "y".to_string()] + ); + } + + // ── fill values ───────────────────────────────────────────────────── + + #[test] + fn a_fill_takes_the_form_array_format_returns() { + assert_eq!( + ::fill_element(Some(&FillValue::Int(-7))), + -7 + ); + assert!(::fill_element(Some(&FillValue::Float(f64::NAN))).is_nan()); + assert_eq!(::fill_element(None), 0); + assert_eq!( + ::fill_element(Some(&FillValue::TimestampNs( + i64::MIN + ))), + TimestampNanosecond(i64::MIN) + ); + } + + // ── AttributeBackend ──────────────────────────────────────────────── + + #[tokio::test] + async fn an_attribute_is_one_value_on_no_axis() { + let backend = AttributeBackend::new("winter".to_string()); + assert_eq!(backend.len(), 1); + assert!(ArrayBackend::::shape(&backend).is_empty()); + assert!(ArrayBackend::::dimensions(&backend).is_empty()); + + let values = backend + .read_subset(ArraySubset::new(vec![], vec![])) .await - .expect("read partial"); - let raw = arr.into_raw_vec_and_offset().0; - assert_eq!(raw, vec![20i32, 30]); + .unwrap(); + assert_eq!(values.ndim(), 0); + assert_eq!( + values.into_raw_vec_and_offset().0, + vec!["winter".to_string()] + ); } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs index d0b3f8d7..65131d77 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs @@ -1,184 +1,178 @@ -//! Conversion between atlas arrays/attributes and `beacon-nd-array` types. +//! The mapping between an Atlas collection and Beacon's ND array model: column +//! names, element types, and the lazy arrays themselves. +//! +//! One mapping, in one place. The Arrow type of a column follows from its +//! [`NdArrayDataType`] through `beacon-nd-array`'s own conversion, so a schema +//! derived here and a batch produced by a scan can never disagree. use std::sync::Arc; -use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; -use atlas::{ArraySchema, Attr, DType, FillValue, MergedSchema}; -use beacon_nd_array::{NdArray, NdArrayD, datatypes::TimestampNanosecond}; +use arrow::datatypes::DataType; +use atlas::{ArraySchema, Attr, DType, DatasetView, FillValue}; +use beacon_nd_array::{ + NdArray, NdArrayD, datatypes::NdArrayDataType, datatypes::TimestampNanosecond, +}; -use crate::backend::{AtlasArrayBackend, AtlasReadable, AttributeBackend}; +use crate::backend::{AtlasArrayBackend, AtlasElement, AttributeBackend}; -/// Arrow type for a scalar atlas [`DType`], or `None` for the list dtypes -/// (`List`, `FixedSizeList`) that have no rank-0 / column analogue in Beacon. +// ─── Column names ──────────────────────────────────────────────────────────── + +/// The column a per-array attribute is surfaced under: `{array}.{attr}`. +pub fn array_attr_column(array: &str, attr: &str) -> String { + format!("{array}.{attr}") +} + +/// The column a dataset-level attribute is surfaced under: `.{attr}`. +/// +/// The leading dot is what netCDF and Zarr use, and it keeps a dataset +/// attribute from colliding with an array of the same name. +pub fn global_attr_column(attr: &str) -> String { + format!(".{attr}") +} + +/// Whether `column` could name a per-array attribute of `array`. /// -/// Kept in lock-step with the [`NdArrayType`](beacon_nd_array::datatypes::NdArrayType) -/// → Arrow mapping the read path produces, so a schema derived here and a batch -/// produced by the scan carry matching types. -fn scalar_dtype_to_arrow(dtype: &DType) -> Option { +/// Used to skip building an attribute map for an array whose attributes the +/// query does not project. +pub fn is_attr_column_of(column: &str, array: &str) -> bool { + column.len() > array.len() + 1 + && column.starts_with(array) + && column.as_bytes()[array.len()] == b'.' +} + +// ─── Element types ─────────────────────────────────────────────────────────── + +/// The ND type of a scalar atlas dtype, or `None` for the list dtypes, which +/// have no rank-0 or column analogue in Beacon. +fn scalar_dtype_to_nd(dtype: &DType) -> Option { Some(match dtype { - DType::Bool => DataType::Boolean, - DType::Int8 => DataType::Int8, - DType::Int16 => DataType::Int16, - DType::Int32 => DataType::Int32, - DType::Int64 => DataType::Int64, - DType::UInt8 => DataType::UInt8, - DType::UInt16 => DataType::UInt16, - DType::UInt32 => DataType::UInt32, - DType::UInt64 => DataType::UInt64, - DType::Float32 => DataType::Float32, - DType::Float64 => DataType::Float64, - DType::String => DataType::Utf8, - DType::Binary => DataType::Binary, - DType::TimestampNs => DataType::Timestamp(TimeUnit::Nanosecond, None), + DType::Bool => NdArrayDataType::Bool, + DType::Int8 => NdArrayDataType::I8, + DType::Int16 => NdArrayDataType::I16, + DType::Int32 => NdArrayDataType::I32, + DType::Int64 => NdArrayDataType::I64, + DType::UInt8 => NdArrayDataType::U8, + DType::UInt16 => NdArrayDataType::U16, + DType::UInt32 => NdArrayDataType::U32, + DType::UInt64 => NdArrayDataType::U64, + DType::Float32 => NdArrayDataType::F32, + DType::Float64 => NdArrayDataType::F64, + DType::String => NdArrayDataType::String, + DType::Binary => NdArrayDataType::Binary, + DType::TimestampNs => NdArrayDataType::Timestamp, DType::List { .. } | DType::FixedSizeList { .. } => return None, }) } -/// Arrow type for an atlas **array** dtype, or `None` for dtypes Beacon can't -/// read as a column. `Bool` is excluded here (atlas's `ArrayElement` isn't -/// implemented for `bool`), matching [`array_to_nd_array`]'s rejection. -pub fn atlas_array_dtype_to_arrow(dtype: &DType) -> Option { +/// The ND type of an atlas **array** dtype, or `None` for one Beacon cannot +/// read as a column. +/// +/// `Bool` is excluded, unlike an attribute: `array-format` implements no +/// element type for `bool`, so no reader can produce the values. Every list +/// dtype is excluded too. +pub fn array_dtype_to_nd(dtype: &DType) -> Option { match dtype { DType::Bool => None, - other => scalar_dtype_to_arrow(other), + other => scalar_dtype_to_nd(other), } } -/// Arrow type for an atlas **attribute** dtype, or `None` for list-valued -/// attributes (no rank-0 analogue). Scalar `Bool` attributes *are* supported. -pub fn atlas_attr_dtype_to_arrow(dtype: &DType) -> Option { - scalar_dtype_to_arrow(dtype) +/// The ND type of an atlas **attribute** dtype, or `None` for a list-valued +/// one. A scalar `Bool` attribute *is* supported: its value comes from the +/// footer rather than from an array. +pub fn attr_dtype_to_nd(dtype: &DType) -> Option { + scalar_dtype_to_nd(dtype) } -/// Build the Arrow schema for a whole atlas store from its collection-wide -/// [`MergedSchema`] — **no per-dataset iteration and no disk I/O**. -/// -/// The merged schema already widens every array/attribute dtype across all -/// datasets (the same union `super_type_schema` would compute), so this is the -/// scale-friendly replacement for opening and typing each dataset in turn. -/// -/// Columns mirror the reader's naming: arrays by name, per-array attributes as -/// `{array}.{attr}`, dataset-level attributes by their bare key. Fields are -/// sorted by name so the layout is stable (adaptation is by name, so order is -/// cosmetic). When `read_dimensions` is `Some`, only arrays whose dimensions are -/// a subset survive; rank-0 attributes (empty dimensions) always survive, so -/// per-array attributes are emitted regardless of their array's dimensionality — -/// matching what the scan produces after `resolve_read_dimensions`. -pub fn atlas_merged_schema_to_arrow( - merged: &MergedSchema, - read_dimensions: Option<&[String]>, -) -> Schema { - let mut fields: Vec = Vec::new(); - - for (name, arr) in &merged.arrays { - let dims_ok = read_dimensions - .map_or(true, |dims| arr.dimension_names.iter().all(|d| dims.contains(d))); - if dims_ok - && let Some(dt) = atlas_array_dtype_to_arrow(&arr.dtype.0) - { - fields.push(Field::new(name, dt, true)); - } - for (attr, ty) in &arr.attributes { - if let Some(dt) = atlas_attr_dtype_to_arrow(&ty.0) { - fields.push(Field::new(format!("{name}.{attr}"), dt, true)); - } - } - } +/// The Arrow type of an atlas array dtype, or `None` when Beacon cannot read +/// it. Derived from [`array_dtype_to_nd`], so it always matches the scan. +pub fn array_dtype_to_arrow(dtype: &DType) -> Option { + array_dtype_to_nd(dtype).map(Into::into) +} - for (key, ty) in &merged.global_attributes { - if let Some(dt) = atlas_attr_dtype_to_arrow(&ty.0) { - fields.push(Field::new(key, dt, true)); - } - } +/// The Arrow type of an atlas attribute dtype, or `None` for a list. +pub fn attr_dtype_to_arrow(dtype: &DType) -> Option { + attr_dtype_to_nd(dtype).map(Into::into) +} - fields.sort_by(|a, b| a.name().cmp(b.name())); - Schema::new(fields) +/// A stable tag for a dtype, for keys that group datasets by shape. +pub(crate) fn dtype_tag(dtype: &DType) -> String { + format!("{dtype:?}") } -/// Convert an atlas array (described by its [`ArraySchema`]) into a lazy -/// [`NdArrayD`] backed by [`AtlasArrayBackend`]. +// ─── Lazy arrays ───────────────────────────────────────────────────────────── + +/// Wrap one atlas array as a lazy [`NdArrayD`] over `view`. /// -/// `Bool`, `FixedSizeList` and `List` dtypes are rejected with an explicit -/// error — atlas's `ArrayElement` isn't implemented for `bool`, and Beacon's -/// ND array model has no analogue for list dtypes. Silently skipping them would -/// propagate dimension mismatches, so the caller (the reader) `warn!`-skips the -/// column instead. +/// Nothing is read here. The shape, the dimension names, the chunk shape and +/// the fill value all come from the collection footer, which the open already +/// held; the values arrive when the engine asks the backend for a subset. /// -/// `fill_value` comes from -/// [`DatasetView::array_fill_value`](atlas::DatasetView::array_fill_value) and -/// is converted to the per-dtype `T` via [`AtlasReadable::fill_element`]. +/// The chunk shape is the one the writer chose. It is what lets the scan cut a +/// dataset on the grid the file actually stores, so one unit of work is one +/// stored chunk. pub fn array_to_nd_array( - atlas: Arc, - dataset_name: &str, + view: Arc, array_name: &str, schema: &ArraySchema, - fill_value: Option, ) -> anyhow::Result> { - let shape = schema.shape.clone(); - let dimensions = schema.dimension_names.clone(); - let chunk_shape = schema.chunk_shape.clone(); + let fill: Option = schema.fill_value.clone().map(Into::into); - macro_rules! mk { + macro_rules! lazy { ($ty:ty) => {{ - let fill: Option<$ty> = fill_value + let fill = fill .as_ref() - .map(|fv| <$ty as AtlasReadable>::fill_element(Some(fv))); + .map(|value| <$ty as AtlasElement>::fill_element(Some(value))); let backend = AtlasArrayBackend::<$ty>::new( - atlas.clone(), - dataset_name.to_string(), + view, array_name.to_string(), - shape.clone(), - dimensions.clone(), - chunk_shape.clone(), + schema.shape.clone(), + schema.dimension_names.clone(), + schema.chunk_shape.clone(), fill, ); - let nd = NdArray::new_with_backend(backend)?; - Ok::, anyhow::Error>(Arc::new(nd)) + Ok(Arc::new(NdArray::new_with_backend(backend)?) as Arc) }}; } match &schema.dtype { + DType::Int8 => lazy!(i8), + DType::Int16 => lazy!(i16), + DType::Int32 => lazy!(i32), + DType::Int64 => lazy!(i64), + DType::UInt8 => lazy!(u8), + DType::UInt16 => lazy!(u16), + DType::UInt32 => lazy!(u32), + DType::UInt64 => lazy!(u64), + DType::Float32 => lazy!(f32), + DType::Float64 => lazy!(f64), + DType::String => lazy!(String), + DType::Binary => lazy!(Vec), + DType::TimestampNs => lazy!(TimestampNanosecond), DType::Bool => Err(anyhow::anyhow!( - "Atlas array '{}' has dtype Bool which is not readable through Beacon \ - (atlas's ArrayElement does not implement bool)", - array_name + "array '{array_name}' is Bool, which atlas stores no elements of" )), - DType::Int8 => mk!(i8), - DType::Int16 => mk!(i16), - DType::Int32 => mk!(i32), - DType::Int64 => mk!(i64), - DType::UInt8 => mk!(u8), - DType::UInt16 => mk!(u16), - DType::UInt32 => mk!(u32), - DType::UInt64 => mk!(u64), - DType::Float32 => mk!(f32), - DType::Float64 => mk!(f64), - DType::String => mk!(String), - DType::Binary => mk!(Vec), - DType::TimestampNs => mk!(TimestampNanosecond), DType::FixedSizeList { .. } => Err(anyhow::anyhow!( - "Atlas array '{}' has unsupported dtype FixedSizeList — Beacon does not model \ - fixed-size lists", - array_name + "array '{array_name}' is a FixedSizeList, which Beacon does not model" )), DType::List { .. } => Err(anyhow::anyhow!( - "Atlas array '{}' has unsupported dtype List — Beacon does not model \ - variable-length lists", - array_name + "array '{array_name}' is a List, which Beacon does not model" )), } } -/// Convert a scalar atlas attribute value into a rank-0 [`NdArrayD`]. +/// Wrap one scalar attribute value as a rank-0 [`NdArrayD`]. /// -/// List-valued attributes have no rank-0 scalar analogue in Beacon's ND array -/// model and are rejected with an error (the caller `warn!`-skips them). +/// A rank-0 array broadcasts onto whatever grid the dataset's own arrays +/// define, so the value repeats across every row the dataset contributes. +/// A list-valued attribute has no such analogue and is refused. pub fn attribute_to_nd_array(attr: &Attr) -> anyhow::Result> { macro_rules! scalar { ($value:expr) => { - Ok(Arc::new(NdArray::new_with_backend(AttributeBackend::new( - $value, - ))?) as Arc) + Ok( + Arc::new(NdArray::new_with_backend(AttributeBackend::new($value))?) + as Arc, + ) }; } @@ -197,20 +191,9 @@ pub fn attribute_to_nd_array(attr: &Attr) -> anyhow::Result> { Attr::String(v) => scalar!(v.clone()), Attr::Binary(v) => scalar!(v.clone()), Attr::TimestampNanoseconds(v) => scalar!(TimestampNanosecond(*v)), - Attr::BoolList(_) - | Attr::Int8List(_) - | Attr::Int16List(_) - | Attr::Int32List(_) - | Attr::Int64List(_) - | Attr::UInt8List(_) - | Attr::UInt16List(_) - | Attr::UInt32List(_) - | Attr::UInt64List(_) - | Attr::Float32List(_) - | Attr::Float64List(_) - | Attr::StringList(_) - | Attr::BinaryList(_) => Err(anyhow::anyhow!( - "list-valued attributes are not representable as Beacon rank-0 arrays" + other => Err(anyhow::anyhow!( + "attribute is a {} list, which has no rank-0 form in Beacon", + dtype_tag(&other.dtype()) )), } } @@ -218,101 +201,115 @@ pub fn attribute_to_nd_array(attr: &Attr) -> anyhow::Result> { #[cfg(test)] mod tests { use super::*; - use atlas::{Atlas, Codec, DType, StoreConfig}; - use beacon_nd_array::{NdArray, datatypes::NdArrayDataType}; - - fn schema_with_dtype(dtype: DType) -> ArraySchema { - ArraySchema { - dtype, - shape: vec![2], - chunk_shape: vec![2], - dimension_names: vec!["x".into()], - codec: Codec::default(), - } + use beacon_nd_array::NdArray; + + // ── column names ──────────────────────────────────────────────────── + + #[test] + fn an_attribute_takes_its_owners_name() { + assert_eq!(array_attr_column("sst", "units"), "sst.units"); + assert_eq!(global_attr_column("Conventions"), ".Conventions"); } - async fn dummy_atlas() -> Arc { - // The rejection branches return before touching the atlas handle, so we - // need a value but never read from it. - let tmp = tempfile::tempdir().expect("temp dir"); - let atlas = Atlas::create_path(tmp.path(), StoreConfig::default()) - .await - .expect("create dummy atlas"); - std::mem::forget(tmp); - Arc::new(atlas) + #[test] + fn an_attribute_column_is_recognized_by_its_array() { + assert!(is_attr_column_of("sst.units", "sst")); + assert!( + !is_attr_column_of("sst", "sst"), + "the array itself is not one" + ); + assert!( + !is_attr_column_of("sst_flag.units", "sst"), + "a prefix is not a name" + ); + assert!(!is_attr_column_of("sst.", "sst"), "an empty key is no key"); } - #[tokio::test] - async fn array_to_nd_array_rejects_bool() { - let atlas = dummy_atlas().await; - let err = array_to_nd_array(atlas, "ds", "flag", &schema_with_dtype(DType::Bool), None) - .expect_err("Bool should be rejected"); - let msg = format!("{err:#}"); - assert!(msg.contains("Bool"), "{msg}"); - assert!(msg.contains("flag"), "{msg}"); + // ── element types ─────────────────────────────────────────────────── + + #[test] + fn every_readable_array_dtype_maps() { + let cases = [ + (DType::Int8, NdArrayDataType::I8), + (DType::Int16, NdArrayDataType::I16), + (DType::Int32, NdArrayDataType::I32), + (DType::Int64, NdArrayDataType::I64), + (DType::UInt8, NdArrayDataType::U8), + (DType::UInt16, NdArrayDataType::U16), + (DType::UInt32, NdArrayDataType::U32), + (DType::UInt64, NdArrayDataType::U64), + (DType::Float32, NdArrayDataType::F32), + (DType::Float64, NdArrayDataType::F64), + (DType::String, NdArrayDataType::String), + (DType::Binary, NdArrayDataType::Binary), + (DType::TimestampNs, NdArrayDataType::Timestamp), + ]; + for (dtype, expected) in cases { + assert_eq!(array_dtype_to_nd(&dtype), Some(expected), "{dtype:?}"); + } } - #[tokio::test] - async fn array_to_nd_array_rejects_list() { - let atlas = dummy_atlas().await; - let err = array_to_nd_array( - atlas, - "ds", - "events", - &schema_with_dtype(DType::List { - child: Box::new(DType::Int32), - }), - None, - ) - .expect_err("List should be rejected"); - let msg = format!("{err:#}"); - assert!(msg.contains("List"), "{msg}"); - assert!(msg.contains("events"), "{msg}"); + /// `array-format` implements no element type for `bool`, so a `Bool` array + /// cannot be read even though the dtype exists. An attribute can. + #[test] + fn a_bool_array_is_refused_but_a_bool_attribute_is_not() { + assert_eq!(array_dtype_to_nd(&DType::Bool), None); + assert_eq!(attr_dtype_to_nd(&DType::Bool), Some(NdArrayDataType::Bool)); } - #[tokio::test] - async fn attribute_bool_round_trips() { - let nd = attribute_to_nd_array(&Attr::Bool(true)).expect("convert"); - assert_eq!(nd.datatype(), NdArrayDataType::Bool); - assert!(nd.shape().is_empty()); - let typed = nd - .as_any() - .downcast_ref::>() - .expect("downcast"); - assert_eq!(typed.clone_into_raw_vec().await, vec![true]); + #[test] + fn list_dtypes_have_no_column() { + let list = DType::List { + child: Box::new(DType::Int32), + }; + let fixed = DType::FixedSizeList { + child: Box::new(DType::Float32), + size: 3, + }; + for dtype in [list, fixed] { + assert_eq!(array_dtype_to_nd(&dtype), None, "{dtype:?}"); + assert_eq!(attr_dtype_to_nd(&dtype), None, "{dtype:?}"); + } } - #[tokio::test] - async fn attribute_int64_round_trips() { - let nd = attribute_to_nd_array(&Attr::Int64(42)).expect("convert"); - assert_eq!(nd.datatype(), NdArrayDataType::I64); - let typed = nd - .as_any() - .downcast_ref::>() - .expect("downcast"); - assert_eq!(typed.clone_into_raw_vec().await, vec![42i64]); + /// The Arrow type follows the ND type, so a schema and a batch agree. + #[test] + fn the_arrow_type_follows_the_nd_type() { + assert_eq!( + array_dtype_to_arrow(&DType::Float64), + Some(DataType::Float64) + ); + assert_eq!(array_dtype_to_arrow(&DType::String), Some(DataType::Utf8)); + assert_eq!( + array_dtype_to_arrow(&DType::TimestampNs), + Some(DataType::Timestamp( + arrow::datatypes::TimeUnit::Nanosecond, + None + )) + ); + assert_eq!(array_dtype_to_arrow(&DType::Bool), None); } + // ── attribute values ──────────────────────────────────────────────── + #[tokio::test] - async fn attribute_string_round_trips() { - let nd = attribute_to_nd_array(&Attr::String("winter".into())).expect("convert"); - assert_eq!(nd.datatype(), NdArrayDataType::String); - let typed = nd - .as_any() - .downcast_ref::>() - .expect("downcast"); - assert_eq!(typed.clone_into_raw_vec().await, vec!["winter".to_string()]); + async fn a_scalar_attribute_is_a_rank_zero_column() { + let nd = attribute_to_nd_array(&Attr::Int64(2024)).unwrap(); + assert_eq!(nd.datatype(), NdArrayDataType::I64); + assert!(nd.shape().is_empty(), "an attribute has no axis"); + let typed = nd.as_any().downcast_ref::>().unwrap(); + assert_eq!(typed.clone_into_raw_vec().await, vec![2024]); } #[tokio::test] - async fn attribute_timestamp_round_trips() { - let nanos = 1_700_000_000_000_000_000i64; - let nd = attribute_to_nd_array(&Attr::TimestampNanoseconds(nanos)).expect("convert"); + async fn a_timestamp_attribute_keeps_its_type() { + let nanos = 1_700_000_000_000_000_000; + let nd = attribute_to_nd_array(&Attr::TimestampNanoseconds(nanos)).unwrap(); assert_eq!(nd.datatype(), NdArrayDataType::Timestamp); let typed = nd .as_any() .downcast_ref::>() - .expect("downcast"); + .unwrap(); assert_eq!( typed.clone_into_raw_vec().await, vec![TimestampNanosecond(nanos)] @@ -320,39 +317,16 @@ mod tests { } #[tokio::test] - async fn attribute_list_rejected() { - let err = attribute_to_nd_array(&Attr::Int32List(vec![1, 2, 3])) - .expect_err("list attribute should be rejected"); - assert!(format!("{err:#}").contains("list-valued")); + async fn a_bool_attribute_is_a_column() { + let nd = attribute_to_nd_array(&Attr::Bool(true)).unwrap(); + assert_eq!(nd.datatype(), NdArrayDataType::Bool); } - /// Two datasets giving the same array *non-numeric* conflicting dtypes - /// (`String` vs `Int64`) still resolve: atlas widens the union to `String`, - /// so the table column is `Utf8` rather than the merge failing or the column - /// being dropped. Pins the assumption the scan relies on — that every merged - /// dtype is something Arrow can cast each dataset's native type *into*. - #[tokio::test] - async fn incompatible_array_dtypes_merge_to_string() { - let tmp = tempfile::tempdir().expect("temp dir"); - crate::reader::test_support::build_incompatible_store(tmp.path()).await; - let atlas = Atlas::open_path(tmp.path()).await.expect("open atlas"); - - let merged = atlas.merged_schema(); - assert_eq!( - merged.arrays.get("value").expect("value in merged schema").dtype.0, - DType::String, - "String ∪ Int64 must widen to String, not fail or drop the array" - ); - - let schema = atlas_merged_schema_to_arrow(&merged, None); - assert_eq!( - schema.field_with_name("value").expect("value field").data_type(), - &DataType::Utf8 - ); - // The column only one dataset declares still appears, at its own type. - assert_eq!( - schema.field_with_name("only_a").expect("only_a field").data_type(), - &DataType::Int32 - ); + #[test] + fn a_list_attribute_is_refused_by_name() { + let error = attribute_to_nd_array(&Attr::Int32List(vec![1, 2, 3])) + .expect_err("a list has no rank-0 form") + .to_string(); + assert!(error.contains("list"), "{error}"); } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs new file mode 100644 index 00000000..77ed9caf --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs @@ -0,0 +1,48 @@ +//! [`AtlasConfig`]: the runtime settings of the Atlas format. + +/// Runtime configuration for the Atlas format. +/// +/// Plain data with sensible defaults; the caller populates it. There is no +/// environment parsing here, so the crate stays reusable and the host decides +/// where the values come from. Each field is the *default* for a runtime, and +/// each can be overridden per table via +/// `CREATE EXTERNAL TABLE ... OPTIONS (...)`. +#[derive(Debug, Clone)] +pub struct AtlasConfig { + /// Whether a read consults the shared reader cache. + /// + /// A collection is immutable, so a cached handle stays valid until its + /// deletion mask changes. The cache saves the footer read and keeps the + /// decompressed blocks of a collection between queries. + pub use_reader_cache: bool, + /// How many opened collections the shared reader cache holds. + /// + /// Each entry owns its own block cache — 256 MiB of decompressed blocks and + /// 64 MiB of raw slabs — so this is also a memory bound: the default of 32 + /// admits up to 10 GiB of cached blocks across every open collection. + pub reader_cache_size: u64, + /// Whether a predicate scan drops the datasets that cannot match, before it + /// reads them. + /// + /// A pure optimization: pruning only ever removes datasets that hold no + /// matching row, and every path fails open. Off trades throughput for + /// skipping the index build. + pub use_pruning: bool, + /// Whether the file analyzer measures a collection's column ranges. + /// + /// The ranges come from the footer, so they cost no array read. The switch + /// exists because a listing of many collections turns even a footer read + /// per collection into real I/O. + pub enable_statistics: bool, +} + +impl Default for AtlasConfig { + fn default() -> Self { + Self { + use_reader_cache: true, + reader_cache_size: 32, + use_pruning: true, + enable_statistics: true, + } + } +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/cache.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/cache.rs deleted file mode 100644 index 274eb49c..00000000 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/cache.rs +++ /dev/null @@ -1,151 +0,0 @@ -//! Cache of opened atlas stores keyed by marker path plus freshness -//! (`last_modified` + `size`). Callers that want an `Arc` go through -//! [`get_or_open_atlas`] so a single store is opened once for as long as its -//! on-disk metadata is unchanged. -//! -//! The cache is owned per-runtime ([`AtlasReaderCache`]) rather than being a -//! process-global static; passing `None` opens directly with no caching. - -use std::sync::Arc; - -use atlas::Atlas; -use moka::future::Cache; -use object_store::{ObjectMeta, ObjectStore, path::Path as OsPath}; - -use crate::util::atlas_store_prefix; - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -struct CacheKey { - path: OsPath, - last_modified: chrono::DateTime, - size: u64, -} - -/// A reader cache for opened atlas stores, sized at construction time. -/// -/// Cloning shares the underlying [`moka`] cache (reference-counted internally), -/// so one instance is shared across the formats, sources and openers a runtime -/// hands a clone to. Per-runtime state — there is no process-global cache. -#[derive(Clone)] -pub struct AtlasReaderCache { - cache: Cache>, -} - -impl AtlasReaderCache { - /// Build a cache holding up to `capacity` opened atlas stores. - pub fn new(capacity: u64) -> Self { - Self { - cache: Cache::builder().max_capacity(capacity).build(), - } - } -} - -// `Atlas` is not `Debug`; the cache is embedded in `Debug` structs -// (formats/sources), so provide an opaque impl. -impl std::fmt::Debug for AtlasReaderCache { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AtlasReaderCache").finish_non_exhaustive() - } -} - -/// Open the atlas store whose marker lives at `marker.location`, over `store`. -/// -/// Atlas opens natively over the [`object_store`] backend: the store prefix is -/// the marker's parent directory, and the metadata variant is auto-detected -/// from the files present. -async fn open_atlas_store( - store: Arc, - marker_path: &OsPath, -) -> datafusion::error::Result> { - let prefix = atlas_store_prefix(marker_path).ok_or_else(|| { - datafusion::error::DataFusionError::Execution(format!( - "Path {marker_path} is not an atlas metadata marker" - )) - })?; - let atlas = Atlas::open(store, prefix.clone()).await.map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to open atlas store at prefix '{prefix}': {e}" - )) - })?; - Ok(Arc::new(atlas)) -} - -/// Return a cached [`Arc`] for `marker`, opening from `store` on miss. -/// -/// When `cache` is `None`, the store is opened directly with no caching. -/// Otherwise freshness is encoded in the cache key — a marker whose -/// `last_modified` or `size` differs from the cached entry produces a new key, -/// forcing a re-open. Concurrent first-readers for the same key coalesce inside -/// [`moka::future::Cache::try_get_with`]. -pub async fn get_or_open_atlas( - cache: Option<&AtlasReaderCache>, - store: Arc, - marker: &ObjectMeta, -) -> datafusion::error::Result> { - let Some(cache) = cache else { - return open_atlas_store(store, &marker.location).await; - }; - - let key = CacheKey { - path: marker.location.clone(), - last_modified: marker.last_modified, - size: marker.size, - }; - let path = marker.location.clone(); - - cache - .cache - .try_get_with(key, async move { open_atlas_store(store, &path).await }) - .await - .map_err(|e: Arc| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to open atlas store via cache: {e}" - )) - }) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::datafusion::test_support::{fixture_marker_object_meta, test_store}; - - #[tokio::test] - async fn cache_returns_same_arc_for_identical_marker() { - let store = test_store().await; - let marker = fixture_marker_object_meta(); - let cache = AtlasReaderCache::new(32); - - let first = get_or_open_atlas(Some(&cache), store.clone(), &marker) - .await - .expect("first open"); - let second = get_or_open_atlas(Some(&cache), store, &marker) - .await - .expect("second open"); - - assert!( - Arc::ptr_eq(&first, &second), - "identical marker must hit the cache" - ); - } - - #[tokio::test] - async fn cache_reopens_when_last_modified_changes() { - let store = test_store().await; - let base = fixture_marker_object_meta(); - let mut bumped = base.clone(); - bumped.last_modified = base.last_modified + chrono::Duration::seconds(1); - let cache = AtlasReaderCache::new(32); - - let first = get_or_open_atlas(Some(&cache), store.clone(), &base) - .await - .expect("first open"); - let second = get_or_open_atlas(Some(&cache), store, &bumped) - .await - .expect("second open"); - - assert!( - !Arc::ptr_eq(&first, &second), - "bumped last_modified must invalidate the cache" - ); - } -} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs index e0f25738..4365536d 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs @@ -1,32 +1,48 @@ -//! Execution metrics for the atlas scan, surfaced through DataFusion's standard -//! metrics reporting (e.g. `EXPLAIN ANALYZE`). +//! What one Atlas scan partition did, reported through DataFusion's metrics. //! //! These complement -//! [`DatasetReadMetrics`](beacon_nd_array::arrow::metrics::DatasetReadMetrics) -//! (output rows/batches and engine-level chunk pruning) with the atlas-specific -//! costs: opening the store, pruning datasets, and building each dataset's lazy -//! backends. All names are `atlas_`-prefixed so they never collide with -//! DataFusion's reserved typed metrics (`output_rows`, `output_batches`, …), -//! which aggregate by name and panic on a variant mismatch. +//! [`ReadMetrics`](beacon_nd_array::arrow::metrics::ReadMetrics), which counts +//! the chunks and rows the shared queue handed out. What it cannot see is the +//! cost of reaching a dataset at all: opening the collection, deciding whether +//! the dataset is worth reading, and building its lazy columns. +//! +//! Every name is `atlas_`-prefixed. DataFusion sums metrics that share a name, +//! and `output_rows` and `output_batches` are already registered for this +//! partition by the scan itself. use datafusion::physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder, Time}; -/// Per-partition timings and counts for one atlas scan partition. All fields are -/// `Arc`-backed handles into the shared [`ExecutionPlanMetricsSet`], so cloning -/// is cheap and every clone accumulates into the same metric. +/// Per-partition timings and counts for one Atlas scan partition. +/// +/// Every field is an `Arc`-backed handle into the shared +/// [`ExecutionPlanMetricsSet`], so a clone is cheap and every clone accumulates +/// into the same metric. #[derive(Debug, Clone)] pub struct AtlasScanMetrics { - /// Wall time opening (or cache-hitting) the atlas store for this partition. + /// Wall time opening collections, or hitting the reader cache for them. pub open_time: Time, - /// Wall time computing which datasets the predicate can match (pruning). + /// Wall time deciding which datasets a predicate can rule out. + /// + /// One partition builds a collection's index and the rest wait on it, so + /// this is the build for one of them and the wait for the others. pub prune_time: Time, - /// Wall time building lazy datasets — metadata, backends, projected - /// attribute values, and the per-dataset schema adapter. + /// Wall time building lazy datasets: resolving the view, reading the + /// projected attribute values out of the footer, wiring the backends, and + /// planning the chunk queue. + /// + /// Array data is read later, as the queue is drained, and `ReadMetrics` + /// counts that. pub dataset_build_time: Time, - /// Datasets this partition opened and scanned. + /// Datasets this partition opened and read. pub datasets_scanned: Count, - /// Datasets this partition skipped because pruning ruled them out. + /// Datasets it skipped because the collection's statistics ruled them out. pub datasets_pruned: Count, + /// Pruning indexes built. One per collection a predicate scan touches, so a + /// number above the collection count means a partition rebuilt one. + pub index_builds: Count, + /// Datasets those indexes covered, which is what the pruning pass looked + /// at rather than read. + pub index_rows: Count, } impl AtlasScanMetrics { @@ -40,6 +56,8 @@ impl AtlasScanMetrics { .counter("atlas_datasets_scanned", partition), datasets_pruned: MetricBuilder::new(metrics) .counter("atlas_datasets_pruned", partition), + index_builds: MetricBuilder::new(metrics).counter("atlas_index_builds", partition), + index_rows: MetricBuilder::new(metrics).counter("atlas_index_rows", partition), } } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs index 901e32c7..3b7695ba 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs @@ -1,99 +1,61 @@ -//! DataFusion integration for atlas stores. +//! The DataFusion integration: discovering Atlas collections, typing them, and +//! planning a scan over their datasets. //! -//! Mirrors the zarr crate: an [`AtlasFormatFactory`] discovers atlas metadata -//! markers, [`AtlasFormat`] infers the merged Arrow schema across a -//! store's datasets and plans the scan, and [`AtlasSource`] opens each store -//! natively over the query's object store and streams every dataset through the -//! shared `beacon-nd-array` engine. +//! [`AtlasFormatFactory`] recognizes a collection in a listing and builds an +//! [`AtlasFormat`] per table. The format infers the collection's schema, then +//! plans a scan whose entries are *datasets* rather than files — see +//! [`source`] for what the openers then do with them. use std::any::Any; +use std::collections::HashMap; use std::sync::Arc; -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{Schema, SchemaRef}; use beacon_datafusion_ext::format_ext::{ DatasetMetadata, FileFormatFactoryExt, SchemaOptions, SchemaUnit, units_over_stores, }; +use beacon_datafusion_ext::format_options::format_option; +use beacon_datafusion_ext::listing_factory::ListingFactory; use beacon_datafusion_ext::type_widening::{LabeledSchema, session_widening}; use datafusion::{ catalog::{Session, memory::DataSourceExec}, common::{GetExt, Statistics, exec_datafusion_err}, datasource::{ file_format::{FileFormat, FileFormatFactory, file_compression_type::FileCompressionType}, - listing::PartitionedFile, + listing::{ListingTableUrl, PartitionedFile}, physical_plan::{ FileGroup, FileScanConfig, FileScanConfigBuilder, FileSinkConfig, FileSource, }, + table_schema::TableSchema, }, + error::{DataFusionError, Result}, physical_expr::LexRequirement, physical_plan::ExecutionPlan, }; use object_store::{ObjectMeta, ObjectStore}; -use crate::datafusion::{ - cache::AtlasReaderCache, options::AtlasOptions, source::AtlasDatasetSlice, source::AtlasSource, +use crate::config::AtlasConfig; +use crate::reader::collection_schema; +use crate::store::{ + ATLAS_MARKER, AtlasReaderCache, get_or_open_atlas, is_atlas_marker, top_level_atlas_markers, }; -use crate::util::{ATLAS_MARKER, top_level_atlas_markers}; -pub mod cache; pub mod metrics; pub mod options; pub mod pruning; pub mod source; +pub mod statistics; pub mod table_function; -pub use cache::{AtlasReaderCache as ReaderCache, get_or_open_atlas}; -pub use options::AtlasOptions as Options; -pub use source::AtlasSource as Source; +pub use options::AtlasOptions; +pub use source::{AtlasEntry, AtlasSource}; pub use table_function::ReadAtlasFunc; -/// Runtime configuration for the atlas format. -/// -/// Plain data with sensible defaults; the caller populates it. The reader-cache -/// capacity is a shared runtime resource, while `use_reader_cache` is a default -/// a table can override via `CREATE EXTERNAL TABLE ... OPTIONS (...)`. -#[derive(Debug, Clone)] -pub struct AtlasConfig { - /// Whether reads consult the shared reader cache by default. - pub use_reader_cache: bool, - /// Capacity (number of opened atlas stores) of the shared reader cache. - pub reader_cache_size: u64, - /// Whether a predicate scan prunes datasets that can't match, using the - /// collection's statistics, before reading them. A pure optimization — off - /// only trades throughput for skipping the pruning-index build. Overridable - /// per table via `CREATE EXTERNAL TABLE ... OPTIONS (use_pruning '…')`. - pub use_pruning: bool, -} - -impl Default for AtlasConfig { - fn default() -> Self { - Self { - use_reader_cache: true, - reader_cache_size: 32, - use_pruning: true, - } - } -} - -/// Split a store's dataset names into up to `partitions` round-robin buckets. -/// -/// Round-robin (`i % parts`) rather than contiguous chunks keeps the buckets -/// balanced when nearby ordinals have similar sizes (a common ingest pattern). -/// Empty buckets are dropped, so a store with fewer datasets than `partitions` -/// simply yields fewer scan partitions. -fn partition_dataset_names(names: Vec, partitions: usize) -> Vec> { - if names.is_empty() { - return Vec::new(); - } - let parts = partitions.max(1).min(names.len()); - let mut buckets: Vec> = vec![Vec::new(); parts]; - for (i, name) in names.into_iter().enumerate() { - buckets[i % parts].push(name); - } - buckets -} +/// The name this format answers to: `STORED AS ATLAS`, `read_atlas`. +pub const ATLAS_FORMAT: &str = "atlas"; -/// Parse a boolean value supplied through a `CREATE EXTERNAL TABLE` option. -fn parse_bool_option(key: &str, value: &str) -> datafusion::error::Result { +/// Parse a boolean supplied through `CREATE EXTERNAL TABLE ... OPTIONS`. +fn parse_bool_option(key: &str, value: &str) -> Result { match value.trim().to_ascii_lowercase().as_str() { "true" | "1" | "yes" | "on" => Ok(true), "false" | "0" | "no" | "off" => Ok(false), @@ -105,11 +67,14 @@ fn parse_bool_option(key: &str, value: &str) -> datafusion::error::Result // ─── Factory ───────────────────────────────────────────────────────────────── +/// Builds an [`AtlasFormat`] per table, over one runtime's settings and one +/// shared reader cache. #[derive(Debug, Clone)] pub struct AtlasFormatFactory { pub options: AtlasOptions, pub config: AtlasConfig, - /// Shared reader cache for this runtime, sized from `config`. + /// The runtime's reader cache, sized from `config` and shared by every + /// format, source and opener this factory builds. cache: AtlasReaderCache, } @@ -123,54 +88,67 @@ impl AtlasFormatFactory { } } - /// Build an [`AtlasFormat`] with the given per-table effective settings, - /// wiring in the shared reader cache when caching is enabled. - fn build_format( + /// A format with this table's effective settings, wired to the shared cache + /// when caching is on. + fn build( &self, options: AtlasOptions, use_reader_cache: bool, use_pruning: bool, ) -> AtlasFormat { - let cache = use_reader_cache.then(|| self.cache.clone()); AtlasFormat::new(options) - .with_cache(cache) + .with_cache(use_reader_cache.then(|| self.cache.clone())) .with_pruning(use_pruning) } + + /// Whether this table wants its columns measured at all. + /// + /// Only the file analyzer measures a collection, through + /// [`FileFormatFactoryExt::create_for_analysis`]. This is the switch that + /// turns even that off, per table or per runtime. + fn statistics_wanted(&self, format_options: &HashMap) -> Result { + match format_option(format_options, "enable_statistics") { + Some(value) => parse_bool_option("enable_statistics", value), + None => Ok(self.config.enable_statistics), + } + } } impl FileFormatFactory for AtlasFormatFactory { fn create( &self, _state: &dyn Session, - format_options: &std::collections::HashMap, - ) -> datafusion::error::Result> { - // Per-table overrides from `CREATE EXTERNAL TABLE ... OPTIONS (...)`, - // defaulting to the runtime config. + format_options: &HashMap, + ) -> Result> { let mut options = self.options.clone(); let mut use_reader_cache = self.config.use_reader_cache; let mut use_pruning = self.config.use_pruning; - if let Some(value) = format_options.get("read_dimensions") { + if let Some(value) = format_option(format_options, "read_dimensions") { options.read_dimensions = Some( value .split(',') - .map(|s| s.trim().to_string()) - .filter(|s| !s.is_empty()) + .map(|dimension| dimension.trim().to_string()) + .filter(|dimension| !dimension.is_empty()) .collect(), ); } - if let Some(value) = format_options.get("use_reader_cache") { + if let Some(value) = format_option(format_options, "use_reader_cache") { use_reader_cache = parse_bool_option("use_reader_cache", value)?; } - if let Some(value) = format_options.get("use_pruning") { + if let Some(value) = format_option(format_options, "use_pruning") { use_pruning = parse_bool_option("use_pruning", value)?; } + // Parsed here only so a bad value is an error at `CREATE EXTERNAL + // TABLE` rather than at the first analysis pass. A query measures + // nothing whatever it says: see `create_for_analysis`. + self.statistics_wanted(format_options)?; - Ok(Arc::new(self.build_format(options, use_reader_cache, use_pruning))) + Ok(Arc::new(self.build(options, use_reader_cache, use_pruning))) } fn default(&self) -> Arc { - Arc::new(self.build_format( + Arc::new(self.build( self.options.clone(), self.config.use_reader_cache, self.config.use_pruning, @@ -184,85 +162,145 @@ impl FileFormatFactory for AtlasFormatFactory { impl GetExt for AtlasFormatFactory { fn get_ext(&self) -> String { - "atlas".to_string() + ATLAS_FORMAT.to_string() } } impl FileFormatFactoryExt for AtlasFormatFactory { - /// Atlas opts into the schema cache for a store read whole. + /// One dataset entry per collection, named by its container object. + /// + /// A collection's datasets are enumerated at plan time, not here: a listing + /// of a data lake would otherwise open every collection it found. + fn discover_datasets(&self, objects: &[ObjectMeta]) -> Result> { + let format = self.get_ext(); + Ok(top_level_atlas_markers(objects) + .into_iter() + .map(|marker| DatasetMetadata::new(marker.location.to_string(), format.clone())) + .collect()) + } + + fn file_format_name(&self) -> String { + self.get_ext() + } + + /// One schema per collection, not per object. /// - /// TODO(#367): cache a dimension-projected read as well. `read_dimensions` - /// decides which arrays are kept, so the same store has one schema per - /// dimension set, and the key would have to carry the set in order. It is - /// left out of this first pass to keep the key simple, and to keep the four - /// nd formats saying the same thing: a read that names dimensions derives - /// its schema, exactly as it did before the cache existed. + /// `infer_schema` reads the container and derives the schema from the + /// footer inside it, so the entry is keyed on the container and depends on + /// everything beside it — the deletion mask included, which changes which + /// datasets the schema covers. + fn schema_units(&self, objects: &[ObjectMeta]) -> Vec { + units_over_stores(objects, &top_level_atlas_markers(objects)) + } + + /// Atlas opts into the schema cache for a collection read whole. + /// + /// TODO(#367): cache a dimension-projected read too. `read_dimensions` + /// decides which arrays survive, so one collection has one schema per + /// dimension set and the key would have to carry the set in order. Left out + /// of this pass to keep the four nd formats saying the same thing. fn schema_options_fingerprint(&self, format: &dyn FileFormat) -> Option { let format = format.as_any().downcast_ref::()?; if format.options.read_dimensions.is_some() { return None; } - Some(SchemaOptions::new("atlas").finish()) + Some(SchemaOptions::new(ATLAS_FORMAT).finish()) } - /// One schema per store, not per object. See the Zarr factory for the same - /// reasoning: `infer_schema` reads the marker at a store's root and derives - /// the schema from the collection behind it, so the entry is keyed on the - /// marker and depends on everything under it. - fn schema_units(&self, objects: &[ObjectMeta]) -> Vec { - units_over_stores(objects, &crate::util::top_level_atlas_markers(objects)) - } - - fn discover_datasets( + /// The same format, with the column measurement switched on. + /// + /// `infer_stats` folds a collection's footer, which is cheap but not free + /// over a listing of thousands. Only the file analyzer asks for it, and a + /// scan prunes from what that recorded. See + /// [`FileFormatFactoryExt::create_for_analysis`]. + fn create_for_analysis( &self, - objects: &[ObjectMeta], - ) -> datafusion::error::Result> { - // One dataset entry per top-level store marker (mirroring zarr, which - // emits one entry per top-level `zarr.json`). The store's individual - // datasets are enumerated at scan time by the opener. - let ext = self.get_ext(); - Ok(top_level_atlas_markers(objects) - .into_iter() - .map(|marker| DatasetMetadata::new(marker.location.to_string(), ext.clone())) - .collect()) - } - - fn file_format_name(&self) -> String { - self.get_ext() + state: &dyn Session, + format_options: &HashMap, + url: &ListingTableUrl, + listing: &ListingFactory, + ) -> Result> { + let wanted = self.statistics_wanted(format_options)?; + let format = self.create_with_native_root(state, format_options, url, listing)?; + let atlas = format + .as_any() + .downcast_ref::() + .ok_or_else(|| { + exec_datafusion_err!("the atlas factory did not produce an AtlasFormat") + })? + .clone(); + Ok(Arc::new(atlas.with_enable_statistics(wanted))) } } // ─── Format ────────────────────────────────────────────────────────────────── -#[derive(Debug, Clone, Default)] +/// Reads one table's worth of Atlas collections. +#[derive(Debug, Clone)] pub struct AtlasFormat { pub options: AtlasOptions, - /// Reader cache to consult, or `None` to bypass caching for this format. + /// The reader cache to consult, or `None` to bypass caching. cache: Option, - /// Whether a predicate scan prunes non-matching datasets before reading. + /// Whether a predicate scan drops the datasets it can rule out. use_pruning: bool, + /// Whether [`FileFormat::infer_stats`] measures a collection's columns. + enable_statistics: bool, +} + +impl Default for AtlasFormat { + fn default() -> Self { + Self::new(AtlasOptions::default()) + } } impl AtlasFormat { pub fn new(options: AtlasOptions) -> Self { + let defaults = AtlasConfig::default(); Self { options, cache: None, - use_pruning: false, + // A query prunes by default: it only ever saves reads. + use_pruning: defaults.use_pruning, + // A query measures nothing. Only the analyzer asks, through + // `create_for_analysis`. + enable_statistics: false, } } - /// Wire in a reader cache (`Some`) or disable caching (`None`). + /// Wire in a reader cache (`Some`), or bypass caching (`None`). pub fn with_cache(mut self, cache: Option) -> Self { self.cache = cache; self } - /// Enable or disable dataset pruning for predicate scans. + /// Drop the datasets a predicate rules out, or read them all. pub fn with_pruning(mut self, use_pruning: bool) -> Self { self.use_pruning = use_pruning; self } + + /// Measure a collection's columns in [`FileFormat::infer_stats`], or report + /// them unknown. + pub fn with_enable_statistics(mut self, enable_statistics: bool) -> Self { + self.enable_statistics = enable_statistics; + self + } +} + +/// Wrap a scan in the nd spine: `NdBroadcastExec` over `NdSourceExec` over the +/// scan. +/// +/// The scan carries its columns `beacon.nd`-encoded, one chunk per row, so +/// `NdSourceExec` decodes them and `NdBroadcastExec` broadcasts them back onto +/// the logical table schema above. +pub fn nd_scan_plan(conf: FileScanConfig) -> Result> { + let scan: Arc = DataSourceExec::from_data_source(conf); + let nd_source = Arc::new(beacon_datafusion_ext::nd::exec::NdSourceExec::try_new( + scan, + )?); + Ok(Arc::new( + beacon_datafusion_ext::nd::exec::NdBroadcastExec::try_new(nd_source)?, + )) } #[async_trait::async_trait] @@ -275,6 +313,7 @@ impl FileFormat for AtlasFormat { None } + /// The container's own name, which is what a listing matches on. fn get_ext(&self) -> String { ATLAS_MARKER.to_string() } @@ -282,118 +321,147 @@ impl FileFormat for AtlasFormat { fn get_ext_with_compression( &self, _file_compression_type: &FileCompressionType, - ) -> datafusion::error::Result { + ) -> Result { Ok(ATLAS_MARKER.to_string()) } + /// The schema of every collection in the listing, merged. + /// + /// Each collection costs one open — a footer read — and the datasets behind + /// it cost no I/O at all. Never enumerate a collection's datasets any other + /// way here: at a million datasets that would turn planning into a scan. async fn infer_schema( &self, state: &dyn Session, store: &Arc, objects: &[ObjectMeta], - ) -> datafusion::error::Result { - let infer_start = std::time::Instant::now(); + ) -> Result { + let started = std::time::Instant::now(); let markers = top_level_atlas_markers(objects); if markers.is_empty() { - return Ok(Arc::new(arrow::datatypes::Schema::empty())); + return Ok(Arc::new(Schema::empty())); } - // Scale note: the schema is derived from each store's collection-wide - // `merged_schema()` — a pre-widened, in-memory summary that costs O(1) - // disk reads (just the metadata already loaded on open), independent of - // the dataset count. Never iterate `list_datasets()` here: at 1M+ - // datasets that turns planning into a full-collection scan. + // One rule for both merges: the datasets inside a collection, and the + // collections of this table. + let widening = session_widening(state); let read_dimensions = self.options.read_dimensions.clone(); - let mut schemas = Vec::new(); + + let mut schemas = Vec::with_capacity(markers.len()); for marker in &markers { - let atlas = get_or_open_atlas(self.cache.as_ref(), store.clone(), marker).await?; - let merged = atlas.merged_schema(); - let schema = - crate::compat::atlas_merged_schema_to_arrow(&merged, read_dimensions.as_deref()); - // The marker names the store, so a refused column names both stores. - schemas.push(LabeledSchema::new( - Arc::new(schema), - marker.location.as_ref(), - )); + let atlas = get_or_open_atlas(self.cache.as_ref(), Arc::clone(store), marker) + .await + .map_err(|e| exec_datafusion_err!("{e}"))?; + let label = marker.location.as_ref(); + let schema = collection_schema(&atlas, read_dimensions.as_deref(), label, &widening) + .await + .map_err(|e| exec_datafusion_err!("{e}"))?; + // The container names the schema, so a refused column names both + // collections. + schemas.push(LabeledSchema::new(schema, label)); } - // Union the stores with the rule of the session. One store is the common - // case, and it gives one schema. - let schema = session_widening(state).merge_schemas(&schemas).map_err(|e| { - exec_datafusion_err!("Failed to merge the schemas of the atlas stores: {}", e) + let schema = widening.merge_schemas(&schemas).map_err(|e| { + exec_datafusion_err!("Failed to merge the schemas of the atlas collections: {e}") })?; tracing::debug!( - elapsed_ms = infer_start.elapsed().as_millis() as u64, - stores = markers.len(), + elapsed_ms = started.elapsed().as_millis() as u64, + collections = markers.len(), fields = schema.fields().len(), "atlas infer_schema", ); Ok(schema) } + /// The column ranges of one collection, folded out of its footer. + /// + /// Reporting unknown rather than erroring is deliberate throughout: absent + /// statistics are always a legal answer, and they only mean Beacon reads + /// what it might have skipped. A listing also hands this method every + /// object it matched, and only a container has a collection behind it. async fn infer_stats( &self, _state: &dyn Session, - _store: &Arc, + store: &Arc, table_schema: SchemaRef, - _object: &ObjectMeta, - ) -> datafusion::error::Result { - Ok(Statistics::new_unknown(&table_schema)) + object: &ObjectMeta, + ) -> Result { + if !self.enable_statistics || !is_atlas_marker(object) { + return Ok(Statistics::new_unknown(&table_schema)); + } + + match get_or_open_atlas(self.cache.as_ref(), Arc::clone(store), object).await { + Ok(atlas) => Ok(statistics::collection_statistics(&atlas, &table_schema)), + Err(e) => { + tracing::debug!(object = %object.location, "not measuring this collection: {e}"); + Ok(Statistics::new_unknown(&table_schema)) + } + } } + /// Plan one entry per dataset, then wrap the scan in the nd spine. + /// + /// Opening each collection here is one metadata read, and the openers reuse + /// the same handle through the reader cache. `list_datasets` is in memory. async fn create_physical_plan( &self, state: &dyn Session, conf: FileScanConfig, - ) -> datafusion::error::Result> { - // Spread each store's datasets across up to `target_partitions` file - // groups so DataFusion scans them on separate cores — the single - // partition a plain listing produces would otherwise pin a 1M-dataset - // store to one thread. Opening the store here is one metadata read, - // cached and reused by the openers; `list_datasets` is in-memory. - let plan_start = std::time::Instant::now(); + ) -> Result> { + beacon_nd_array::arrow::morsel::reject_partition_columns("Atlas", &conf)?; + + let started = std::time::Instant::now(); let object_store = state .runtime_env() .object_store(conf.object_store_url.clone())?; - let target_partitions = state.config().target_partitions().max(1); - let mut markers: Vec = Vec::new(); - for group in &conf.file_groups { - for file in group.files() { - markers.push(file.object_meta.clone()); - } - } - let markers = top_level_atlas_markers(&markers); + let listed: Vec = conf + .file_groups + .iter() + .flat_map(|group| group.files()) + .map(|file| file.object_meta.clone()) + .collect(); + let markers = top_level_atlas_markers(&listed); - let mut total_datasets = 0usize; - let mut file_groups: Vec = Vec::new(); + let mut datasets = 0usize; + let mut file_groups: Vec = Vec::with_capacity(markers.len()); for marker in &markers { - let atlas = get_or_open_atlas(self.cache.as_ref(), object_store.clone(), marker).await?; + let atlas = get_or_open_atlas(self.cache.as_ref(), Arc::clone(&object_store), marker) + .await + .map_err(|e| exec_datafusion_err!("{e}"))?; let names = atlas.list_datasets(); - total_datasets += names.len(); - for slice in partition_dataset_names(names, target_partitions) { - // `From` keeps the marker's freshness (last_modified - // + size) so the opener's cache key matches this plan-time open. - let mut file = PartitionedFile::from(marker.clone()); - file.extensions = Some(Arc::new(AtlasDatasetSlice { names: slice })); - file_groups.push(FileGroup::new(vec![file])); - } + datasets += names.len(); + + let entries: Vec = names + .into_iter() + .enumerate() + .map(|(position, dataset)| { + // `From` keeps the container's freshness, so + // the opener's cache key matches this plan-time open. + let mut entry = PartitionedFile::from(marker.clone()); + entry.extensions = Some(Arc::new(AtlasEntry { dataset, position })); + entry + }) + .collect(); + file_groups.push(FileGroup::new(entries)); } tracing::debug!( - elapsed_ms = plan_start.elapsed().as_millis() as u64, - stores = markers.len(), - datasets = total_datasets, - partitions = file_groups.len(), - "atlas create_physical_plan partitioning", + elapsed_ms = started.elapsed().as_millis() as u64, + collections = markers.len(), + datasets, + "atlas create_physical_plan", ); - let table_schema = datafusion::datasource::table_schema::TableSchema::new( - conf.file_schema().clone(), - conf.table_partition_cols().clone(), - ); - // Preserve a projection that the scan pushed down into the incoming - // source — rebuilding the source below would otherwise drop it. + // The scan carries nd columns, so the source's schema is the encoded + // form of the logical table schema. + let encoded = Arc::new(beacon_datafusion_ext::nd::encoded_schema( + conf.file_schema(), + )); + let table_schema = TableSchema::new(encoded, conf.table_partition_cols().clone()); + // Preserve a projection already pushed into the incoming source; + // rebuilding it below would otherwise drop it. let projection = conf.file_source().projection().cloned(); + let source = AtlasSource::new(self.options.read_dimensions.clone(), table_schema) .with_cache(self.cache.clone()) .with_pruning(self.use_pruning) @@ -402,7 +470,16 @@ impl FileFormat for AtlasFormat { .with_file_groups(file_groups) .with_source(Arc::new(source)) .build(); - Ok(DataSourceExec::from_data_source(conf)) + + nd_scan_plan(conf) + } + + fn file_source(&self, table_schema: TableSchema) -> Arc { + Arc::new( + AtlasSource::new(self.options.read_dimensions.clone(), table_schema) + .with_cache(self.cache.clone()) + .with_pruning(self.use_pruning), + ) } async fn create_writer_physical_plan( @@ -411,678 +488,732 @@ impl FileFormat for AtlasFormat { _state: &dyn Session, _conf: FileSinkConfig, _order_requirements: Option, - ) -> datafusion::error::Result> { - Err(datafusion::error::DataFusionError::NotImplemented( - "Writing atlas datasets is not supported".to_string(), + ) -> Result> { + Err(DataFusionError::NotImplemented( + "an atlas collection is written once, by `atlas create`, and Beacon does not write one" + .to_string(), )) } - - fn file_source( - &self, - table_schema: datafusion::datasource::table_schema::TableSchema, - ) -> Arc { - Arc::new( - AtlasSource::new(self.options.read_dimensions.clone(), table_schema) - .with_cache(self.cache.clone()) - .with_pruning(self.use_pruning), - ) - } } #[cfg(test)] -pub(crate) mod test_support { - //! Shared helpers for integration tests across the datafusion module. - - use crate::reader::test_support::build_two_dataset_store; - use crate::util::ATLAS_MARKER; - use object_store::{ObjectMeta, ObjectStore, local::LocalFileSystem, path::Path as OsPath}; - use std::path::PathBuf; - use std::sync::Arc; - - /// Fixture directory (under the crate-local test root) for the store. - pub const FIXTURE_DIR: &str = "two_datasets.atlas"; - - /// The local root fixtures are built under — a crate-local dir in the OS - /// temp dir. These tests need *a* root, not the application's. - fn datasets_root() -> PathBuf { - std::env::temp_dir().join("beacon-arrow-atlas-datasets") +mod tests { + use super::*; + use crate::test_support; + use arrow::datatypes::DataType; + use datafusion::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, + }; + use datafusion::physical_plan::ExecutionPlan; + use datafusion::prelude::{SessionConfig, SessionContext}; + use std::path::Path; + + /// Register the collection in `dir` as `name`, through a listing table. + async fn register(ctx: &SessionContext, dir: &Path, name: &str) { + let format: Arc = Arc::new(AtlasFormat::default()); + register_with(ctx, dir, name, format).await; } - /// Ensure the fixture store exists under the test root. Idempotent and - /// race-free across concurrent `#[tokio::test]` invocations. - async fn ensure_fixture() -> PathBuf { - static FIXTURE: tokio::sync::OnceCell = tokio::sync::OnceCell::const_new(); - FIXTURE - .get_or_init(|| async { - let dst = datasets_root().join(FIXTURE_DIR); - let marker = dst.join(ATLAS_MARKER); - if !marker.exists() { - if dst.exists() { - std::fs::remove_dir_all(&dst).expect("cleanup partial fixture"); - } - std::fs::create_dir_all(&dst).expect("create fixture dir"); - build_two_dataset_store(&dst).await; - } - dst - }) + async fn register_with( + ctx: &SessionContext, + dir: &Path, + name: &str, + format: Arc, + ) { + let directory = dir.to_string_lossy().replace('\\', "/"); + let url = ListingTableUrl::parse(format!("file://{directory}/")).unwrap(); + let listing = ListingOptions::new(format).with_file_extension(ATLAS_MARKER); + let config = ListingTableConfig::new(url) + .with_listing_options(listing) + .infer_schema(&ctx.state()) .await - .clone() + .expect("the collection types"); + ctx.register_table(name, Arc::new(ListingTable::try_new(config).unwrap())) + .unwrap(); } - /// An object store rooted at the test datasets root, plus the ensured - /// fixture. The marker location is `{FIXTURE_DIR}/atlas.json` relative to - /// this store. - pub async fn test_store() -> Arc { - ensure_fixture().await; - Arc::new(LocalFileSystem::new_with_prefix(datasets_root()).unwrap()) + fn context(partitions: usize) -> SessionContext { + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(partitions)) } - /// `ObjectMeta` for the fixture's marker, relative to [`test_store`]. - pub fn fixture_marker_object_meta() -> ObjectMeta { - ObjectMeta { - location: OsPath::from(format!("{FIXTURE_DIR}/{ATLAS_MARKER}")), - last_modified: Default::default(), - size: 0, - e_tag: None, - version: None, - } + async fn rows(ctx: &SessionContext, sql: &str) -> usize { + ctx.sql(sql) + .await + .unwrap() + .collect() + .await + .unwrap() + .iter() + .map(|batch| batch.num_rows()) + .sum() } -} -#[cfg(test)] -mod tests { - use super::test_support::{fixture_marker_object_meta, test_store}; - use super::*; - use datafusion::datasource::file_format::FileFormat; - use datafusion::datasource::listing::{ - ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, - }; - use datafusion::prelude::SessionContext; - use object_store::local::LocalFileSystem; - use object_store::path::Path as OsPath; + async fn count(ctx: &SessionContext, sql: &str) -> i64 { + use arrow::array::Int64Array; + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("a count is an i64") + .value(0) + } // ── discovery ─────────────────────────────────────────────────────── #[test] - fn factory_get_ext_is_atlas() { + fn the_factory_answers_to_atlas() { let factory = AtlasFormatFactory::new(Default::default(), Default::default()); assert_eq!(factory.get_ext(), "atlas"); assert_eq!(factory.file_format_name(), "atlas"); + assert_eq!(AtlasFormat::default().get_ext(), "data.atlas"); } #[test] - fn discover_datasets_emits_one_entry_per_store() { - let factory = AtlasFormatFactory::new(Default::default(), Default::default()); - let objects = vec![ - ObjectMeta { - location: OsPath::from("store_a/atlas.json"), - last_modified: Default::default(), - size: 0, - e_tag: None, - version: None, - }, - // A nested store marker must NOT become its own dataset. - ObjectMeta { - location: OsPath::from("store_a/inner/atlas.json"), - last_modified: Default::default(), - size: 0, - e_tag: None, - version: None, - }, + fn one_dataset_entry_per_collection() { + fn object(path: &str) -> ObjectMeta { ObjectMeta { - location: OsPath::from("store_b/atlas.msgpack"), + location: object_store::path::Path::from(path), last_modified: Default::default(), size: 0, e_tag: None, version: None, - }, - ]; - let mut datasets = factory.discover_datasets(&objects).expect("discover"); - datasets.sort_by(|a, b| a.file_path.cmp(&b.file_path)); - assert_eq!(datasets.len(), 2, "{datasets:?}"); - assert_eq!(datasets[0].file_path, "store_a/atlas.json"); - assert_eq!(datasets[1].file_path, "store_b/atlas.msgpack"); - assert!(datasets.iter().all(|d| d.format == "atlas")); - } + } + } - #[test] - fn discover_datasets_ignores_non_markers() { let factory = AtlasFormatFactory::new(Default::default(), Default::default()); - let objects = vec![ObjectMeta { - location: OsPath::from("some/other.nc"), - last_modified: Default::default(), - size: 0, - e_tag: None, - version: None, - }]; - assert!(factory.discover_datasets(&objects).unwrap().is_empty()); - } - - // ── schema inference over the object store ────────────────────────── - - #[tokio::test] - async fn infer_schema_unions_columns_across_datasets() { - let store = test_store().await; - let format = AtlasFormat::default(); - let ctx = SessionContext::new(); - - let schema = format - .infer_schema(&ctx.state(), &store, &[fixture_marker_object_meta()]) - .await - .expect("infer"); + let discovered = factory + .discover_datasets(&[ + object("a/data.atlas"), + object("a/deleted.mask"), + object("b/data.atlas"), + object("b/notes.txt"), + ]) + .unwrap(); - let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - for expected in ["temperature", "cycle", "season", "year"] { - assert!(names.contains(&expected), "missing {expected} in {names:?}"); - } + let paths: Vec<&str> = discovered.iter().map(|d| d.file_path.as_str()).collect(); + assert_eq!(paths, vec!["a/data.atlas", "b/data.atlas"]); + assert!(discovered.iter().all(|d| d.format == "atlas")); } - #[tokio::test] - async fn infer_schema_empty_objects_returns_empty_schema() { - let store = test_store().await; - let format = AtlasFormat::default(); - let ctx = SessionContext::new(); - let schema = format - .infer_schema(&ctx.state(), &store, &[]) - .await - .expect("infer"); - assert!(schema.fields().is_empty()); - } + // ── reading, end to end ───────────────────────────────────────────── #[tokio::test] - async fn file_source_returns_atlas_type() { - let format = AtlasFormat::default(); - let source = format.file_source( - datafusion::datasource::table_schema::TableSchema::from_file_schema(Arc::new( - arrow::datatypes::Schema::empty(), - )), - ); - assert_eq!(source.file_type(), "atlas"); - } - - // ── end-to-end through DataFusion + ListingTable ──────────────────── - - /// Register the fixture store as a table backed by [`AtlasFormat`] over a - /// `file://` object store (a `LocalFileSystem` DataFusion supplies). - async fn register_example(ctx: &SessionContext) { - // Ensure the fixture exists on disk, then point a ListingTable at it. - let _ = test_store().await; - let store_dir = std::env::temp_dir() - .join("beacon-arrow-atlas-datasets") - .join(super::test_support::FIXTURE_DIR); - let store_dir = store_dir.to_string_lossy().replace('\\', "/"); - let table_path = ListingTableUrl::parse(format!("file:///{store_dir}/")).unwrap(); + async fn every_dataset_of_a_collection_is_read() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let ctx = context(1); + register(&ctx, tmp.path(), "obs").await; - let format: Arc = Arc::new(AtlasFormat::default()); - let listing_options = ListingOptions::new(format).with_file_extension("atlas.json"); - let config = ListingTableConfig::new(table_path) - .with_listing_options(listing_options) - .infer_schema(&ctx.state()) - .await - .unwrap(); - let table = ListingTable::try_new(config).unwrap(); - ctx.register_table("atlas_t", Arc::new(table)).unwrap(); + // winter contributes 4 rows and summer 3. + assert_eq!(rows(&ctx, "SELECT temperature FROM obs").await, 7); } #[tokio::test] - async fn reads_all_datasets_through_datafusion() { - let _ = LocalFileSystem::new(); // ensure the local store type is linked - let ctx = SessionContext::new(); - register_example(&ctx).await; + async fn count_star_counts_every_dataset() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let ctx = context(1); + register(&ctx, tmp.path(), "obs").await; - let batches = ctx - .sql("SELECT temperature FROM atlas_t") - .await - .unwrap() - .collect() - .await - .unwrap(); - let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - // winter (4) + summer (3) temperature values. - assert_eq!(rows, 7); + assert_eq!(count(&ctx, "SELECT COUNT(*) FROM obs").await, 7); } - /// The same store through [`FastObjectTable`], which is what `read_atlas` - /// builds. - /// - /// An atlas store is a directory, but the reader never opens one: it takes - /// the `atlas.json` marker and resolves the store from there, reading every - /// dataset when no slice is attached. So the table hands it that object like - /// any other file. Every other test here goes through `ListingTable`, which - /// would not notice if that stopped being true. + /// The plan is the nd spine over the scan, in that nesting order. #[tokio::test] - async fn reads_all_datasets_through_the_fast_object_table() { - use beacon_datafusion_ext::fast_object::FastObjectTable; - use beacon_datafusion_ext::type_widening::ArrowTypeWidening; - use datafusion::execution::SessionStateBuilder; - use datafusion::prelude::SessionConfig; - - let _ = test_store().await; - let store_dir = std::env::temp_dir() - .join("beacon-arrow-atlas-datasets") - .join(super::test_support::FIXTURE_DIR); - let store_dir = store_dir.to_string_lossy().replace('\\', "/"); - let url = ListingTableUrl::parse(format!("file:///{store_dir}/")).unwrap(); - - let state = SessionStateBuilder::new() - .with_config( - SessionConfig::new() - .with_target_partitions(4) - .with_extension(ArrowTypeWidening::default_extension()), - ) - .with_default_features() - .build(); - let ctx = SessionContext::new_with_state(state); + async fn the_plan_is_the_nd_spine_over_the_scan() { + use datafusion::physical_plan::displayable; - let table = FastObjectTable::try_new( - &ctx.state(), - Arc::new(AtlasFormat::default()), - vec![url], - ) - .await - .expect("an atlas store registers as a table"); - ctx.register_table("atlas_fast", Arc::new(table)).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let ctx = context(1); + register(&ctx, tmp.path(), "obs").await; - let batches = ctx - .sql("SELECT temperature FROM atlas_fast") + let plan = ctx + .sql("SELECT temperature FROM obs") .await .unwrap() - .collect() + .create_physical_plan() .await .unwrap(); - let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - // The same winter (4) + summer (3) values the ListingTable path reads. - assert_eq!(rows, 7); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); + + let broadcast = rendered.find("NdBroadcastExec"); + let source = rendered.find("NdSourceExec"); + let scan = rendered.find("DataSourceExec"); + assert!( + broadcast.is_some() && source.is_some() && scan.is_some(), + "the spine must be present:\n{rendered}" + ); + assert!( + broadcast < source && source < scan, + "expected NdBroadcastExec over NdSourceExec over DataSourceExec:\n{rendered}" + ); } #[tokio::test] - async fn projection_prunes_columns_through_datafusion() { - let ctx = SessionContext::new(); - register_example(&ctx).await; + async fn a_projection_reaches_the_result() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let ctx = context(1); + register(&ctx, tmp.path(), "obs").await; - let df = ctx.sql("SELECT temperature FROM atlas_t").await.unwrap(); - let names: Vec = df + let df = ctx.sql("SELECT temperature FROM obs").await.unwrap(); + let columns: Vec = df .schema() .fields() .iter() - .map(|f| f.name().clone()) + .map(|field| field.name().clone()) .collect(); - assert_eq!(names, vec!["temperature".to_string()]); + assert_eq!(columns, vec!["temperature".to_string()]); } + /// An attribute rides along as a constant column on every row its dataset + /// contributes. #[tokio::test] - async fn predicate_pushdown_prunes_rows_through_datafusion() { - let ctx = SessionContext::new(); - register_example(&ctx).await; - - let rows: usize = ctx - .sql("SELECT temperature FROM atlas_t WHERE temperature > 1000000") - .await - .unwrap() - .collect() - .await - .unwrap() - .iter() - .map(|b| b.num_rows()) - .sum(); - assert_eq!(rows, 0, "no temperature exceeds 1e6"); - } + async fn an_attribute_is_constant_across_its_datasets_rows() { + use arrow::array::{Array, StringArray}; - #[tokio::test] - async fn count_star_counts_every_dataset_row() { - let ctx = SessionContext::new(); - register_example(&ctx).await; + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let ctx = context(1); + register(&ctx, tmp.path(), "obs").await; - use arrow::array::Int64Array; let batches = ctx - .sql("SELECT COUNT(*) AS n FROM atlas_t") + .sql(r#"SELECT ".season" AS season FROM obs WHERE temperature < 10"#) .await .unwrap() .collect() .await .unwrap(); - let n = batches[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - // winter contributes 4 rows, summer 3. - assert_eq!(n, 7, "COUNT(*) must count rows from every dataset"); - } - // ── cross-dataset dtype widening: cast + null-fill ────────────────── - - /// Register the widening fixture (`a.value: Int16`, `b.value: Float32`, - /// `a.flag: Int32` only) as a table over its own `file://` store. - async fn register_widening(ctx: &SessionContext) -> tempfile::TempDir { - let tmp = tempfile::tempdir().unwrap(); - crate::reader::test_support::build_widening_store(tmp.path()).await; - let dir = tmp.path().to_string_lossy().replace('\\', "/"); - let table_path = ListingTableUrl::parse(format!("file:///{dir}/")).unwrap(); - let format: Arc = Arc::new(AtlasFormat::default()); - let listing_options = ListingOptions::new(format).with_file_extension("atlas.json"); - let config = ListingTableConfig::new(table_path) - .with_listing_options(listing_options) - .infer_schema(&ctx.state()) - .await - .unwrap(); - ctx.register_table("w", Arc::new(ListingTable::try_new(config).unwrap())) - .unwrap(); - tmp + let mut seen = Vec::new(); + for batch in &batches { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("a text column"); + for row in 0..column.len() { + seen.push(column.value(row).to_string()); + } + } + // Only winter's four rows are below 10 degrees. + assert_eq!(seen, vec!["winter".to_string(); 4]); } + // ── datasets that disagree ────────────────────────────────────────── + #[tokio::test] - async fn widened_array_dtype_is_cast_from_each_dataset() { - use arrow::array::Float32Array; - use arrow::datatypes::DataType; + async fn a_widened_column_is_cast_from_each_dataset() { + use arrow::array::Float64Array; - let ctx = SessionContext::new(); - let _tmp = register_widening(&ctx).await; + let tmp = tempfile::tempdir().unwrap(); + test_support::widening(tmp.path()).await; + let ctx = context(1); + register(&ctx, tmp.path(), "w").await; - // Int16 ∪ Float32 widens to Float32 (int16 is exactly representable), - // so the table's `value` column is Float32 and each dataset is cast up. let df = ctx.sql("SELECT value FROM w ORDER BY value").await.unwrap(); assert_eq!( - df.schema().field_with_unqualified_name("value").unwrap().data_type(), - &DataType::Float32, - "merged value column must be the widened super-type" + df.schema() + .field_with_unqualified_name("value") + .unwrap() + .data_type(), + &DataType::Float64 ); let batches = df.collect().await.unwrap(); - let mut vals: Vec = Vec::new(); - for b in &batches { - let col = b.column(0).as_any().downcast_ref::().unwrap(); - for i in 0..col.len() { - vals.push(col.value(i)); - } + let mut values = Vec::new(); + for batch in &batches { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + values.extend(column.iter().flatten()); } - // a.value = [1,2] (Int16 → f32), b.value = [3.5,4.5] (Float32). - assert_eq!(vals, vec![1.0, 2.0, 3.5, 4.5], "each dataset cast up to the super-type"); + // a.value = [1, 2] as Int16, b.value = [3.5, 4.5] as Float32. + assert_eq!(values, vec![1.0, 2.0, 3.5, 4.5]); } + /// A dataset that lacks a projected column contributes its rows with that + /// column null, rather than dropping them. #[tokio::test] - async fn missing_column_is_null_filled_per_dataset() { - let ctx = SessionContext::new(); - let _tmp = register_widening(&ctx).await; + async fn a_column_one_dataset_lacks_is_null_filled() { + let tmp = tempfile::tempdir().unwrap(); + test_support::widening(tmp.path()).await; + let ctx = context(1); + register(&ctx, tmp.path(), "w").await; - // `flag` exists only in dataset `a` (2 rows); dataset `b`'s 2 rows must - // null-fill it. let batches = ctx - .sql("SELECT flag FROM w") + .sql("SELECT value, flag FROM w") .await .unwrap() .collect() .await .unwrap(); - let (rows, nulls): (usize, usize) = batches - .iter() - .map(|b| (b.num_rows(), b.column(0).null_count())) - .fold((0, 0), |(r, n), (br, bn)| (r + br, n + bn)); - assert_eq!(rows, 4, "both datasets contribute rows"); - assert_eq!(nulls, 2, "dataset b's rows null-fill the missing flag column"); + let (rows, nulls) = batches.iter().fold((0, 0), |(rows, nulls), batch| { + ( + rows + batch.num_rows(), + nulls + batch.column(1).null_count(), + ) + }); + assert_eq!(rows, 4, "both datasets contribute their rows"); + assert_eq!(nulls, 2, "dataset b declares no flag"); } - /// A collection whose datasets give the same array *non-numeric* conflicting - /// dtypes still reads: atlas widens `String` ∪ `Int64` to `String`, and the - /// integer dataset is cast into `Utf8` rather than the scan failing. Guards - /// the assumption that every merged dtype is castable-into from each - /// dataset's native type. - #[tokio::test] - async fn incompatible_dtype_union_reads_both_datasets_as_strings() { - use arrow::array::{Array, StringArray}; - use arrow::datatypes::DataType; + // ── the deletion mask ─────────────────────────────────────────────── - let ctx = SessionContext::new(); + /// A deleted dataset is gone from the result, and its rows with it. + #[tokio::test] + async fn a_deleted_dataset_is_not_read() { let tmp = tempfile::tempdir().unwrap(); - crate::reader::test_support::build_incompatible_store(tmp.path()).await; - let dir = tmp.path().to_string_lossy().replace('\\', "/"); - let table_path = ListingTableUrl::parse(format!("file:///{dir}/")).unwrap(); - let format: Arc = Arc::new(AtlasFormat::default()); - let listing_options = ListingOptions::new(format).with_file_extension("atlas.json"); - let config = ListingTableConfig::new(table_path) - .with_listing_options(listing_options) - .infer_schema(&ctx.state()) + test_support::two_datasets(tmp.path()).await; + test_support::open(tmp.path()) .await - .unwrap(); - ctx.register_table("m", Arc::new(ListingTable::try_new(config).unwrap())) - .unwrap(); - - let df = ctx.sql("SELECT value FROM m ORDER BY value").await.unwrap(); - assert_eq!( - df.schema().field_with_unqualified_name("value").unwrap().data_type(), - &DataType::Utf8, - "String wins the union, so the column is Utf8" - ); + .delete_dataset("winter") + .await + .expect("delete winter"); - let batches = df.collect().await.unwrap(); - let mut vals: Vec = Vec::new(); - for b in &batches { - let col = b.column(0).as_any().downcast_ref::().unwrap(); - for i in 0..col.len() { - vals.push(col.value(i).to_string()); - } - } - // a.value = ["x","y"] (native String); b.value = [1,2] (Int64, stringified). - assert_eq!(vals, vec!["1", "2", "x", "y"], "both datasets contribute, integers cast to text"); + let ctx = context(1); + register(&ctx, tmp.path(), "obs").await; + // Summer's three rows alone. + assert_eq!(rows(&ctx, "SELECT temperature FROM obs").await, 3); } - // ── partition splitting ───────────────────────────────────────────── + // ── dividing the scan ─────────────────────────────────────────────── - #[test] - fn partition_dataset_names_round_robins_and_drops_empties() { - let names: Vec = (0..5).map(|i| format!("d{i}")).collect(); - let buckets = partition_dataset_names(names, 3); - assert_eq!(buckets.len(), 3); - // round-robin: [d0,d3], [d1,d4], [d2] - assert_eq!(buckets[0], vec!["d0", "d3"]); - assert_eq!(buckets[1], vec!["d1", "d4"]); - assert_eq!(buckets[2], vec!["d2"]); - - // Fewer datasets than partitions → one bucket per dataset, no empties. - let two = partition_dataset_names(vec!["x".into(), "y".into()], 8); - assert_eq!(two.len(), 2); - - assert!(partition_dataset_names(Vec::new(), 4).is_empty()); + /// Every row is read exactly once, however many partitions share the + /// collection. A dataset popped twice is a row returned twice, and one + /// popped by nobody is a row lost; neither raises an error. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn a_partitioned_scan_reads_every_row_once() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 12).await; + + for partitions in [1_usize, 2, 4, 8] { + let ctx = context(partitions); + register(&ctx, tmp.path(), "ranged").await; + assert_eq!( + rows(&ctx, "SELECT temperature FROM ranged").await, + 48, + "partitions={partitions}: 12 datasets of 4 rows" + ); + assert_eq!( + count(&ctx, "SELECT COUNT(*) FROM ranged").await, + 48, + "partitions={partitions}: and the count agrees" + ); + } } - // ── dataset pruning toggle ────────────────────────────────────────── + /// The scan is planned across every partition, through the queue. + #[tokio::test] + async fn a_collection_is_planned_across_every_partition() { + use datafusion::physical_plan::ExecutionPlanProperties; - /// Register the ranged fixture (`d{i}.temperature ∈ [10i, 10i+3]`, `n` - /// datasets) with pruning `use_pruning` on or off. - async fn register_ranged(ctx: &SessionContext, n: usize, use_pruning: bool) -> tempfile::TempDir { let tmp = tempfile::tempdir().unwrap(); - crate::reader::test_support::build_ranged_store(tmp.path(), n).await; - let dir = tmp.path().to_string_lossy().replace('\\', "/"); - let table_path = ListingTableUrl::parse(format!("file:///{dir}/")).unwrap(); - let format: Arc = - Arc::new(AtlasFormat::default().with_pruning(use_pruning)); - let listing_options = ListingOptions::new(format).with_file_extension("atlas.json"); - let config = ListingTableConfig::new(table_path) - .with_listing_options(listing_options) - .infer_schema(&ctx.state()) - .await - .unwrap(); - ctx.register_table("ranged", Arc::new(ListingTable::try_new(config).unwrap())) - .unwrap(); - tmp - } + test_support::ranged(tmp.path(), 12).await; + let ctx = context(4); + register(&ctx, tmp.path(), "ranged").await; - async fn ranged_row_count(use_pruning: bool, predicate: &str) -> usize { - let ctx = SessionContext::new(); - let _tmp = register_ranged(&ctx, 10, use_pruning).await; - ctx.sql(&format!("SELECT temperature FROM ranged WHERE {predicate}")) + let plan = ctx + .sql("SELECT temperature FROM ranged") .await .unwrap() - .collect() + .create_physical_plan() .await - .unwrap() - .iter() - .map(|b| b.num_rows()) - .sum() + .unwrap(); + + // The count comes off the scan, not the plan root: DataFusion adds a + // round robin above a single-partition scan either way. + let mut scan = Arc::clone(&plan); + while let Some(child) = scan.children().first() { + scan = Arc::clone(child); + } + assert_eq!( + scan.output_partitioning().partition_count(), + 4, + "the datasets divide over the partitions:\n{}", + datafusion::physical_plan::displayable(plan.as_ref()).indent(false) + ); } + // ── predicates ────────────────────────────────────────────────────── + #[tokio::test] - async fn pruning_matches_unpruned_results() { - // `> 45` keeps d5..d9 → 5 datasets × 4 rows = 20 rows, whether or not - // pruning is on. Pruning only changes which datasets get opened. - assert_eq!(ranged_row_count(true, "temperature > 45").await, 20); - assert_eq!(ranged_row_count(false, "temperature > 45").await, 20); - - // An impossible predicate → 0 rows both ways. - assert_eq!(ranged_row_count(true, "temperature > 100000").await, 0); - assert_eq!(ranged_row_count(false, "temperature > 100000").await, 0); + async fn a_predicate_keeps_only_the_rows_that_match() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 10).await; + let ctx = context(2); + register(&ctx, tmp.path(), "ranged").await; + + // d5..d9 hold [50..53] … [90..93]: 20 rows above 45. + assert_eq!( + rows( + &ctx, + "SELECT temperature FROM ranged WHERE temperature > 45" + ) + .await, + 20 + ); + assert_eq!( + rows( + &ctx, + "SELECT temperature FROM ranged WHERE temperature > 100000" + ) + .await, + 0, + "a predicate nothing meets returns nothing" + ); } - #[tokio::test] - async fn pruning_on_mixed_dtype_column_end_to_end() { - // End-to-end proof that pruning casts a mixed-dtype column to the merged - // table type before filtering: `value` is Int16 in `a`, Float32 in `b` - // (merged Float32). `value > 3` prunes `a` and keeps `b` — and the - // result is byte-identical with pruning on and off. - use arrow::array::Float32Array; + // ── pruning, end to end ───────────────────────────────────────────── - async fn values(use_pruning: bool) -> Vec { - let tmp = tempfile::tempdir().unwrap(); - crate::reader::test_support::build_widening_store(tmp.path()).await; - let dir = tmp.path().to_string_lossy().replace('\\', "/"); - let table_path = ListingTableUrl::parse(format!("file:///{dir}/")).unwrap(); - let format: Arc = - Arc::new(AtlasFormat::default().with_pruning(use_pruning)); - let listing_options = ListingOptions::new(format).with_file_extension("atlas.json"); - let ctx = SessionContext::new(); - let config = ListingTableConfig::new(table_path) - .with_listing_options(listing_options) - .infer_schema(&ctx.state()) - .await - .unwrap(); - ctx.register_table("w", Arc::new(ListingTable::try_new(config).unwrap())) - .unwrap(); + /// Register the collection twice, once pruning and once not. + async fn register_pruning(ctx: &SessionContext, dir: &Path, name: &str, use_pruning: bool) { + let format: Arc = + Arc::new(AtlasFormat::default().with_pruning(use_pruning)); + register_with(ctx, dir, name, format).await; + } - let batches = ctx - .sql("SELECT value FROM w WHERE value > 3 ORDER BY value") - .await - .unwrap() - .collect() - .await - .unwrap(); - let mut out = Vec::new(); - for b in &batches { - let col = b.column(0).as_any().downcast_ref::().unwrap(); - for i in 0..col.len() { - out.push(col.value(i)); - } - } - out + async fn values(ctx: &SessionContext, sql: &str) -> Vec { + use arrow::array::Float32Array; + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + let mut out = Vec::new(); + for batch in &batches { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("a float column"); + out.extend(column.iter().flatten()); } + out + } - // Only b's values exceed 3; a's [1,2] (Int16, cast to f32) are pruned out. - assert_eq!(values(true).await, vec![3.5, 4.5]); - assert_eq!(values(false).await, vec![3.5, 4.5], "toggle must not change results"); + /// The switch changes what is read, never what is returned. + /// + /// This is the property pruning has to hold above all others: it drops + /// datasets that cannot contain a matching row, and nothing else. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn pruning_does_not_change_the_answer() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 10).await; + + for predicate in [ + "temperature > 45", + "temperature < 25", + "temperature > 1000", + "temperature >= 0", + "temperature > 45 AND temperature < 75", + ] { + let sql = + format!("SELECT temperature FROM ranged WHERE {predicate} ORDER BY temperature"); + + let on = context(4); + register_pruning(&on, tmp.path(), "ranged", true).await; + let off = context(4); + register_pruning(&off, tmp.path(), "ranged", false).await; + + assert_eq!( + values(&on, &sql).await, + values(&off, &sql).await, + "pruning changed the answer for `{predicate}`" + ); + } } - #[tokio::test] - async fn pruning_across_many_partitions_is_correct() { - // Many partitions all share one memoized prune result per store; the - // union must still be exactly the matching rows, none dropped or doubled. - use datafusion::prelude::SessionConfig; - let ctx = - SessionContext::new_with_config(SessionConfig::new().with_target_partitions(8)); - let _tmp = register_ranged(&ctx, 10, true).await; - let rows: usize = ctx + /// Find the scan's metrics by the names only this format registers. + fn atlas_metrics( + plan: &Arc, + ) -> Option { + if let Some(metrics) = plan.metrics() + && metrics.sum_by_name("atlas_datasets_scanned").is_some() + { + return Some(metrics); + } + plan.children().into_iter().find_map(atlas_metrics) + } + + /// The scan reports what it read and what it skipped, and the two add up to + /// the collection. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn the_metrics_report_what_was_read_and_what_was_skipped() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 10).await; + let ctx = context(4); + register_pruning(&ctx, tmp.path(), "ranged", true).await; + + let plan = ctx .sql("SELECT temperature FROM ranged WHERE temperature > 45") .await .unwrap() - .collect() + .create_physical_plan() .await - .unwrap() - .iter() - .map(|b| b.num_rows()) - .sum(); - assert_eq!(rows, 20, "d5..d9 × 4 rows, across 8 partitions"); + .unwrap(); + datafusion::physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()) + .await + .unwrap(); + + let metrics = atlas_metrics(&plan).expect("the atlas scan reports metrics"); + let sum = |name: &str| metrics.sum_by_name(name).map(|value| value.as_usize()); + + // d5..d9 hold values above 45; d0..d4 cannot. + assert_eq!(sum("atlas_datasets_scanned"), Some(5)); + assert_eq!(sum("atlas_datasets_pruned"), Some(5)); + assert_eq!( + sum("atlas_index_rows"), + Some(10), + "the index covered them all" + ); + assert!(metrics.sum_by_name("atlas_prune_time").is_some()); } - #[tokio::test] - async fn scan_metrics_report_pruned_and_scanned_counts() { - use datafusion::physical_plan::metrics::MetricsSet; - use datafusion::physical_plan::{ExecutionPlan, collect}; + /// One index per collection, however many partitions share it. + /// + /// Every partition's opener holds the same memo, so the first to reach the + /// collection builds the index and the rest await it. Without that, a + /// twenty-four-partition scan would build it twenty-four times. + #[tokio::test(flavor = "multi_thread", worker_threads = 8)] + async fn one_index_is_built_per_collection() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 24).await; + let ctx = context(8); + register_pruning(&ctx, tmp.path(), "ranged", true).await; - let ctx = SessionContext::new(); - let _tmp = register_ranged(&ctx, 10, true).await; let plan = ctx - .sql("SELECT temperature FROM ranged WHERE temperature > 45") + .sql("SELECT temperature FROM ranged WHERE temperature > 100") .await .unwrap() .create_physical_plan() .await .unwrap(); - collect(plan.clone(), ctx.task_ctx()).await.unwrap(); - - // Find the scan node exposing the atlas metrics (aggregated over partitions). - fn find(plan: &Arc) -> Option { - if let Some(m) = plan.metrics() - && m.sum_by_name("atlas_datasets_scanned").is_some() - { - return Some(m); - } - plan.children().into_iter().find_map(find) - } - let m = find(&plan).expect("atlas scan metrics present"); - let sum = |name: &str| m.sum_by_name(name).map(|v| v.as_usize()); + datafusion::physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()) + .await + .unwrap(); - // `> 45` keeps d5..d9 (5) and prunes d0..d4 (5), summed across partitions. - assert_eq!(sum("atlas_datasets_scanned"), Some(5)); - assert_eq!(sum("atlas_datasets_pruned"), Some(5)); - // Timers are registered. - assert!(m.sum_by_name("atlas_open_time").is_some()); - assert!(m.sum_by_name("atlas_prune_time").is_some()); - assert!(m.sum_by_name("atlas_dataset_build_time").is_some()); + let metrics = atlas_metrics(&plan).expect("the atlas scan reports metrics"); + assert_eq!( + metrics + .sum_by_name("atlas_index_builds") + .map(|value| value.as_usize()), + Some(1), + "eight partitions must share one index" + ); } - #[tokio::test] - async fn pruning_off_by_option_still_correct() { - // The same table, pruning disabled via the format flag, returns every - // matching row — the toggle must never change results. - let ctx = SessionContext::new(); - let _tmp = register_ranged(&ctx, 10, false).await; - use arrow::array::Float32Array; - let batches = ctx - .sql("SELECT temperature FROM ranged WHERE temperature > 45 ORDER BY temperature") + /// A scan with no predicate builds no index at all. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn a_scan_without_a_predicate_builds_no_index() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 6).await; + let ctx = context(4); + register_pruning(&ctx, tmp.path(), "ranged", true).await; + + let plan = ctx + .sql("SELECT temperature FROM ranged") .await .unwrap() - .collect() + .create_physical_plan() .await .unwrap(); - let mut vals = Vec::new(); - for b in &batches { - let col = b.column(0).as_any().downcast_ref::().unwrap(); - for i in 0..col.len() { - vals.push(col.value(i)); - } + datafusion::physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()) + .await + .unwrap(); + + let metrics = atlas_metrics(&plan).expect("the atlas scan reports metrics"); + assert_eq!( + metrics + .sum_by_name("atlas_index_builds") + .map(|value| value.as_usize()), + Some(0) + ); + assert_eq!( + metrics + .sum_by_name("atlas_datasets_pruned") + .map(|value| value.as_usize()), + Some(0) + ); + } + + // ── measuring a collection ────────────────────────────────────────── + + /// A query never measures a collection, whatever the option says. Only the + /// analyzer asks, through `create_for_analysis`. + #[test] + fn a_query_never_measures_a_collection() { + let factory = AtlasFormatFactory::new(Default::default(), Default::default()); + let ctx = SessionContext::new(); + let on = HashMap::from([("enable_statistics".to_string(), "true".to_string())]); + + for options in [HashMap::new(), on] { + let format = factory.create(&ctx.state(), &options).unwrap(); + assert!( + !format + .as_any() + .downcast_ref::() + .unwrap() + .enable_statistics, + "a format built for a query measures nothing" + ); } - assert!(vals.iter().all(|v| *v > 45.0)); - assert_eq!(vals.len(), 20); } + /// The analyzer layers the per-table option over the runtime default. + #[test] + fn analysis_layers_the_statistics_option_over_the_runtime() { + let measured = |config: AtlasConfig, options: HashMap| { + let ctx = SessionContext::new(); + let listing = Arc::new(ListingFactory::dynamic()); + let url = ListingTableUrl::parse("file:///tmp/").unwrap(); + AtlasFormatFactory::new(Default::default(), config) + .create_for_analysis(&ctx.state(), &options, &url, &listing) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .enable_statistics + }; + let off = HashMap::from([("enable_statistics".to_string(), "false".to_string())]); + let on = HashMap::from([("enable_statistics".to_string(), "yes".to_string())]); + + assert!(measured(AtlasConfig::default(), HashMap::new())); + assert!(!measured(AtlasConfig::default(), off)); + + let disabled = AtlasConfig { + enable_statistics: false, + ..Default::default() + }; + assert!(!measured(disabled.clone(), HashMap::new())); + assert!(measured(disabled, on), "one table can turn them back on"); + } + + // ── dimensions ────────────────────────────────────────────────────── + #[tokio::test] - async fn partitioned_scan_reads_every_dataset_row() { - // With target_partitions > 1 the store's datasets are split across - // partitions; the union of all partitions must still be every row. - use datafusion::prelude::SessionConfig; - let ctx = - SessionContext::new_with_config(SessionConfig::new().with_target_partitions(4)); - register_example(&ctx).await; - - let rows: usize = ctx - .sql("SELECT temperature FROM atlas_t") - .await - .unwrap() - .collect() + async fn read_dimensions_narrow_the_table() { + let tmp = tempfile::tempdir().unwrap(); + test_support::chunked_grid(tmp.path()).await; + let ctx = context(1); + let format: Arc = Arc::new(AtlasFormat::new(AtlasOptions { + read_dimensions: Some(vec!["lat".to_string()]), + })); + register_with(&ctx, tmp.path(), "grid", format).await; + + let columns: Vec = ctx + .table_provider("grid") .await .unwrap() + .schema() + .fields() .iter() - .map(|b| b.num_rows()) - .sum(); - assert_eq!(rows, 7, "partitioned scan must not drop or duplicate rows"); + .map(|field| field.name().clone()) + .collect(); + assert!( + !columns.contains(&"temperature".to_string()), + "a 2-D array does not fit a 1-D grid: {columns:?}" + ); + } + + // ── the table this crate is registered as ─────────────────────────── + + /// The same collection through `FastObjectTable`, which is what + /// `read_atlas` builds. + /// + /// A collection is one file, and the reader takes it as the marker it is; + /// every other test here goes through `ListingTable`, which would not + /// notice if that stopped being true. + #[tokio::test] + async fn a_collection_reads_through_the_fast_object_table() { + use beacon_datafusion_ext::fast_object::FastObjectTable; + use beacon_datafusion_ext::type_widening::ArrowTypeWidening; + use datafusion::execution::SessionStateBuilder; + + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let state = SessionStateBuilder::new() + .with_config( + SessionConfig::new() + .with_target_partitions(4) + .with_extension(ArrowTypeWidening::default_extension()), + ) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + + let directory = tmp.path().to_string_lossy().replace('\\', "/"); + let url = ListingTableUrl::parse(format!("file://{directory}/")).unwrap(); + let table = + FastObjectTable::try_new(&ctx.state(), Arc::new(AtlasFormat::default()), vec![url]) + .await + .expect("a collection registers as a table"); + ctx.register_table("obs", Arc::new(table)).unwrap(); + + assert_eq!(rows(&ctx, "SELECT temperature FROM obs").await, 7); + } + + // ── refusals ──────────────────────────────────────────────────────── + + /// A dataset lives inside a container, not at a path, so no `PARTITIONED + /// BY` value can be read off it. Saying so beats returning the column + /// silently empty. + #[tokio::test] + async fn a_partitioned_table_is_refused_by_name() { + use datafusion::datasource::physical_plan::FileScanConfigBuilder; + use datafusion::execution::object_store::ObjectStoreUrl; + + let table_schema = TableSchema::new( + Arc::new(arrow::datatypes::Schema::empty()), + vec![Arc::new(arrow::datatypes::Field::new( + "year", + DataType::Utf8, + false, + ))], + ); + let source = AtlasSource::new(None, table_schema); + let conf = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + Arc::new(source) as Arc, + ) + .build(); + + let ctx = SessionContext::new(); + let error = AtlasFormat::default() + .create_physical_plan(&ctx.state(), conf) + .await + .expect_err("a partitioned atlas table is refused") + .to_string(); + assert!(error.contains("Atlas"), "{error}"); + assert!(error.contains("year"), "{error}"); + } + + #[test] + fn an_unparseable_option_is_an_error() { + let error = parse_bool_option("use_reader_cache", "maybe") + .unwrap_err() + .to_string(); + assert!(error.contains("use_reader_cache"), "{error}"); + assert!(error.contains("maybe"), "{error}"); + } + + /// A read that names dimensions stays out of the schema cache, because the + /// key does not carry the dimension set. See the `TODO(#367)` above. + #[test] + fn a_dimension_projected_read_is_not_schema_cached() { + let factory = AtlasFormatFactory::new(Default::default(), Default::default()); + assert!( + factory + .schema_options_fingerprint(&AtlasFormat::default()) + .is_some() + ); + assert!( + factory + .schema_options_fingerprint(&AtlasFormat::new(AtlasOptions { + read_dimensions: Some(vec!["time".to_string()]), + })) + .is_none() + ); } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/options.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/options.rs index 04782656..d6b488de 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/options.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/options.rs @@ -1,10 +1,17 @@ -/// Configuration options for the atlas file format in DataFusion queries. -#[derive(Debug, Default, Clone)] +//! Per-query settings of the Atlas format. + +/// Settings that change *what* a scan reads, as opposed to how fast it does so. +/// +/// The runtime settings live in [`AtlasConfig`](crate::AtlasConfig). These come +/// from the query: `read_atlas(paths, dimensions)` sets them, and so does +/// `CREATE EXTERNAL TABLE ... OPTIONS ('read_dimensions' '…')`. +#[derive(Debug, Default, Clone, PartialEq, Eq)] pub struct AtlasOptions { - /// Optional list of dimension names used to filter which arrays are read. + /// The dimensions the table reads, or `None` to pick a broadcast-compatible + /// default per dataset. /// - /// When `Some`, only arrays whose dimensions match the listed dimensions - /// are kept (applied via [`beacon_nd_array::projection::DatasetProjection`]). - /// When `None` (default), all arrays are read. + /// An array survives only when every one of its dimensions is in the list, + /// so this is how a query drops the wide grids of a collection and keeps its + /// coordinates. pub read_dimensions: Option>, } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs index 61fc168b..b88aff90 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs @@ -1,25 +1,35 @@ -//! Optional dataset-level predicate pruning. +//! Dropping the datasets a predicate cannot match, from the collection footer. //! -//! atlas co-locates each column's per-dataset statistics in one `.af` stats -//! file, and [`Atlas::pruning_index`](atlas::Atlas::pruning_index) pivots the -//! requested columns into flat, per-dataset (min, max, null_count, row_count) -//! buffers with a single read per column. Feeding those to DataFusion's -//! [`PruningPredicate`] yields, for an arbitrary predicate, the set of datasets -//! whose ranges could still satisfy it — so a selective `WHERE` over a 1M+ -//! dataset store skips opening the datasets that provably can't match. +//! # One index, not a decision per dataset //! -//! Pruning is a pure optimization: it only ever drops datasets that cannot -//! contain a matching row, and every path **fails open** (returns the full -//! input on any error, unsupported predicate, or missing statistic) so a query -//! can never lose a real row to a pruning hiccup. The not-fully-consumed filter -//! kept above the scan re-checks whatever survives. +//! A collection can hold millions of datasets. Evaluating a predicate against +//! each one in turn would cost millions of evaluations, and each would need its +//! own `DatasetView` — a linear scan of the footer — so the pass would be +//! quadratic before it did any work. +//! +//! Instead the first opener that reaches a collection builds one +//! [`PruningIndex`] over it: one row per live dataset, and one column of typed +//! Arrow statistics per column the predicate names. DataFusion's +//! [`PruningPredicate`] then evaluates the whole collection in one vectorised +//! pass, and the result is a bit per dataset that every partition reads. +//! +//! The statistics come from the footer the open already held, so the build +//! costs no I/O. An array column costs one linear pass through +//! [`Atlas::array_stats_by_dataset`], with no view and no name lookup at all. +//! +//! # Pruning is only ever an optimization +//! +//! Every path here fails open: an error, a predicate the engine cannot use, a +//! column with no statistics, or a bound that will not cast all leave the +//! datasets in. A dataset that survives is still filtered row by row above the +//! scan, so a hiccup here costs time and never a row. -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::Arc; use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; -use arrow::datatypes::{DataType, SchemaRef}; -use atlas::{Atlas, ColumnKey, MergedSchema, PruningIndex, StatVal}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use atlas::{Atlas, Attr, StatValue}; use datafusion::common::Column; use datafusion::common::pruning::PruningStatistics; use datafusion::physical_expr::PhysicalExpr; @@ -27,34 +37,78 @@ use datafusion::physical_expr::utils::collect_columns; use datafusion::physical_optimizer::pruning::PruningPredicate; use datafusion::scalar::ScalarValue; -/// Which of a store's datasets a predicate could match — the result of pruning -/// the whole store once, so every partition's opener can reuse it. -#[derive(Debug, Clone)] +/// Above this many datasets, an attribute column stays out of the index. +/// +/// An array column costs one footer pass. An attribute has no bulk accessor in +/// the reader, so its value needs one `DatasetView` per dataset and each of +/// those is a linear scan — the pass is quadratic. It is worth paying on a +/// collection of thousands and not on one of millions, and the only cost of +/// skipping it is that a predicate on that attribute prunes nothing. +/// +/// `Atlas::attribute_by_dataset` upstream would remove the limit. +const ATTRIBUTE_INDEX_LIMIT: usize = 100_000; + +// ─── What a scan does with the answer ──────────────────────────────────────── + +/// Which datasets of one collection a predicate could still match. +#[derive(Debug)] pub enum CandidateFilter { - /// Pruning didn't apply (fail-open) — keep every dataset. + /// Pruning did not apply. Every dataset is read. KeepAll, - /// Only these dataset names could match; everything else is prunable. - Only(HashSet), + /// One bit per dataset, in the order the plan listed them. + Rows { + kept: BooleanArray, + /// The dataset at each row, so a listing that changed under the plan is + /// detected rather than mis-indexed. + names: Vec, + }, } impl CandidateFilter { - /// Restrict `names` to the datasets that survive pruning, preserving order. - pub fn retain(&self, names: Vec) -> Vec { + /// Whether the dataset at `position` is worth reading. + /// + /// The name is checked against the row it indexes. A collection is + /// immutable, but its deletion mask is not, so a delete between the plan + /// and the open would shift every row after it. A mismatch keeps the + /// dataset: the filter above the scan decides it either way. + pub fn keeps(&self, position: usize, dataset: &str) -> bool { match self { - CandidateFilter::KeepAll => names, - CandidateFilter::Only(set) => { - names.into_iter().filter(|n| set.contains(n)).collect() - } + Self::KeepAll => true, + Self::Rows { kept, names } => match names.get(position) { + Some(name) if name == dataset => kept.value(position), + _ => true, + }, + } + } + + /// Whether an index was built at all, as opposed to pruning not applying. + pub fn is_index(&self) -> bool { + matches!(self, Self::Rows { .. }) + } + + /// How many datasets this filter drops. For diagnostics. + pub fn pruned(&self) -> usize { + match self { + Self::KeepAll => 0, + Self::Rows { kept, .. } => kept.len() - kept.true_count(), + } + } + + /// How many rows the index behind this filter holds. + pub fn rows(&self) -> usize { + match self { + Self::KeepAll => 0, + Self::Rows { kept, .. } => kept.len(), } } } -/// Per-store memo of the [`CandidateFilter`], so a store is pruned **once** per -/// query rather than once per scan partition. +/// Each collection's [`CandidateFilter`], computed once per scan. /// -/// Keyed by marker path; the predicate and schema are fixed for a given source, -/// so the marker identifies the result. Concurrent partition openers coalesce on -/// the shared [`moka`] entry — the first computes, the rest await it. +/// Keyed by the container's path. The predicate and the schema are fixed for a +/// scan, so the container identifies the answer. Every partition's opener holds +/// a clone of this cache, and the clones share one store: the first opener to +/// reach a collection builds the index while the rest await the same future. #[derive(Clone)] pub struct PruneCache { cache: moka::future::Cache>, @@ -67,8 +121,8 @@ impl PruneCache { } } - /// Return the memoized filter for `key`, computing it via `init` on first - /// use. Concurrent callers for the same key share one computation. + /// The memoized filter for `key`, computing it with `init` on first use. + /// Concurrent callers for one key share the one computation. pub async fn get_or_compute(&self, key: String, init: F) -> Arc where F: std::future::Future>, @@ -89,391 +143,724 @@ impl std::fmt::Debug for PruneCache { } } -/// Prune a whole store: the set of dataset names whose statistics leave them -/// able to satisfy `predicate`. Fails open to [`CandidateFilter::KeepAll`] on any -/// error, unsupported predicate, or unmapped column, so a real row is never lost. -/// -/// `schema` must contain every column the predicate references with its table -/// type — the scan's projected output schema does (a not-fully-consumed filter -/// forces its columns into the projection). -pub async fn candidate_set( - atlas: &Arc, - predicate: &Arc, - schema: &SchemaRef, -) -> CandidateFilter { - match try_candidates(atlas, predicate, schema).await { - Ok(Some(set)) => CandidateFilter::Only(set), - Ok(None) | Err(_) => CandidateFilter::KeepAll, +// ─── The index ─────────────────────────────────────────────────────────────── + +/// One column's statistics, one row per dataset. +struct StatColumn { + min: ArrayRef, + max: ArrayRef, + null_count: ArrayRef, + row_count: ArrayRef, +} + +/// A collection's statistics, pivoted into columns of equal length. +struct PruningIndex { + rows: usize, + columns: HashMap, +} + +impl PruningStatistics for PruningIndex { + fn min_values(&self, column: &Column) -> Option { + self.columns.get(column.name()).map(|c| Arc::clone(&c.min)) + } + + fn max_values(&self, column: &Column) -> Option { + self.columns.get(column.name()).map(|c| Arc::clone(&c.max)) + } + + fn null_counts(&self, column: &Column) -> Option { + self.columns + .get(column.name()) + .map(|c| Arc::clone(&c.null_count)) + } + + fn row_counts(&self, column: &Column) -> Option { + self.columns + .get(column.name()) + .map(|c| Arc::clone(&c.row_count)) + } + + fn num_containers(&self) -> usize { + self.rows + } + + fn contained( + &self, + _column: &Column, + _values: &std::collections::HashSet, + ) -> Option { + // An attribute's value is exact, so an `IN` list could prune on one. + // Not yet: every column here reports a range, and a range says nothing + // about membership. + None } } -/// Keep only the datasets in `names` that could satisfy `predicate`. A thin -/// convenience over [`candidate_set`] for one-shot callers and tests. -pub async fn retain_candidates( - atlas: &Arc, - names: Vec, - predicate: &Arc, - schema: &SchemaRef, -) -> Vec { - candidate_set(atlas, predicate, schema).await.retain(names) +// ─── Building it ───────────────────────────────────────────────────────────── + +/// The logical schema behind an nd-encoded one. +/// +/// A scan carries its columns as `beacon.nd` structs, and a predicate is +/// written against the values inside them. A field whose type does not decode +/// keeps its own type, which simply leaves it unprunable. +pub fn logical_schema(encoded: &Schema) -> SchemaRef { + let fields: Vec> = encoded + .fields() + .iter() + .map(|field| { + let value_type = beacon_datafusion_ext::nd::encoding::nd_value_type(field.data_type()) + .unwrap_or_else(|_| field.data_type().clone()); + Arc::new(Field::new(field.name(), value_type, true)) + }) + .collect(); + Arc::new(Schema::new(fields)) } -async fn try_candidates( +/// Which datasets of `atlas` could satisfy `predicate`. +/// +/// `logical_schema` must type every column the predicate names, which the +/// scan's own projected schema does: a filter that stays above the scan forces +/// its columns into the projection. +/// +/// Fails open to [`CandidateFilter::KeepAll`] on anything it cannot prove. +pub async fn candidate_filter( atlas: &Arc, predicate: &Arc, - schema: &SchemaRef, -) -> datafusion::error::Result>> { - let Ok(pruning_predicate) = PruningPredicate::try_new(predicate.clone(), schema.clone()) else { - return Ok(None); // predicate shape unsupported by the pruning engine + logical_schema: &SchemaRef, +) -> CandidateFilter { + let Ok(pruning) = PruningPredicate::try_new(Arc::clone(predicate), Arc::clone(logical_schema)) + else { + // The engine cannot use this predicate shape. + return CandidateFilter::KeepAll; }; - let referenced = collect_columns(pruning_predicate.orig_expr()); + let referenced = collect_columns(pruning.orig_expr()); if referenced.is_empty() { - return Ok(None); + return CandidateFilter::KeepAll; + } + + let names = atlas.list_datasets(); + if names.is_empty() { + return CandidateFilter::KeepAll; } - // Map referenced table columns to atlas column keys, skipping any we can't - // resolve (they simply won't contribute statistics → never prune on them). - let merged = atlas.merged_schema(); - let keyed: Vec<(String, ColumnKey)> = referenced + // The pivot is pure CPU over data already in memory. A million rows is real + // work, so it does not run on the async runtime. + let atlas = Arc::clone(atlas); + let schema = Arc::clone(logical_schema); + let wanted: Vec = referenced .iter() - .filter_map(|col| column_key(&merged, col.name()).map(|k| (col.name().to_string(), k))) + .map(|column| column.name().to_string()) .collect(); - if keyed.is_empty() { - return Ok(None); - } - - // `pruning_index` reads and pivots each column's stats file concurrently - // (atlas bounds it to the CPU count). We then pack each column's per-dataset - // stats into DataFusion `ScalarValue` arrays — also one task per column, so - // the whole build is column-parallel end to end. - let keys: Vec = keyed.iter().map(|(_, k)| k.clone()).collect(); - let index = Arc::new(atlas.pruning_index(&keys).await.map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to build atlas pruning index: {e}" - )) - })?); - let rows = index.rows(); - - let mut handles = Vec::with_capacity(keyed.len()); - for (name, key) in keyed { - let index = index.clone(); - let schema = schema.clone(); - handles.push(tokio::task::spawn_blocking(move || { - pack_column(&index, &key, &name, &schema).map(|packed| (name, packed)) - })); - } - let mut columns: HashMap = HashMap::new(); - for handle in handles { - if let Ok(Some((name, packed))) = handle.await { - columns.insert(name, packed); - } - } - let stats = AtlasPruningStatistics { - num_containers: rows, - columns, + let built = tokio::task::spawn_blocking(move || { + let index = build_index(&atlas, &names, &wanted, &schema); + (names, index) + }) + .await; + + let Ok((names, index)) = built else { + return CandidateFilter::KeepAll; }; - let mask = pruning_predicate.prune(&stats).map_err(|e| { - datafusion::error::DataFusionError::Execution(format!("Failed to prune atlas datasets: {e}")) - })?; - - // Row ordinals the predicate couldn't rule out → their dataset names. - let mut candidates: HashSet = HashSet::new(); - for (row, keep) in mask.iter().enumerate() { - if *keep - && let Some(name) = index.dataset_name(row) - { - candidates.insert(name.to_string()); + if index.columns.is_empty() { + // Nothing the predicate names has statistics, so nothing can be ruled + // out. + return CandidateFilter::KeepAll; + } + + match pruning.prune(&index) { + Ok(kept) => CandidateFilter::Rows { + kept: BooleanArray::from(kept), + names, + }, + Err(e) => { + tracing::debug!("atlas pruning fell back to reading every dataset: {e}"); + CandidateFilter::KeepAll } } - Ok(Some(candidates)) } -/// Pack one column's per-dataset statistics into [`PackedStats`] (min/max as -/// `ScalarValue` arrays, plus count arrays). Pure CPU; run on a blocking task so -/// a wide, million-row pack doesn't stall the async runtime. `None` if the -/// column is absent or its datasets disagree on type. -/// -/// Each dataset stores its stats in its *own* dtype, so a column whose datasets -/// disagree (e.g. `Int16` in one, `Float32` in another) yields mixed `StatVal` -/// variants. Every value is cast to `target` — the column's type in the merged -/// (super-typed) table schema, which is also the type -/// [`PruningPredicate`](datafusion::physical_optimizer::pruning::PruningPredicate) -/// compares against — so the per-container min/max arrays are homogeneous and -/// the comparison is well-defined. -fn pack_column( - index: &PruningIndex, - key: &ColumnKey, - name: &str, +/// Pivot the footer into one [`StatColumn`] per column that has statistics. +fn build_index( + atlas: &Atlas, + names: &[String], + wanted: &[String], schema: &SchemaRef, -) -> Option { - let view = index.view(key)?; - let rows = index.rows(); - // `target` is the merged/super-typed table type for this column (the schema - // is derived from `Atlas::merged_schema`), so casting to it lifts every - // dataset's native-typed stat onto the common comparison type. - let target = schema - .field_with_name(name) - .map(|f| f.data_type().clone()) - .unwrap_or(DataType::Null); - let null_scalar = ScalarValue::try_from(&target).unwrap_or(ScalarValue::Null); - - let mut mins = Vec::with_capacity(rows); - let mut maxes = Vec::with_capacity(rows); - let mut null_counts: Vec> = Vec::with_capacity(rows); - let mut row_counts: Vec> = Vec::with_capacity(rows); - for row in 0..rows { - if view.is_present(row) { - mins.push(stat_to_scalar(view.min(row), &target, &null_scalar)); - maxes.push(stat_to_scalar(view.max(row), &target, &null_scalar)); - null_counts.push(Some(view.null_count(row))); - row_counts.push(Some(view.row_count(row))); +) -> PruningIndex { + // Where each dataset sits, so a footer pass in write order can be scattered + // into rows without a search. + let row_of: HashMap<&str, usize> = names + .iter() + .enumerate() + .map(|(row, name)| (name.as_str(), row)) + .collect(); + let arrays = atlas.list_arrays(); + + let mut columns = HashMap::new(); + for column in wanted { + let Ok(field) = schema.field_with_name(column) else { + continue; + }; + let target = field.data_type(); + + let packed = if arrays.iter().any(|array| array == column) { + Some(pack_array_column( + atlas, + names.len(), + &row_of, + column, + target, + )) } else { - // Column absent from this dataset: leave counts unknown so the - // predicate can't prune it (its rows read back null-filled, and the - // filter above the scan decides). Never claim "all null" — that - // could wrongly drop an `IS NULL` match. - mins.push(null_scalar.clone()); - maxes.push(null_scalar.clone()); - null_counts.push(None); - row_counts.push(None); + pack_attribute_column(atlas, names, &row_of, column, target) + }; + if let Some(packed) = packed { + columns.insert(column.clone(), packed); } } - let min = ScalarValue::iter_to_array(mins).ok()?; - let max = ScalarValue::iter_to_array(maxes).ok()?; - Some(PackedStats { + PruningIndex { + rows: names.len(), + columns, + } +} + +/// One array column, from one linear pass over the footer. +/// +/// A dataset with no entry for the array keeps a null bound and unknown counts. +/// That is what a dataset which does not declare the array looks like, and it +/// is also what one that declared it and never wrote it looks like — both must +/// stay in, and a null does exactly that. +fn pack_array_column( + atlas: &Atlas, + rows: usize, + row_of: &HashMap<&str, usize>, + column: &str, + target: &DataType, +) -> StatColumn { + let null = ScalarValue::try_from(target).unwrap_or(ScalarValue::Null); + let mut mins = vec![null.clone(); rows]; + let mut maxes = vec![null.clone(); rows]; + let mut null_counts: Vec> = vec![None; rows]; + let mut row_counts: Vec> = vec![None; rows]; + + for (dataset, stats) in atlas.array_stats_by_dataset(column) { + let Some(&row) = row_of.get(dataset.as_str()) else { + continue; + }; + mins[row] = stat_to_scalar(stats.min.as_ref(), target, &null); + maxes[row] = stat_to_scalar(stats.max.as_ref(), target, &null); + null_counts[row] = Some(stats.null_count); + row_counts[row] = Some(stats.row_count); + } + + StatColumn { + min: scalars_to_array(mins, rows, target), + max: scalars_to_array(maxes, rows, target), + null_count: Arc::new(UInt64Array::from(null_counts)), + row_count: Arc::new(UInt64Array::from(row_counts)), + } +} + +/// One attribute column, or `None` when it is not worth the pass. +/// +/// An attribute's value is exact, so it is both the minimum and the maximum of +/// its dataset. That prunes an equality on a dataset-level attribute — the +/// platform a file came from, say — out of the footer alone. +fn pack_attribute_column( + atlas: &Atlas, + names: &[String], + row_of: &HashMap<&str, usize>, + column: &str, + target: &DataType, +) -> Option { + if names.len() > ATTRIBUTE_INDEX_LIMIT { + tracing::debug!( + datasets = names.len(), + column, + "not indexing an attribute over a collection this large; see ATTRIBUTE_INDEX_LIMIT" + ); + return None; + } + + let rows = names.len(); + let null = ScalarValue::try_from(target).unwrap_or(ScalarValue::Null); + let mut values = vec![null.clone(); rows]; + let mut null_counts: Vec> = vec![None; rows]; + let mut seen = false; + + for name in names { + let Some(&row) = row_of.get(name.as_str()) else { + continue; + }; + let Ok(view) = atlas.dataset(name) else { + continue; + }; + let Some(value) = attribute_of(&view, column) else { + continue; + }; + let Some(scalar) = attr_to_scalar(&value) else { + continue; + }; + values[row] = scalar.cast_to(target).unwrap_or_else(|_| null.clone()); + // One value, and it is not the fill of anything. + null_counts[row] = Some(0); + seen = true; + } + + if !seen { + return None; + } + + let min = scalars_to_array(values.clone(), rows, target); + let max = scalars_to_array(values, rows, target); + Some(StatColumn { min, max, null_count: Arc::new(UInt64Array::from(null_counts)), - row_count: Arc::new(UInt64Array::from(row_counts)), + // An attribute is one value broadcast over whatever grid the dataset + // has, so its row count is not the dataset's. Unknown is honest. + row_count: Arc::new(UInt64Array::from(vec![None::; rows])), }) } -/// Resolve a table column name to the atlas [`ColumnKey`] whose statistics back -/// it, using the merged schema to disambiguate. `None` if it maps to nothing -/// prunable (an unknown name, or a `{array}.{attr}` whose parts don't resolve). -fn column_key(merged: &MergedSchema, name: &str) -> Option { - if merged.arrays.contains_key(name) { - return Some(ColumnKey::array(name)); - } - if merged.global_attributes.contains_key(name) { - return Some(ColumnKey::global_attr(name)); - } - // `{array}.{attr}` — try each `.` split, since both an array name and an - // attribute key may themselves contain dots. - for (i, c) in name.char_indices() { - if c == '.' { - let (array, rest) = name.split_at(i); - let attr = &rest[1..]; - if let Some(marr) = merged.arrays.get(array) - && marr.attributes.contains_key(attr) - { - return Some(ColumnKey::array_attr(array, attr)); +/// The attribute a column name refers to, dataset-level or per-array. +fn attribute_of(view: &atlas::DatasetView, column: &str) -> Option { + if let Some(key) = column.strip_prefix('.') { + return view.get_attribute(key); + } + // An array name and an attribute key may both hold dots, so every split is + // a candidate. + for (index, character) in column.char_indices() { + if character == '.' { + let (array, rest) = column.split_at(index); + if let Some(value) = view.get_array_attribute(array, &rest[1..]) { + return Some(value); } } } None } -/// Convert an atlas [`StatVal`] to a `ScalarValue` cast to `target`, or -/// `null_scalar` when absent or on any conversion failure. -fn stat_to_scalar( - value: Option<&StatVal>, - target: &DataType, - null_scalar: &ScalarValue, -) -> ScalarValue { +/// Pack scalars into one typed array, or a column of nulls when they will not. +fn scalars_to_array(values: Vec, rows: usize, target: &DataType) -> ArrayRef { + ScalarValue::iter_to_array(values) + .unwrap_or_else(|_| arrow::array::new_null_array(target, rows)) +} + +/// An atlas statistic as a scalar of the table's own type. +/// +/// A value that will not cast, and a `NaN` bound, both read as null. `NaN` +/// sorts last under `total_cmp`, so a `NaN` maximum says nothing about the +/// values below it, and claiming it as a bound would drop rows. +fn stat_to_scalar(value: Option<&StatValue>, target: &DataType, null: &ScalarValue) -> ScalarValue { let canonical = match value { - Some(StatVal::Int(x)) => ScalarValue::Int64(Some(*x)), - Some(StatVal::UInt(x)) => ScalarValue::UInt64(Some(*x)), - Some(StatVal::Float(x)) => ScalarValue::Float64(Some(*x)), - Some(StatVal::Bytes(b)) => match std::str::from_utf8(b) { - Ok(s) => ScalarValue::Utf8(Some(s.to_string())), - Err(_) => ScalarValue::Binary(Some(b.clone())), + Some(StatValue::Int(v)) => ScalarValue::Int64(Some(*v)), + Some(StatValue::UInt(v)) => ScalarValue::UInt64(Some(*v)), + Some(StatValue::Float(v)) if v.is_nan() => return null.clone(), + Some(StatValue::Float(v)) => ScalarValue::Float64(Some(*v)), + Some(StatValue::TimestampNs(v)) => ScalarValue::TimestampNanosecond(Some(*v), None), + Some(StatValue::Bytes(bytes)) => match std::str::from_utf8(bytes) { + Ok(text) => ScalarValue::Utf8(Some(text.to_string())), + Err(_) => ScalarValue::Binary(Some(bytes.clone())), }, - Some(StatVal::TimestampNs(x)) => ScalarValue::TimestampNanosecond(Some(*x), None), - None => return null_scalar.clone(), + None => return null.clone(), }; - canonical.cast_to(target).unwrap_or_else(|_| null_scalar.clone()) + canonical.cast_to(target).unwrap_or_else(|_| null.clone()) } -struct PackedStats { - min: ArrayRef, - max: ArrayRef, - null_count: ArrayRef, - row_count: ArrayRef, +/// An attribute value as a scalar, or `None` for a list, which bounds nothing. +fn attr_to_scalar(attr: &Attr) -> Option { + Some(match attr { + Attr::Bool(v) => ScalarValue::Boolean(Some(*v)), + Attr::Int8(v) => ScalarValue::Int8(Some(*v)), + Attr::Int16(v) => ScalarValue::Int16(Some(*v)), + Attr::Int32(v) => ScalarValue::Int32(Some(*v)), + Attr::Int64(v) => ScalarValue::Int64(Some(*v)), + Attr::UInt8(v) => ScalarValue::UInt8(Some(*v)), + Attr::UInt16(v) => ScalarValue::UInt16(Some(*v)), + Attr::UInt32(v) => ScalarValue::UInt32(Some(*v)), + Attr::UInt64(v) => ScalarValue::UInt64(Some(*v)), + Attr::Float32(v) if v.is_nan() => return None, + Attr::Float32(v) => ScalarValue::Float32(Some(*v)), + Attr::Float64(v) if v.is_nan() => return None, + Attr::Float64(v) => ScalarValue::Float64(Some(*v)), + Attr::String(v) => ScalarValue::Utf8(Some(v.clone())), + Attr::Binary(v) => ScalarValue::Binary(Some(v.clone())), + Attr::TimestampNanoseconds(v) => ScalarValue::TimestampNanosecond(Some(*v), None), + _ => return None, + }) } -/// [`PruningStatistics`] with one container per dataset row slot, backed by the -/// atlas pruning index. -struct AtlasPruningStatistics { - num_containers: usize, - columns: HashMap, -} +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Column as ColumnExpr, Literal}; -impl PruningStatistics for AtlasPruningStatistics { - fn min_values(&self, column: &Column) -> Option { - self.columns.get(column.name()).map(|c| c.min.clone()) + fn schema(name: &str, data_type: DataType) -> SchemaRef { + Arc::new(Schema::new(vec![Field::new(name, data_type, true)])) } - fn max_values(&self, column: &Column) -> Option { - self.columns.get(column.name()).map(|c| c.max.clone()) + fn binary(column: &str, op: Operator, value: ScalarValue) -> Arc { + Arc::new(BinaryExpr::new( + Arc::new(ColumnExpr::new(column, 0)), + op, + Arc::new(Literal::new(value)), + )) } - fn null_counts(&self, column: &Column) -> Option { - self.columns.get(column.name()).map(|c| c.null_count.clone()) + /// The datasets a predicate leaves in, in listing order. + async fn kept( + atlas: &Arc, + predicate: Arc, + schema: SchemaRef, + ) -> Vec { + let filter = candidate_filter(atlas, &predicate, &schema).await; + atlas + .list_datasets() + .into_iter() + .enumerate() + .filter(|(position, name)| filter.keeps(*position, name)) + .map(|(_, name)| name) + .collect() } - fn row_counts(&self, column: &Column) -> Option { - self.columns.get(column.name()).map(|c| c.row_count.clone()) - } + // ── the index over array statistics ───────────────────────────────── - fn num_containers(&self) -> usize { - self.num_containers - } + /// The ranged fixture gives dataset `d{i}` the values `[10i, 10i+3]`, so a + /// threshold has an answer that can be written down. + #[tokio::test] + async fn only_the_datasets_whose_range_reaches_the_threshold_survive() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 10).await; + let atlas = test_support::open(tmp.path()).await; + let schema = schema("temperature", DataType::Float32); - fn contained( - &self, - _column: &Column, - _values: &HashSet, - ) -> Option { - None + let survivors = kept( + &atlas, + binary( + "temperature", + Operator::Gt, + ScalarValue::Float32(Some(45.0)), + ), + schema, + ) + .await; + assert_eq!(survivors, vec!["d5", "d6", "d7", "d8", "d9"]); } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::datatypes::{Field, Schema}; - use atlas::Atlas; - use datafusion::logical_expr::Operator; - use datafusion::physical_expr::expressions::{binary, col, lit}; - use object_store::ObjectStore; - use object_store::local::LocalFileSystem; - use object_store::path::Path as OsPath; - async fn ranged_atlas(n: usize) -> (tempfile::TempDir, Arc) { + #[tokio::test] + async fn a_predicate_nothing_can_meet_prunes_everything() { let tmp = tempfile::tempdir().unwrap(); - crate::reader::test_support::build_ranged_store(tmp.path(), n).await; - let store: Arc = - Arc::new(LocalFileSystem::new_with_prefix(tmp.path()).unwrap()); - let atlas = Atlas::open(store, OsPath::from("")).await.unwrap(); - (tmp, Arc::new(atlas)) - } - - fn temperature_schema() -> SchemaRef { - Arc::new(Schema::new(vec![Field::new( - "temperature", - DataType::Float32, - true, - )])) - } - - fn gt(schema: &SchemaRef, threshold: f32) -> Arc { - binary( - col("temperature", schema).unwrap(), - Operator::Gt, - lit(ScalarValue::Float32(Some(threshold))), - schema, + test_support::ranged(tmp.path(), 6).await; + let atlas = test_support::open(tmp.path()).await; + + let survivors = kept( + &atlas, + binary( + "temperature", + Operator::Gt, + ScalarValue::Float32(Some(10_000.0)), + ), + schema("temperature", DataType::Float32), ) - .unwrap() + .await; + assert!(survivors.is_empty(), "{survivors:?}"); } #[tokio::test] - async fn retains_only_datasets_whose_range_can_match() { - // 10 datasets: d{i}.temperature ∈ [10i, 10i+3]. - let (_tmp, atlas) = ranged_atlas(10).await; - let schema = temperature_schema(); - let names = atlas.list_datasets(); + async fn a_predicate_everything_meets_prunes_nothing() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 6).await; + let atlas = test_support::open(tmp.path()).await; - // `> 45` matches d5..d9 (ranges 50-53 … 90-93); d0..d4 (max 43) are pruned. - let kept = retain_candidates(&atlas, names.clone(), >(&schema, 45.0), &schema).await; - assert_eq!(kept, vec!["d5", "d6", "d7", "d8", "d9"]); + let survivors = kept( + &atlas, + binary( + "temperature", + Operator::GtEq, + ScalarValue::Float32(Some(0.0)), + ), + schema("temperature", DataType::Float32), + ) + .await; + assert_eq!(survivors, atlas.list_datasets()); } + /// The index holds one row per live dataset, in listing order, and the + /// filter reads it by position. #[tokio::test] - async fn impossible_predicate_prunes_everything() { - let (_tmp, atlas) = ranged_atlas(6).await; - let schema = temperature_schema(); - let kept = - retain_candidates(&atlas, atlas.list_datasets(), >(&schema, 1_000.0), &schema).await; - assert!(kept.is_empty(), "no dataset reaches 1000: {kept:?}"); + async fn the_index_holds_one_row_per_live_dataset() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 10).await; + let atlas = test_support::open(tmp.path()).await; + + let filter = candidate_filter( + &atlas, + &binary( + "temperature", + Operator::Gt, + ScalarValue::Float32(Some(45.0)), + ), + &schema("temperature", DataType::Float32), + ) + .await; + assert_eq!(filter.rows(), 10); + assert_eq!(filter.pruned(), 5); } + /// A deleted dataset has no row at all, and the rows after it shift up. + /// That is why the filter checks the name it was given against the row it + /// indexes. #[tokio::test] - async fn permissive_predicate_keeps_everything() { - let (_tmp, atlas) = ranged_atlas(6).await; - let schema = temperature_schema(); - let all = atlas.list_datasets(); - let kept = retain_candidates(&atlas, all.clone(), >(&schema, -1.0), &schema).await; - assert_eq!(kept, all, "every dataset can exceed -1"); + async fn a_delete_shifts_the_rows_and_the_name_check_catches_it() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 6).await; + let atlas = test_support::open(tmp.path()).await; + atlas.delete_dataset("d0").await.unwrap(); + + let filter = candidate_filter( + &atlas, + &binary( + "temperature", + Operator::Gt, + ScalarValue::Float32(Some(45.0)), + ), + &schema("temperature", DataType::Float32), + ) + .await; + assert_eq!(filter.rows(), 5, "the deleted dataset has no row"); + + // Row 0 is now d1. A plan made before the delete would ask about d0 + // there, and that must not read as d1's answer. + assert!( + filter.keeps(0, "d0"), + "a name that does not match its row is kept, not mis-indexed" + ); } + // ── mixed and awkward types ───────────────────────────────────────── + + /// Two datasets that type one array differently still prune: every bound is + /// cast to the column's table type before it is compared. #[tokio::test] - async fn mixed_dtype_column_casts_to_merged_type_then_prunes() { - // The widening fixture stores `value` as Int16 in `a` ([1,2]) and - // Float32 in `b` ([3.5,4.5]); the merged column type is Float32. Pruning - // must cast `a`'s Int16 stats up to Float32 before comparing, or the two - // datasets' min/max buffers wouldn't even be a single Arrow array. + async fn a_mixed_dtype_column_is_cast_before_it_is_compared() { let tmp = tempfile::tempdir().unwrap(); - crate::reader::test_support::build_widening_store(tmp.path()).await; - let store: Arc = - Arc::new(LocalFileSystem::new_with_prefix(tmp.path()).unwrap()); - let atlas = Arc::new(Atlas::open(store, OsPath::from("")).await.unwrap()); - - let schema = Arc::new(Schema::new(vec![Field::new( - "value", - DataType::Float32, - true, - )])); - let value_gt = |t: f32| { - binary( - col("value", &schema).unwrap(), - Operator::Gt, - lit(ScalarValue::Float32(Some(t))), - &schema, - ) - .unwrap() - }; - let names = atlas.list_datasets(); // [a, b] + test_support::widening(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + // Int16 and Float32 merge to Float64 under the session rule. + let schema = schema("value", DataType::Float64); - // `> 3`: a (max 2, cast from Int16) is pruned; b (max 4.5) survives. + // a holds [1, 2] and b holds [3.5, 4.5]. assert_eq!( - retain_candidates(&atlas, names.clone(), &value_gt(3.0), &schema).await, + kept( + &atlas, + binary("value", Operator::Gt, ScalarValue::Float64(Some(3.0))), + Arc::clone(&schema) + ) + .await, vec!["b"] ); - // `> 10`: neither can match. + assert_eq!( + kept( + &atlas, + binary("value", Operator::Lt, ScalarValue::Float64(Some(3.0))), + Arc::clone(&schema) + ) + .await, + vec!["a"] + ); assert!( - retain_candidates(&atlas, names.clone(), &value_gt(10.0), &schema) - .await - .is_empty() + kept( + &atlas, + binary("value", Operator::Gt, ScalarValue::Float64(Some(100.0))), + schema + ) + .await + .is_empty() + ); + } + + /// A dataset-level attribute is exact, so an equality on it prunes from the + /// footer alone. + #[tokio::test] + async fn an_attribute_predicate_prunes() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 6).await; + let atlas = test_support::open(tmp.path()).await; + + let survivors = kept( + &atlas, + binary( + ".platform", + Operator::Eq, + ScalarValue::Utf8(Some("p3".to_string())), + ), + schema(".platform", DataType::Utf8), + ) + .await; + assert_eq!(survivors, vec!["d3"]); + } + + // ── failing open ──────────────────────────────────────────────────── + + #[tokio::test] + async fn a_column_with_no_statistics_prunes_nothing() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 4).await; + let atlas = test_support::open(tmp.path()).await; + + let survivors = kept( + &atlas, + binary("ghost", Operator::Gt, ScalarValue::Float32(Some(0.0))), + schema("ghost", DataType::Float32), + ) + .await; + assert_eq!(survivors, atlas.list_datasets()); + } + + /// A column one dataset declares and another does not: the one without it + /// has no bound, so it stays in and its rows are decided above the scan. + #[tokio::test] + async fn a_dataset_that_lacks_the_column_is_never_pruned_on_it() { + let tmp = tempfile::tempdir().unwrap(); + test_support::widening(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + // Only `a` declares `flag`, and it holds [7, 8]. + let survivors = kept( + &atlas, + binary("flag", Operator::Gt, ScalarValue::Int32(Some(100))), + schema("flag", DataType::Int32), + ) + .await; + assert_eq!(survivors, vec!["b"], "a is ruled out, b cannot be"); + } + + #[tokio::test] + async fn a_collection_with_no_datasets_prunes_nothing() { + let tmp = tempfile::tempdir().unwrap(); + test_support::empty(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let filter = candidate_filter( + &atlas, + &binary("temperature", Operator::Gt, ScalarValue::Float32(Some(0.0))), + &schema("temperature", DataType::Float32), + ) + .await; + assert!(matches!(filter, CandidateFilter::KeepAll)); + } + + // ── the pieces ────────────────────────────────────────────────────── + + #[test] + fn a_nan_bound_is_no_bound() { + let null = ScalarValue::Float64(None); + let nan = stat_to_scalar(Some(&StatValue::Float(f64::NAN)), &DataType::Float64, &null); + assert!(nan.is_null(), "NaN sorts last, so it bounds nothing"); + } + + #[test] + fn a_bound_that_will_not_cast_is_no_bound() { + let null = ScalarValue::Int32(None); + let text = stat_to_scalar( + Some(&StatValue::Bytes(b"not a number".to_vec())), + &DataType::Int32, + &null, + ); + assert!(text.is_null()); + } + + #[test] + fn a_text_bound_survives_as_text() { + let null = ScalarValue::Utf8(None); + let text = stat_to_scalar( + Some(&StatValue::Bytes(b"argo".to_vec())), + &DataType::Utf8, + &null, ); - // `> 1.5`: a's cast max (2.0) still qualifies, so both survive. + assert_eq!(text, ScalarValue::Utf8(Some("argo".to_string()))); + } + + #[test] + fn a_list_attribute_bounds_nothing() { + assert!(attr_to_scalar(&Attr::Int32List(vec![1, 2])).is_none()); + assert!(attr_to_scalar(&Attr::Float64(f64::NAN)).is_none()); assert_eq!( - retain_candidates(&atlas, names.clone(), &value_gt(1.5), &schema).await, - names + attr_to_scalar(&Attr::Int64(7)), + Some(ScalarValue::Int64(Some(7))) ); } - #[tokio::test] - async fn unmappable_column_fails_open() { - // A predicate on a column with no atlas statistics must keep every input. - let (_tmp, atlas) = ranged_atlas(4).await; - let schema = Arc::new(Schema::new(vec![Field::new( - "ghost", - DataType::Float32, - true, - )])); - let pred = binary( - col("ghost", &schema).unwrap(), - Operator::Gt, - lit(ScalarValue::Float32(Some(0.0))), - &schema, + /// The whole point of an index: a collection of hundreds of thousands of + /// datasets is judged in one vectorised pass. + /// + /// The index is built by hand here. Writing that many real datasets would + /// take minutes and prove nothing extra — what this pins is that the + /// evaluation is one pass over Arrow arrays rather than a decision per + /// dataset. + #[test] + fn a_large_index_is_judged_in_one_pass() { + use arrow::array::Float64Array; + + const ROWS: usize = 200_000; + const THRESHOLD: f64 = 199_000.0; + + // Row i covers [i, i + 1], so exactly the rows above the threshold + // survive. + let mins: Float64Array = (0..ROWS).map(|row| Some(row as f64)).collect(); + let maxes: Float64Array = (0..ROWS).map(|row| Some(row as f64 + 1.0)).collect(); + let counts: UInt64Array = (0..ROWS).map(|_| Some(0u64)).collect(); + let rows: UInt64Array = (0..ROWS).map(|_| Some(1u64)).collect(); + + let index = PruningIndex { + rows: ROWS, + columns: HashMap::from([( + "temperature".to_string(), + StatColumn { + min: Arc::new(mins), + max: Arc::new(maxes), + null_count: Arc::new(counts), + row_count: Arc::new(rows), + }, + )]), + }; + + let pruning = PruningPredicate::try_new( + binary( + "temperature", + Operator::Gt, + ScalarValue::Float64(Some(THRESHOLD)), + ), + schema("temperature", DataType::Float64), ) - .unwrap(); - let all = atlas.list_datasets(); - let kept = retain_candidates(&atlas, all.clone(), &pred, &schema).await; - assert_eq!(kept, all, "unknown column must not prune anything"); + .expect("the predicate is prunable"); + + let kept = pruning.prune(&index).expect("one pass over the index"); + assert_eq!(kept.len(), ROWS); + // Row i survives when its maximum, i + 1, exceeds the threshold, so the + // survivors are the rows from the threshold onward. + let expected = ROWS - THRESHOLD as usize; + assert_eq!(kept.iter().filter(|keep| **keep).count(), expected); + } + + /// The scan's schema is nd-encoded; a predicate is written against the + /// values inside it. + #[test] + fn the_logical_schema_unwraps_the_encoding() { + let logical = Schema::new(vec![Field::new("temperature", DataType::Float32, true)]); + let encoded = beacon_datafusion_ext::nd::encoded_schema(&logical); + assert_ne!( + encoded.field(0).data_type(), + &DataType::Float32, + "the encoded form is a struct" + ); + assert_eq!( + logical_schema(&encoded).field(0).data_type(), + &DataType::Float32 + ); } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs index 23153697..4063e181 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs @@ -1,123 +1,135 @@ -//! DataFusion [`FileSource`]/[`FileOpener`] for atlas stores. +//! The DataFusion [`FileSource`] and [`FileOpener`] for Atlas collections. //! -//! Each `PartitionedFile` is one atlas store's metadata marker plus a slice of -//! that store's dataset names (attached as [`AtlasDatasetSlice`] in the file's -//! `extensions` by [`AtlasFormat::create_physical_plan`](super::AtlasFormat)). -//! The opener opens the store once over the query's object store and reads only -//! its assigned datasets — so a store's datasets are spread across DataFusion -//! partitions and scanned on every core in parallel. Each dataset is built with -//! just the projected columns and streamed through the shared `beacon-nd-array` -//! engine (predicate row-masking via [`PushdownFilter`]). +//! # One dataset is one unit of work +//! +//! A plan entry is a *dataset*, not a collection: [`AtlasFormat`] lists each +//! collection at plan time and emits one [`PartitionedFile`] per dataset, with +//! the dataset's name in [`PartitionedFile::extensions`]. Every entry carries +//! the collection's own marker, so the opener knows which container to open and +//! the reader cache keys on the same object the plan did. +//! +//! Those entries go into one [`MorselSource`], and each partition holds a +//! standing entry pointing at it. A partition takes the next dataset when it is +//! free, and helps drain an open one when none is left. Balance follows +//! completion, so a collection of a million small datasets and a collection of +//! four large ones both divide over every core. +//! +//! [`AtlasFormat`]: super::AtlasFormat use std::any::Any; use std::sync::Arc; use std::time::Instant; use arrow::datatypes::SchemaRef; -use arrow::record_batch::{RecordBatch, RecordBatchOptions}; -use beacon_nd_array::{ - arrow::{ - batch::any_dataset_as_record_batch_stream, metrics::DatasetReadMetrics, - pushdown_filter::PushdownFilter, schema::any_dataset_to_arrow_schema, - }, - dataset::resolve_read_dimensions, - projection::DatasetProjection, +use atlas::{Atlas, DatasetView}; +use beacon_nd_array::arrow::{ + file_read::FileRead, + metrics::ReadMetrics, + morsel::{MorselSource, OpenFile, morsel_scan}, + partition::FilePartitions, }; -use beacon_datafusion_ext::scan_adapt::batch_adapter_factory; use datafusion::{ config::ConfigOptions, datasource::{ listing::PartitionedFile, physical_plan::{FileOpenFuture, FileOpener, FileScanConfig, FileSource}, - schema_adapter::SchemaAdapterFactory, table_schema::TableSchema, }, - physical_expr::{PhysicalExpr, conjunction, projection::ProjectionExprs}, + error::{DataFusionError, Result}, + physical_expr::{ + PhysicalExpr, conjunction, projection::ProjectionExprs, utils::collect_columns, + }, physical_plan::{ filter_pushdown::{FilterPushdownPropagation, PushedDown}, metrics::ExecutionPlanMetricsSet, }, }; -use futures::future; -use futures::{StreamExt, TryStreamExt, stream::BoxStream}; -use object_store::{ObjectMeta, ObjectStore}; +use futures::FutureExt; +use object_store::ObjectStore; -use crate::datafusion::cache::AtlasReaderCache; +use crate::compat; use crate::datafusion::metrics::AtlasScanMetrics; -use crate::datafusion::pruning::PruneCache; - -/// How many of a partition's datasets the opener reads concurrently. Overlaps -/// per-dataset I/O and decompression within a partition; cross-partition -/// parallelism comes from DataFusion running the partitions on separate cores. -const ATLAS_DATASET_CONCURRENCY: usize = 8; +use crate::datafusion::pruning::{CandidateFilter, PruneCache, candidate_filter, logical_schema}; +use crate::reader::{dataset_from_view, project_read_dimensions}; +use crate::store::{AtlasReaderCache, get_or_open_atlas}; -/// The slice of a store's dataset names assigned to one scan partition. +/// Which dataset of a collection one plan entry stands for. /// -/// Attached to each [`PartitionedFile::extensions`] by -/// [`AtlasFormat::create_physical_plan`](super::AtlasFormat). When absent (e.g. -/// a source built outside the physical plan), the opener falls back to reading -/// every dataset in the store. +/// Attached to [`PartitionedFile::extensions`] by +/// [`AtlasFormat::create_physical_plan`](super::AtlasFormat). The `position` is +/// the dataset's index in the `list_datasets()` call the plan made, which is +/// the row a collection-wide pruning index keys on. #[derive(Debug, Clone)] -pub struct AtlasDatasetSlice { - pub names: Vec, +pub struct AtlasEntry { + /// The dataset's name, as the collection footer states it. + pub dataset: String, + /// Its row in the plan-time listing. + pub position: usize, } -/// DataFusion [`FileSource`] for atlas stores. +/// DataFusion [`FileSource`] for Atlas collections. #[derive(Debug, Clone)] pub struct AtlasSource { - schema_adapter_factory: Option>, table_schema: TableSchema, execution_plan_metrics: ExecutionPlanMetricsSet, batch_size: usize, predicate: Option>, read_dimensions: Option>, - /// Reader cache to consult for this scan. `None` disables caching. + projection: Option, + /// The reader cache to consult, or `None` to open every collection afresh. cache: Option, - /// Whether to prune non-matching datasets before reading them. + /// Whether a predicate scan drops the datasets it can rule out. use_pruning: bool, - /// Per-query memo of each store's pruning result, shared across the scan's - /// partition openers so a store is pruned once, not once per partition. + /// Each collection's pruning result, computed once for this scan and shared + /// by every partition's opener. prune_cache: PruneCache, - /// Projection pushed down by the scan, applied on top of the table schema. - projection: Option, + /// The scan's dataset queue, when it is planned morsel-driven. See + /// [`morsel_scan`]. + morsel: Option>, } impl AtlasSource { pub fn new(read_dimensions: Option>, table_schema: TableSchema) -> Self { Self { - schema_adapter_factory: None, table_schema, execution_plan_metrics: ExecutionPlanMetricsSet::new(), batch_size: usize::MAX, predicate: None, read_dimensions, + projection: None, cache: None, use_pruning: false, prune_cache: PruneCache::new(), - projection: None, + morsel: None, } } - /// Returns a copy of this source that consults `cache` (when `Some`) for - /// opened atlas stores. The format wires in the runtime's shared cache here. + /// Consult `cache` for opened collections, or open them afresh with `None`. pub fn with_cache(mut self, cache: Option) -> Self { self.cache = cache; self } - /// Enable or disable dataset pruning for this scan. + /// Drop the datasets a predicate rules out, or read them all. pub fn with_pruning(mut self, use_pruning: bool) -> Self { self.use_pruning = use_pruning; self } - /// Returns a copy of this source carrying the given projection. Used to - /// preserve a pushed-down projection when the format rebuilds the source in - /// `create_physical_plan`. + /// Carry a projection the scan pushed down. + /// + /// The format rebuilds the source in `create_physical_plan`, and without + /// this the projection pushed into the old one would be lost. pub fn with_projection(mut self, projection: Option) -> Self { self.projection = projection; self } + + /// The datasets this scan's queue holds, when it is planned morsel-driven. + #[cfg(test)] + pub(crate) fn morsel_datasets(&self) -> Option { + self.morsel.as_ref().map(|source| source.files()) + } } impl FileSource for AtlasSource { @@ -126,24 +138,36 @@ impl FileSource for AtlasSource { object_store: Arc, base_config: &FileScanConfig, partition: usize, - ) -> datafusion::error::Result> { + ) -> Result> { let projected_schema = base_config.projected_schema()?; + let read_metrics = ReadMetrics::new(&self.execution_plan_metrics, partition); + let scan_metrics = AtlasScanMetrics::new(&self.execution_plan_metrics, partition); - Ok(Arc::new(AtlasOpener { + let datasets = Arc::new(AtlasDatasets { object_store, - projected_schema, - batch_size: self.batch_size, - metrics: self.execution_plan_metrics.clone(), - partition, - read_dimensions: self.read_dimensions.clone(), - predicate: self.predicate.clone(), cache: self.cache.clone(), + // A predicate is written against the values, not the encoding the + // scan carries them in. + logical_schema: logical_schema(&projected_schema), + projected_schema, use_pruning: self.use_pruning, prune_cache: self.prune_cache.clone(), + read_dimensions: self.read_dimensions.clone(), + batch_size: self.batch_size, + predicate: self.predicate.clone(), + read_metrics: read_metrics.clone(), + scan_metrics, + }); + + Ok(Arc::new(AtlasOpener { + datasets, + morsel: self.morsel.clone(), + partition, + read_metrics, })) } - fn as_any(&self) -> &dyn std::any::Any { + fn as_any(&self) -> &dyn Any { self } @@ -158,25 +182,50 @@ impl FileSource for AtlasSource { }) } - /// Whether a scan may split one file across partitions. It may not. + /// Put every dataset of the scan in one queue, and point each partition at + /// it. /// - /// Atlas divides its own work, and it divides it by dataset name. - /// `create_physical_plan` opens each store, lists its datasets, and gives - /// every file group one slice of those names in - /// [`PartitionedFile::extensions`]. Every slice carries the *same* - /// `object_meta`, because they all name the same store. + /// Nothing is assigned here. A dataset's cost is the cells the query keeps, + /// which no plan-time number states: two datasets of one collection differ + /// by orders of magnitude, and a predicate prunes them unevenly. So the + /// partitions divide the queue as they drain it. /// - /// DataFusion's partitioner splits by byte range and copies the extensions - /// into each share. Two shares of one slice would therefore carry the same - /// [`AtlasDatasetSlice`], and the opener reads by name and never looks at a - /// byte range, so each share would return the same datasets over again. - /// - /// Declining the split does not cost parallelism. The name slices already - /// give one partition per share of a store's datasets. - /// - /// [`PartitionedFile::extensions`]: datafusion::datasource::listing::PartitionedFile::extensions - fn supports_repartitioning(&self) -> bool { - false + /// `repartition_file_min_size` is ignored, as it is for Zarr. It was the + /// size a file had to reach before sharing it was worth the seek, and a + /// queue makes no such bet. An atlas entry has no size of its own anyway: + /// every dataset of a collection reports the container's. + fn repartitioned( + &self, + target_partitions: usize, + _repartition_file_min_size: usize, + output_ordering: Option, + config: &FileScanConfig, + ) -> Result> { + if output_ordering.is_some() || target_partitions <= 1 { + // A partition holding an arbitrary share of the datasets cannot + // emit its rows in collection order. + return Ok(None); + } + + if let Some((morsel, file_groups)) = morsel_scan(&config.file_groups, target_partitions) { + tracing::debug!( + "AtlasSource morsel scan: {} datasets over {target_partitions} partitions", + morsel.files() + ); + let mut config = config.clone(); + config.file_groups = file_groups; + // The openers are built from the config's source, so the queue has + // to travel with it. + config.file_source = Arc::new(Self { + morsel: Some(morsel), + ..self.clone() + }); + return Ok(Some(config)); + } + + // The queue declined: one partition, or no datasets. Keeping the scan + // as planned is the answer to both. + Ok(None) } fn metrics(&self) -> &ExecutionPlanMetricsSet { @@ -187,20 +236,6 @@ impl FileSource for AtlasSource { "atlas" } - fn with_schema_adapter_factory( - &self, - factory: Arc, - ) -> datafusion::error::Result> { - Ok(Arc::new(Self { - schema_adapter_factory: Some(factory), - ..self.clone() - })) - } - - fn schema_adapter_factory(&self) -> Option> { - self.schema_adapter_factory.clone() - } - fn projection(&self) -> Option<&ProjectionExprs> { self.projection.as_ref() } @@ -208,27 +243,29 @@ impl FileSource for AtlasSource { fn try_pushdown_projection( &self, projection: &ProjectionExprs, - ) -> datafusion::error::Result>> { - // Merge with any projection already pushed down, then record it on a new - // source. `FileScanConfig::projected_schema` reads this back via - // `projection()`, and the opener's schema adapter applies it per dataset. + ) -> Result>> { let merged = match &self.projection { Some(existing) => existing.try_merge(projection)?, None => projection.clone(), }; - - let source = Self { + Ok(Some(Arc::new(Self { projection: Some(merged), ..self.clone() - }; - Ok(Some(Arc::new(source))) + }))) } + /// Take the filters as a hint, and leave them above the scan. + /// + /// The scan uses a predicate twice: to skip a chunk whose coordinates + /// cannot hold a matching row, and (once pruning lands) to skip a whole + /// dataset whose footer statistics cannot. Neither is exact — both work in + /// whole chunks and whole datasets — so the filter above the scan still + /// decides each row, and `PushedDown::No` is what says so. fn try_pushdown_filters( &self, filters: Vec>, _config: &ConfigOptions, - ) -> datafusion::error::Result>> { + ) -> Result>> { let predicate = match self.predicate.clone() { Some(existing) => conjunction(std::iter::once(existing).chain(filters.clone())), None => conjunction(filters.clone()), @@ -247,452 +284,450 @@ impl FileSource for AtlasSource { } } -// ─── FileOpener ──────────────────────────────────────────────────────────── +// ─── The opener ────────────────────────────────────────────────────────────── +/// One partition's opener. struct AtlasOpener { - object_store: Arc, - projected_schema: SchemaRef, - batch_size: usize, - metrics: ExecutionPlanMetricsSet, + /// How one dataset is opened, for the queue to call. + datasets: Arc, + /// The scan's queue, when it is planned morsel-driven. `Some` means the + /// entry `FileStream` hands this opener stands for the whole scan. + morsel: Option>, partition: usize, - read_dimensions: Option>, - predicate: Option>, - cache: Option, - use_pruning: bool, - prune_cache: PruneCache, + read_metrics: ReadMetrics, } -impl AtlasOpener { - #[allow(clippy::too_many_arguments)] - async fn read_task( - object_store: Arc, - object_meta: ObjectMeta, - assigned_names: Option>, - projected_schema: SchemaRef, - batch_size: usize, - metrics: ExecutionPlanMetricsSet, - partition: usize, - read_dimensions: Option>, - predicate: Option>, - cache: Option, - use_pruning: bool, - prune_cache: PruneCache, - ) -> datafusion::error::Result>> { - let scan_metrics = AtlasScanMetrics::new(&metrics, partition); - let read_metrics = DatasetReadMetrics::new(&metrics, partition); - - let open_start = Instant::now(); - let atlas = - crate::datafusion::cache::get_or_open_atlas(cache.as_ref(), object_store, &object_meta) - .await?; - scan_metrics.open_time.add_elapsed(open_start); - let object_path = object_meta.location.clone(); - - // The physical plan assigns each partition its slice of dataset names; - // fall back to the whole store if a source was built without one. - let names = assigned_names.unwrap_or_else(|| atlas.list_datasets()); - let assigned = names.len(); - - // Skip datasets the predicate proves can't match, before opening them. - // The store is pruned once per query (memoized by marker path) and every - // partition's opener reuses the result; only the projected schema is - // needed — it carries the predicate columns (a filter that stays above - // the scan forces them in). Fails open. - let names = match (use_pruning, &predicate) { - (true, Some(pred)) => { - let prune_start = Instant::now(); - let atlas = atlas.clone(); - let pred = pred.clone(); - let schema = projected_schema.clone(); - let filter = prune_cache - .get_or_compute(object_path.to_string(), async move { - Arc::new( - crate::datafusion::pruning::candidate_set(&atlas, &pred, &schema).await, - ) - }) - .await; - let kept = filter.retain(names); - scan_metrics.prune_time.add_elapsed(prune_start); - kept - } - _ => names, - }; - scan_metrics.datasets_pruned.add(assigned - names.len()); - scan_metrics.datasets_scanned.add(names.len()); - - let stream = futures::stream::iter(names) - .map(move |dataset_name| { - let atlas = atlas.clone(); - let projected_schema = projected_schema.clone(); - let read_dimensions = read_dimensions.clone(); - let read_metrics = read_metrics.clone(); - let scan_metrics = scan_metrics.clone(); - let object_path = object_path.clone(); - let predicate = predicate.clone(); - async move { - read_one_dataset( - atlas, - dataset_name, - object_path, - projected_schema, - read_dimensions, - predicate, - batch_size, - Some(read_metrics), - scan_metrics, - ) - .await - } - }) - .buffer_unordered(ATLAS_DATASET_CONCURRENCY) - .try_flatten() - .boxed(); +impl FileOpener for AtlasOpener { + fn open(&self, file: PartitionedFile) -> Result { + // A morsel-driven scan hands every partition the same standing entry. + // It is not a dataset: the datasets are in the queue, and this + // partition reads whatever it hands out until the scan is done. + if let Some(morsel) = &self.morsel { + let stream = morsel.stream( + self.partition, + Arc::clone(&self.datasets), + Some(self.read_metrics.clone()), + ); + return Ok(futures::future::ready(Ok(stream)).boxed()); + } - Ok(stream) + // One partition, or no datasets: `FileStream` walks the real entries + // and this opener reads each one whole. + let datasets = Arc::clone(&self.datasets); + let metrics = self.read_metrics.clone(); + Ok(async move { + let read = datasets.open(&file).await?; + Ok(read.stream(Some(metrics))) + } + .boxed()) } } -/// Build and stream one atlas dataset's projected columns, adapting each batch -/// onto the query's projected output schema. -#[allow(clippy::too_many_arguments)] -async fn read_one_dataset( - atlas: Arc, - dataset_name: String, - object_path: object_store::path::Path, +/// How one Atlas dataset becomes a planned [`FileRead`]. +/// +/// This is everything a [`MorselSource`] needs of the format. The queue holds +/// the datasets; this says what opening one means. +struct AtlasDatasets { + object_store: Arc, + cache: Option, + /// The scan's output schema, nd-encoded. Its field *names* are the columns + /// to keep, and the encoding leaves names alone. projected_schema: SchemaRef, + /// The same schema with the encoding unwrapped, which is what a predicate + /// and the pruning engine are written against. + logical_schema: SchemaRef, + use_pruning: bool, + prune_cache: PruneCache, read_dimensions: Option>, - predicate: Option>, batch_size: usize, - metrics: Option, + predicate: Option>, + read_metrics: ReadMetrics, scan_metrics: AtlasScanMetrics, -) -> datafusion::error::Result>> { - // Time everything up to handing back the (still-lazy) stream: opening the - // view, reading metadata + projected attributes, wiring backends, and - // building the schema adapter. Array *data* is read later, as the returned - // stream is polled, and is counted by `DatasetReadMetrics` (rows/batches). - let build_start = Instant::now(); - - // Push the query's projection straight into the dataset build: only the - // requested columns get backends, and only their attribute values are read - // from disk. `projected_schema` is the scan's output schema, so its field - // names are the exact column set to keep (it also carries any predicate - // columns, since a not-fully-consumed filter stays above the scan). A - // dataset simply omits any name it doesn't declare. An empty projection is - // `COUNT(*)` — no column names, so nothing is built here. - let projected_names: Vec = projected_schema - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(); +} - let projected = crate::reader::dataset_from_atlas(atlas.clone(), &dataset_name, Some(&projected_names)) - .await - .map_err(|e| { - tracing::warn!(dataset = %dataset_name, path = %object_path, error = %e, "failed to read atlas dataset"); - datafusion::error::DataFusionError::Execution(format!( - "Failed to read atlas dataset '{dataset_name}' at {object_path}: {e}" - )) - })?; - - // Apply explicit dimensions, or narrow to a broadcast-compatible default so - // a mix of incompatible dimension sets can't fail the scan. No log label: - // this runs per dataset (logging happens in schema inference). - let projected = match resolve_read_dimensions(&projected, read_dimensions, None) { - Some(dims) => projected - .project(&DatasetProjection::new_with_dimension_projection(dims)) - .map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to project atlas dataset '{dataset_name}' with dimensions: {e}" - )) - })?, - None => projected, - }; - - let stream = if !projected.dataset().arrays.is_empty() { - // Adapt each batch onto the scan's output schema. The dataset's arrays - // carry their own (per-dataset) dtypes, which the merged table schema - // may have widened — so the source schema reflects the native types and - // the adapter casts each column up to the super-type and null-fills any - // projected column this dataset doesn't declare, keyed by name. - stream_adapted(projected, projected_schema, predicate, batch_size, metrics)? - } else { - // This dataset declares *none* of the projected columns. Its rows must - // still appear, null-filled to the table schema (union semantics): - // selecting a column the dataset lacks yields its rows as nulls, and - // `COUNT(*)` (an empty projection) counts them. Establish the row count - // from the dataset's largest readable array, then let the adapter - // null-fill every projected column. The predicate is dropped here — - // those columns are all null, and the not-fully-consumed filter kept - // above the scan re-checks them, so correctness holds. - match driver_dataset(&atlas, &dataset_name, &object_path).await? { - Some(driver) => stream_adapted(driver, projected_schema, None, batch_size, metrics)?, - None => { - // No readable array either (attribute-only / bool / list - // dataset): contribute a single broadcast row. - let batch = null_row_batch(&projected_schema)?; - futures::stream::once(async move { Ok(batch) }).boxed() - } +impl AtlasDatasets { + /// Which datasets of one collection this scan's predicate can still match. + /// + /// Built once per collection per scan. Every partition's opener shares one + /// memo, so the first to arrive builds the index while the rest await it; + /// each then reads its own dataset's bit out of the result. + /// + /// Without a predicate, or with pruning off, nothing is ruled out and no + /// index is built. + async fn candidates(&self, atlas: &Arc, marker: &str) -> Arc { + let Some(predicate) = self.predicate.clone().filter(|_| self.use_pruning) else { + return Arc::new(CandidateFilter::KeepAll); + }; + // DataFusion offers a scan its filters even when it has none, and an + // empty conjunction is the literal `true`. Such a predicate names no + // column, so it can rule nothing out and is not worth a pass. + if collect_columns(&predicate).is_empty() { + return Arc::new(CandidateFilter::KeepAll); } - }; - scan_metrics.dataset_build_time.add_elapsed(build_start); - Ok(stream) -} + let started = Instant::now(); + let atlas = Arc::clone(atlas); + let schema = Arc::clone(&self.logical_schema); + let metrics = self.scan_metrics.clone(); + let filter = self + .prune_cache + .get_or_compute(marker.to_string(), async move { + let filter = Arc::new(candidate_filter(&atlas, &predicate, &schema).await); + // Only an index that exists is a build. Pruning that did not + // apply read nothing and judged nothing. + if filter.is_index() { + metrics.index_builds.add(1); + metrics.index_rows.add(filter.rows()); + } + filter + }) + .await; + self.scan_metrics.prune_time.add_elapsed(started); + filter + } -/// Adapt a dataset's record-batch stream onto `projected_schema`: cast each -/// column to the table's (possibly widened) type and null-fill any projected -/// column the dataset lacks, matching by name. -fn stream_adapted( - dataset: beacon_nd_array::dataset::AnyDataset, - projected_schema: SchemaRef, - predicate: Option>, - batch_size: usize, - metrics: Option, -) -> datafusion::error::Result>> { - let source_schema: SchemaRef = Arc::new(any_dataset_to_arrow_schema(&dataset).map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to derive Arrow schema for atlas dataset: {e}" - )) - })?); - let adapter = batch_adapter_factory(projected_schema).make_adapter(&source_schema)?; - - let pushdown_filter: Option = predicate.map(PushdownFilter::new); - let stream = any_dataset_as_record_batch_stream(dataset, batch_size, pushdown_filter, metrics) - .map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Error reading atlas dataset as Arrow stream: {e}" - )) - }) - .and_then(move |batch| { - let mapped = adapter.adapt_batch(&batch).map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to adapt atlas batch schema: {e}" - )) - }); - future::ready(mapped) - }) - .boxed(); + /// The columns to build for one dataset, or `None` to build every one. + /// + /// The projection reaches the build, so an unprojected array gets no + /// backend and an unprojected attribute is never read out of the footer. + /// + /// Two cases need care. A `COUNT(*)` projects nothing, and building nothing + /// would leave the read with no grid to count; it takes the widest array of + /// the dataset instead, which is what states the row count. And a predicate + /// column is added to the set: the filter above the scan forces such a + /// column into the projection today, but the chunk pruning inside the read + /// matches columns by name and would silently stop pruning if that ever + /// changed. + fn projected_names(&self, view: &DatasetView) -> Option> { + if self.projected_schema.fields().is_empty() { + return count_driver(view).map(|driver| vec![driver]); + } - Ok(stream) + let mut names: Vec = self + .projected_schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(); + + if let Some(predicate) = &self.predicate { + for column in collect_columns(predicate) { + if !names.iter().any(|name| name == column.name()) { + names.push(column.name().to_string()); + } + } + } + Some(names) + } } -/// Build a single-column dataset over `dataset_name`'s largest readable array, -/// used only to establish the dataset's row count when the query projects no -/// column it declares. `None` if the dataset has no readable array. -async fn driver_dataset( - atlas: &Arc, - dataset_name: &str, - object_path: &object_store::path::Path, -) -> datafusion::error::Result> { - let view = atlas.open_dataset(dataset_name).await.map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to open atlas dataset '{dataset_name}' at {object_path}: {e}" - )) - })?; - let driver = view - .schema() +/// The array a `COUNT(*)` reads to establish a dataset's row count: the widest +/// one Beacon can read. +/// +/// From the footer alone, so choosing it costs no I/O and no backend. `None` +/// for a dataset with no readable array, and the caller then builds what there +/// is — an attribute-only dataset contributes the one row its scalars define. +fn count_driver(view: &DatasetView) -> Option { + view.schema() .arrays .iter() - .filter(|(_, s)| crate::compat::atlas_array_dtype_to_arrow(&s.dtype).is_some()) - .max_by_key(|(_, s)| s.shape.iter().product::()) - .map(|(name, _)| name.clone()); - - let Some(driver) = driver else { - return Ok(None); - }; - let names = [driver]; - let dataset = crate::reader::dataset_from_atlas(atlas.clone(), dataset_name, Some(&names)) - .await - .map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to read atlas dataset '{dataset_name}' at {object_path}: {e}" - )) - })?; - Ok(Some(dataset)) -} - -/// A single all-null row shaped as `schema` (0 columns → a 0-column, 1-row -/// batch, the `COUNT(*)` unit for a scalar-only dataset). -fn null_row_batch(schema: &SchemaRef) -> datafusion::error::Result { - let columns: Vec = schema - .fields() - .iter() - .map(|f| arrow::array::new_null_array(f.data_type(), 1)) - .collect(); - RecordBatch::try_new_with_options( - schema.clone(), - columns, - &RecordBatchOptions::new().with_row_count(Some(1)), - ) - .map_err(|e| { - datafusion::error::DataFusionError::Execution(format!("Failed to build null row batch: {e}")) - }) + .filter(|(_, schema)| compat::array_dtype_to_nd(&schema.dtype).is_some()) + .max_by_key(|(_, schema)| schema.shape.iter().product::()) + .map(|(name, _)| name.clone()) } -impl FileOpener for AtlasOpener { - fn open(&self, file: PartitionedFile) -> datafusion::error::Result { - let assigned_names = file +#[async_trait::async_trait] +impl OpenFile for AtlasDatasets { + async fn open(&self, file: &PartitionedFile) -> Result> { + let entry = file .extensions .as_ref() - .and_then(|ext| (ext.as_ref() as &dyn Any).downcast_ref::()) - .map(|slice| slice.names.clone()); - let fut = Self::read_task( - self.object_store.clone(), - file.object_meta, - assigned_names, - self.projected_schema.clone(), + .and_then(|extension| (extension.as_ref() as &dyn Any).downcast_ref::()) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "the atlas scan entry at '{}' names no dataset", + file.object_meta.location + )) + })?; + + let open_start = Instant::now(); + let atlas = get_or_open_atlas( + self.cache.as_ref(), + Arc::clone(&self.object_store), + &file.object_meta, + ) + .await + .map_err(|e| DataFusionError::Execution(format!("{e}")))?; + self.scan_metrics.open_time.add_elapsed(open_start); + + // One index per collection decides this, and the first opener to reach + // the collection builds it. A dataset ruled out costs one pop and no + // read at all. + if !self + .candidates(&atlas, file.object_meta.location.as_ref()) + .await + .keeps(entry.position, &entry.dataset) + { + self.scan_metrics.datasets_pruned.add(1); + return Ok(FileRead::skipped()); + } + + let build_start = Instant::now(); + let view = Arc::new(atlas.dataset(&entry.dataset).map_err(|e| { + DataFusionError::Execution(format!( + "Failed to open atlas dataset '{}' of '{}': {e}", + entry.dataset, file.object_meta.location + )) + })?); + + let projected = self.projected_names(&view); + let dataset = dataset_from_view(view, projected.as_deref()) + .await + .map_err(|e| DataFusionError::Execution(format!("{e}")))?; + // Explicit dimensions, or a broadcast-compatible default. No log label: + // this runs per dataset, and schema inference already logged the choice. + let dataset = project_read_dimensions(dataset, self.read_dimensions.clone(), None) + .map_err(|e| DataFusionError::Execution(format!("{e}")))?; + + let read = FileRead::plan( + dataset, + Arc::clone(&self.projected_schema), self.batch_size, - self.metrics.clone(), - self.partition, - self.read_dimensions.clone(), self.predicate.clone(), - self.cache.clone(), - self.use_pruning, - self.prune_cache.clone(), - ); - Ok(Box::pin(fut)) + // A dataset lives inside a container, not at a path, so no + // `PARTITIONED BY` value can be read off it. The format refuses + // such a table outright. + FilePartitions::none(), + Some(&self.read_metrics), + ) + .await?; + + self.scan_metrics + .dataset_build_time + .add_elapsed(build_start); + self.scan_metrics.datasets_scanned.add(1); + Ok(read) } } #[cfg(test)] -mod repartition_tests { - //! Atlas divides its own work, by dataset name, so it must refuse - //! DataFusion's byte-range split. +mod tests { + use super::*; + use datafusion::datasource::physical_plan::FileScanConfigBuilder; + use datafusion::execution::object_store::ObjectStoreUrl; - use std::sync::Arc; + fn source() -> AtlasSource { + AtlasSource::new( + None, + TableSchema::from_file_schema(Arc::new(arrow::datatypes::Schema::empty())), + ) + } - use datafusion::datasource::listing::PartitionedFile; - use datafusion::datasource::physical_plan::{FileScanConfigBuilder, FileSource}; - use datafusion::datasource::table_schema::TableSchema; - use datafusion::execution::object_store::ObjectStoreUrl; + fn entry(dataset: &str, position: usize) -> PartitionedFile { + let mut file = PartitionedFile::new("obs/data.atlas", 4096); + file.extensions = Some(Arc::new(AtlasEntry { + dataset: dataset.to_string(), + position, + })); + file + } + + /// Every dataset of the scan goes into one queue, and each partition gets a + /// standing entry pointing at it. + #[test] + fn the_datasets_go_into_one_queue() { + const PARTITIONS: usize = 4; - use super::AtlasSource; + let source = source(); + let mut builder = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + Arc::new(source.clone()) as Arc, + ); + for (position, name) in ["a", "b", "c", "d", "e"].iter().enumerate() { + builder = builder.with_file(entry(name, position)); + } + let config = builder.build(); - /// A store never splits by byte range. - /// - /// Every name slice of a store carries the same `object_meta`, and the - /// partitioner copies `extensions` into each share it makes. Two shares of - /// one slice would therefore read the same datasets twice. The count would - /// grow with `target_partitions`, silently. + let planned = source + .repartitioned(PARTITIONS, 10 * 1024 * 1024, None, &config) + .unwrap() + .expect("the datasets are planned across the partitions"); + + assert_eq!(planned.file_groups.len(), PARTITIONS); + for group in &planned.file_groups { + assert_eq!(group.len(), 1, "one standing entry per partition"); + } + let planned = planned + .file_source() + .as_any() + .downcast_ref::() + .expect("the config carries an AtlasSource"); + assert_eq!( + planned.morsel_datasets(), + Some(5), + "and the queue holds every dataset" + ); + } + + /// An ordered scan cannot share: a partition holding an arbitrary share of + /// the datasets cannot emit its rows in collection order. #[test] - fn a_store_never_splits_by_byte_range() { - let table_schema = - TableSchema::from_file_schema(Arc::new(arrow::datatypes::Schema::empty())); - let source = AtlasSource::new(None, table_schema); + fn an_ordered_scan_is_left_alone() { + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::expressions::Column; + + let source = source(); let config = FileScanConfigBuilder::new( ObjectStoreUrl::local_filesystem(), Arc::new(source.clone()) as Arc, ) - // Comfortably over the partitioner's minimum split size. - .with_file(PartitionedFile::new("store.atlas", 64 * 1024 * 1024)) + .with_file(entry("a", 0)) .build(); - assert!(!source.supports_repartitioning()); + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new("time", 0), + ))]); assert!( - source.repartitioned(4, 1, None, &config).unwrap().is_none(), - "an atlas store must not split by byte range" + source + .repartitioned(4, 0, ordering, &config) + .unwrap() + .is_none() ); } -} - -#[cfg(test)] -mod adapter_tests { - //! The per-dataset schema adaptation contract, exercised directly. - //! - //! [`stream_adapted`] leans entirely on `batch_adapter_factory`: a dataset is - //! read at its own native dtypes and every batch is mapped onto the merged - //! table schema. These pin that mapping — cast up, null-fill by name — without - //! building an atlas store or a DataFusion session, so a behaviour change in - //! the adapter surfaces here rather than as a wrong query result. - - use std::sync::Arc; - - use arrow::array::{Array, Int16Array, Int64Array, StringArray}; - use arrow::datatypes::{DataType, Field, Schema}; - use arrow::record_batch::RecordBatch; - use beacon_datafusion_ext::scan_adapt::batch_adapter_factory; - - /// Map `batch` (in `source`) onto `target`, exactly as `stream_adapted` does. - fn adapt( - source: Arc, - target: Arc, - batch: RecordBatch, - ) -> datafusion::error::Result { - batch_adapter_factory(target) - .make_adapter(&source)? - .adapt_batch(&batch) - } #[test] - fn casts_a_narrower_dataset_dtype_up_to_the_merged_type() { - // The widening case: one dataset stores Int16, the collection merged to Float32. - let source = Arc::new(Schema::new(vec![Field::new("value", DataType::Int16, true)])); - let target = Arc::new(Schema::new(vec![Field::new("value", DataType::Float32, true)])); - let batch = RecordBatch::try_new( - source.clone(), - vec![Arc::new(Int16Array::from(vec![1, 2]))], + fn one_partition_divides_nothing() { + let source = source(); + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + Arc::new(source.clone()) as Arc, ) - .expect("source batch"); + .with_file(entry("a", 0)) + .build(); - let out = adapt(source, target, batch).expect("adapt"); - assert_eq!(out.schema().field(0).data_type(), &DataType::Float32); - let col = out - .column(0) - .as_any() - .downcast_ref::() - .expect("Float32 column"); - assert_eq!(col.values(), &[1.0, 2.0]); + assert!(source.repartitioned(1, 0, None, &config).unwrap().is_none()); } - #[test] - fn null_fills_a_column_the_dataset_does_not_declare() { - // The dataset has `value`; the table also has `flag`, which this dataset - // lacks. Its rows must survive with `flag` null — union semantics. - let source = Arc::new(Schema::new(vec![Field::new("value", DataType::Int64, true)])); - let target = Arc::new(Schema::new(vec![ - Field::new("flag", DataType::Int32, true), - Field::new("value", DataType::Int64, true), - ])); - let batch = RecordBatch::try_new( - source.clone(), - vec![Arc::new(Int64Array::from(vec![10, 20]))], - ) - .expect("source batch"); - - let out = adapt(source, target, batch).expect("adapt"); - assert_eq!(out.num_rows(), 2, "rows are kept, not dropped"); - let flag = out.column(out.schema().index_of("flag").expect("flag")); - assert_eq!(flag.null_count(), 2, "missing column is entirely null"); - let value = out - .column(out.schema().index_of("value").expect("value")) - .as_any() - .downcast_ref::() - .expect("Int64 column"); - assert_eq!(value.values(), &[10, 20], "declared column is untouched"); + /// An entry that names no dataset is a bug in the planner, not bad input, + /// and the error says which collection it came from. + #[tokio::test] + async fn an_entry_without_a_dataset_is_an_internal_error() { + let datasets = AtlasDatasets { + object_store: Arc::new(object_store::memory::InMemory::new()), + cache: None, + projected_schema: Arc::new(arrow::datatypes::Schema::empty()), + logical_schema: Arc::new(arrow::datatypes::Schema::empty()), + use_pruning: false, + prune_cache: PruneCache::new(), + read_dimensions: None, + batch_size: usize::MAX, + predicate: None, + read_metrics: ReadMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + scan_metrics: AtlasScanMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + }; + + let error = datasets + .open(&PartitionedFile::new("obs/data.atlas", 1)) + .await + .expect_err("an entry must name its dataset") + .to_string(); + assert!(error.contains("names no dataset"), "{error}"); + assert!(error.contains("obs/data.atlas"), "{error}"); } - #[test] - fn stringifies_a_numeric_dataset_when_the_merge_widened_to_utf8() { - // Non-numeric conflict: atlas merges String ∪ Int64 to String, so the - // integer dataset must be cast *into* Utf8 rather than erroring. This is - // what makes a mixed-dtype collection readable at all. - let source = Arc::new(Schema::new(vec![Field::new("value", DataType::Int64, true)])); - let target = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); - let batch = RecordBatch::try_new( - source.clone(), - vec![Arc::new(Int64Array::from(vec![1, 2]))], - ) - .expect("source batch"); + // ── which columns a dataset is built with ─────────────────────────── - let out = adapt(source, target, batch).expect("Int64 -> Utf8 must be castable"); - let col = out - .column(0) - .as_any() - .downcast_ref::() - .expect("Utf8 column"); + use crate::test_support; + + /// One dataset's view. It owns what it needs, so the collection handle + /// behind it may go. + async fn view(dir: &std::path::Path, dataset: &str) -> DatasetView { + test_support::open(dir) + .await + .dataset(dataset) + .expect("the dataset") + } + + fn datasets_wanting( + projected: Vec<&str>, + predicate: Option>, + ) -> AtlasDatasets { + let fields: Vec = projected + .into_iter() + .map(|name| arrow::datatypes::Field::new(name, arrow::datatypes::DataType::Null, true)) + .collect(); + let projected_schema = Arc::new(arrow::datatypes::Schema::new(fields)); + AtlasDatasets { + object_store: Arc::new(object_store::memory::InMemory::new()), + cache: None, + logical_schema: Arc::clone(&projected_schema), + projected_schema, + use_pruning: false, + prune_cache: PruneCache::new(), + read_dimensions: None, + batch_size: usize::MAX, + predicate, + read_metrics: ReadMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + scan_metrics: AtlasScanMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + } + } + + #[tokio::test] + async fn a_scan_builds_the_columns_it_projects() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let datasets = datasets_wanting(vec!["temperature"], None); + let names = datasets + .projected_names(&view(tmp.path(), "winter").await) + .expect("a projection"); + assert_eq!(names, vec!["temperature".to_string()]); + } + + /// A predicate column joins the set even when the projection leaves it out. + /// The chunk pruning inside the read matches by name, so a missing column + /// would silently stop it pruning. + #[tokio::test] + async fn a_predicate_column_is_built_too() { + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; + use datafusion::scalar::ScalarValue; + + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let predicate: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("cycle", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(20)))), + )); + let datasets = datasets_wanting(vec!["temperature"], Some(predicate)); + let names = datasets + .projected_names(&view(tmp.path(), "winter").await) + .expect("a projection"); assert_eq!( - (0..col.len()).map(|i| col.value(i)).collect::>(), - vec!["1", "2"], + names, + vec!["temperature".to_string(), "cycle".to_string()], + "the predicate's column is kept alongside the projection" + ); + } + + /// `COUNT(*)` projects nothing and reads one array: the row count is a + /// property of the grid, and building every column to find it would read + /// every attribute of the dataset for nothing. + #[tokio::test] + async fn a_count_reads_the_widest_array_alone() { + let tmp = tempfile::tempdir().unwrap(); + test_support::chunked_grid(tmp.path()).await; + + let datasets = datasets_wanting(vec![], None); + let names = datasets + .projected_names(&view(tmp.path(), "grid").await) + .expect("a driver"); + assert_eq!(names.len(), 1); + assert!( + names[0] == "temperature" || names[0] == "sparse", + "either 4x6 array states the row count: {names:?}" ); } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs new file mode 100644 index 00000000..42507449 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs @@ -0,0 +1,292 @@ +//! The column ranges of a whole collection, for the file analyzer. +//! +//! A collection reports one range per column, folded over its live datasets out +//! of the footer. It costs no array read: the writer computed each dataset's +//! minimum and maximum while it staged the data, and the open already holds +//! them. +//! +//! # Why a wrong answer here is worse than no answer +//! +//! A recorded range prunes whole collections before a scan opens them, so a +//! range that is too narrow silently deletes matching rows from an answer. +//! Every path below reports unknown unless it can *prove* the bound. Unknown is +//! always legal: it only means Beacon reads what it might have skipped. + +use arrow::datatypes::{DataType, Schema}; +use atlas::{Atlas, StatValue}; +use datafusion::common::{ColumnStatistics, Statistics, stats::Precision}; +use datafusion::scalar::ScalarValue; + +/// The statistics of one collection, in `table_schema` order. +pub fn collection_statistics(atlas: &Atlas, table_schema: &Schema) -> Statistics { + let arrays = atlas.list_arrays(); + let live = atlas.dataset_count(); + + let mut statistics = Statistics::default(); + for field in table_schema.fields() { + let range = if arrays.iter().any(|array| array == field.name()) { + column_range(atlas, field.name(), field.data_type(), live) + } else { + // An attribute column would need one `DatasetView` per dataset, and + // each of those is a linear scan of the footer. The query-time + // pruning index pays that where it is bounded and worth it; a + // background pass over every collection is not the place. + None + }; + + statistics = statistics.add_column_statistics(match range { + Some((min, max)) => ColumnStatistics::new_unknown() + .with_min_value(Precision::Exact(min)) + .with_max_value(Precision::Exact(max)), + None => ColumnStatistics::new_unknown(), + }); + } + statistics +} + +/// The range of one array column over every live dataset, or `None`. +/// +/// # Every live dataset must report +/// +/// A dataset that declares an array and never writes it has no statistics +/// entry, and its cells read back as the array's fill — or, when it declares +/// none, as zeros that nothing nulls. Folding only the datasets that *do* +/// report would then produce a range those zeros sit outside of, and pruning +/// would drop the collection for a query that matches them. +/// +/// The footer cannot tell "declares it and never wrote it" from "does not +/// declare it" without a view per dataset, which is a linear scan each. So the +/// count is the proof: a bound is claimed only when every live dataset reported +/// one. A uniform collection — which is what `atlas create` writes, and what +/// this format exists for — satisfies that; a heterogeneous one goes unknown +/// and is read in full. +fn column_range( + atlas: &Atlas, + column: &str, + target: &DataType, + live: usize, +) -> Option<(ScalarValue, ScalarValue)> { + let per_dataset = atlas.array_stats_by_dataset(column); + if per_dataset.is_empty() || per_dataset.len() != live { + return None; + } + + let mut low: Option = None; + let mut high: Option = None; + for (_, stats) in per_dataset { + // One dataset without a bound leaves the column unbounded: its values + // may lie anywhere. + let min = bound(stats.min.as_ref(), target)?; + let max = bound(stats.max.as_ref(), target)?; + low = Some(match low { + Some(held) => smaller(held, min)?, + None => min, + }); + high = Some(match high { + Some(held) => larger(held, max)?, + None => max, + }); + } + Some((low?, high?)) +} + +/// One statistic as a scalar of the table's type, or `None` when it proves +/// nothing: absent, `NaN` — which sorts last and bounds nothing — or a value +/// that will not cast. +fn bound(value: Option<&StatValue>, target: &DataType) -> Option { + let canonical = match value? { + StatValue::Int(v) => ScalarValue::Int64(Some(*v)), + StatValue::UInt(v) => ScalarValue::UInt64(Some(*v)), + StatValue::Float(v) if v.is_nan() => return None, + StatValue::Float(v) => ScalarValue::Float64(Some(*v)), + StatValue::TimestampNs(v) => ScalarValue::TimestampNanosecond(Some(*v), None), + StatValue::Bytes(bytes) => match std::str::from_utf8(bytes) { + Ok(text) => ScalarValue::Utf8(Some(text.to_string())), + Err(_) => ScalarValue::Binary(Some(bytes.clone())), + }, + }; + let cast = canonical.cast_to(target).ok()?; + // A cast that lands on null has lost the value, and a null bounds nothing. + if cast.is_null() { None } else { Some(cast) } +} + +/// The lower of two bounds, or `None` when they do not compare. +fn smaller(held: ScalarValue, next: ScalarValue) -> Option { + match held.partial_cmp(&next)? { + std::cmp::Ordering::Greater => Some(next), + _ => Some(held), + } +} + +/// The higher of two bounds, or `None` when they do not compare. +fn larger(held: ScalarValue, next: ScalarValue) -> Option { + match held.partial_cmp(&next)? { + std::cmp::Ordering::Less => Some(next), + _ => Some(held), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support; + use arrow::datatypes::Field; + + fn schema(fields: Vec) -> Schema { + Schema::new(fields) + } + + fn range( + statistics: &Statistics, + index: usize, + ) -> (Precision, Precision) { + let column = &statistics.column_statistics[index]; + (column.min_value.clone(), column.max_value.clone()) + } + + /// A uniform collection — every dataset holding every array — reports the + /// union of its datasets' ranges. + #[tokio::test] + async fn a_uniform_collection_reports_the_union_of_its_datasets() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 10).await; + let atlas = test_support::open(tmp.path()).await; + + let statistics = collection_statistics( + &atlas, + &schema(vec![Field::new("temperature", DataType::Float32, true)]), + ); + let (min, max) = range(&statistics, 0); + // d0 starts at 0 and d9 ends at 93. + assert_eq!(min, Precision::Exact(ScalarValue::Float32(Some(0.0)))); + assert_eq!(max, Precision::Exact(ScalarValue::Float32(Some(93.0)))); + } + + /// A column only some datasets declare goes unknown, because a dataset that + /// declared it and never wrote it is indistinguishable from one that never + /// declared it, and the first reads back as values this fold cannot see. + #[tokio::test] + async fn a_column_not_every_dataset_reports_goes_unknown() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let statistics = collection_statistics( + &atlas, + &schema(vec![ + // winter alone declares `cycle`. + Field::new("cycle", DataType::Int32, true), + // both declare `temperature`. + Field::new("temperature", DataType::Float32, true), + ]), + ); + assert_eq!(range(&statistics, 0).0, Precision::Absent); + assert_eq!( + range(&statistics, 1).0, + Precision::Exact(ScalarValue::Float32(Some(1.0))), + "a column every dataset reports still bounds" + ); + assert_eq!( + range(&statistics, 1).1, + Precision::Exact(ScalarValue::Float32(Some(22.0))), + ); + } + + /// A column two datasets type differently is folded on the table's type. + #[tokio::test] + async fn a_mixed_dtype_column_folds_on_the_table_type() { + let tmp = tempfile::tempdir().unwrap(); + test_support::widening(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let statistics = collection_statistics( + &atlas, + &schema(vec![Field::new("value", DataType::Float64, true)]), + ); + let (min, max) = range(&statistics, 0); + // a holds [1, 2] as Int16 and b holds [3.5, 4.5] as Float32. + assert_eq!(min, Precision::Exact(ScalarValue::Float64(Some(1.0)))); + assert_eq!(max, Precision::Exact(ScalarValue::Float64(Some(4.5)))); + } + + /// A column the collection does not hold, and a column whose values are + /// text, both report unknown rather than a guess. + #[tokio::test] + async fn an_unknown_column_reports_unknown() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 4).await; + let atlas = test_support::open(tmp.path()).await; + + let statistics = collection_statistics( + &atlas, + &schema(vec![ + Field::new("ghost", DataType::Float32, true), + Field::new(".platform", DataType::Utf8, true), + ]), + ); + assert_eq!(range(&statistics, 0).0, Precision::Absent); + assert_eq!( + range(&statistics, 1).0, + Precision::Absent, + "an attribute is not measured here" + ); + } + + #[tokio::test] + async fn an_empty_collection_bounds_nothing() { + let tmp = tempfile::tempdir().unwrap(); + test_support::empty(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let statistics = collection_statistics( + &atlas, + &schema(vec![Field::new("temperature", DataType::Float32, true)]), + ); + assert_eq!(range(&statistics, 0).0, Precision::Absent); + } + + /// A deleted dataset counts toward nothing: neither the fold nor the count + /// that guards it. + #[tokio::test] + async fn a_deleted_dataset_leaves_the_range() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 10).await; + let atlas = test_support::open(tmp.path()).await; + atlas.delete_dataset("d9").await.unwrap(); + + // Reopen, so the handle reads the mask that was just written. + let atlas = test_support::open(tmp.path()).await; + let statistics = collection_statistics( + &atlas, + &schema(vec![Field::new("temperature", DataType::Float32, true)]), + ); + assert_eq!( + range(&statistics, 0).1, + Precision::Exact(ScalarValue::Float32(Some(83.0))), + "d9 held the values up to 93 and is gone" + ); + } + + // ── the pieces ────────────────────────────────────────────────────── + + #[test] + fn a_nan_bound_proves_nothing() { + assert!(bound(Some(&StatValue::Float(f64::NAN)), &DataType::Float64).is_none()); + } + + #[test] + fn a_bound_that_will_not_cast_proves_nothing() { + assert!( + bound( + Some(&StatValue::Bytes(b"argo".to_vec())), + &DataType::Float64 + ) + .is_none() + ); + } + + #[test] + fn an_absent_bound_proves_nothing() { + assert!(bound(None, &DataType::Float64).is_none()); + } +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/table_function.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/table_function.rs index 5ce3433a..6de67997 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/table_function.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/table_function.rs @@ -1,22 +1,24 @@ +//! `read_atlas(paths)` and `read_atlas(paths, dimensions)`. + use std::collections::HashMap; use std::sync::{Arc, Weak}; use arrow::datatypes::{DataType, Field}; +use beacon_common::table_function::BeaconTableFunctionImpl; use beacon_datafusion_ext::fast_object::FastObjectTable; +use beacon_datafusion_ext::listing_factory::ListingFactory; use datafusion::{ - catalog::TableFunctionImpl, - common::plan_err, + catalog::{TableFunctionImpl, TableProvider}, + common::{plan_datafusion_err, plan_err}, + error::Result, prelude::{Expr, SessionContext}, scalar::ScalarValue, }; -use beacon_common::table_function::BeaconTableFunctionImpl; - -/// Format identity the atlas factory is registered under (its `get_ext`). -const ATLAS_FORMAT: &str = "atlas"; +use crate::datafusion::ATLAS_FORMAT; +/// Reads the Atlas collections that match one or more glob patterns. pub struct ReadAtlasFunc { - // Session Reference runtime_handle: tokio::runtime::Handle, session_ctx: Weak, } @@ -43,9 +45,10 @@ impl BeaconTableFunctionImpl for ReadAtlasFunc { fn description(&self) -> Option { Some( - "Reads atlas stores. Each path must point to an atlas.json marker file \ - (exact path or glob like **/atlas.json). Optional second arg filters \ - arrays to those matching the listed dimensions." + "Reads Atlas collections. Each path names a 'data.atlas' container file, exactly or \ + through a glob such as '**/data.atlas'. The optional second argument lists the \ + dimensions to read, and an array survives only when the list holds every one of its \ + own." .to_string(), ) } @@ -54,7 +57,7 @@ impl BeaconTableFunctionImpl for ReadAtlasFunc { "read_atlas".to_string() } - fn arguments(&self) -> Option> { + fn arguments(&self) -> Option> { Some(vec![ Field::new( "glob_paths", @@ -71,81 +74,70 @@ impl BeaconTableFunctionImpl for ReadAtlasFunc { } impl TableFunctionImpl for ReadAtlasFunc { - fn call( - &self, - args: &[datafusion::prelude::Expr], - ) -> datafusion::error::Result> { + fn call(&self, args: &[Expr]) -> Result> { let glob_paths = beacon_common::table_function::parse_glob_paths_arg(args, "read_atlas")?; let mut dimensions: Vec = vec![]; - if let Some(dimensions_arg) = args.get(1) { - if let Expr::Literal(ScalarValue::List(values), _) = dimensions_arg { - let string_array = values.as_ref().values(); - match string_array - .as_any() - .downcast_ref::() - { - Some(str_arr) => { - dimensions = str_arr - .iter() - .filter_map(|opt_str| opt_str.map(|s| s.to_string())) - .collect(); - } - None => { - return plan_err!( - "read_atlas second argument must be a List of dimension names" - ); - } + if let Some(argument) = args.get(1) + && let Expr::Literal(ScalarValue::List(values), _) = argument + { + match values + .as_ref() + .values() + .as_any() + .downcast_ref::() + { + Some(names) => { + dimensions = names + .iter() + .filter_map(|name| name.map(str::to_string)) + .collect(); + } + None => { + return plan_err!( + "read_atlas second argument must be a List of dimension names" + ); } } } - tracing::debug!("read_atlas glob paths: {:?}", glob_paths); + tracing::debug!("read_atlas glob paths: {glob_paths:?}"); - let session_ctx = self.session_ctx.upgrade().ok_or_else(|| { - datafusion::common::plan_datafusion_err!("session context has been dropped") - })?; + let session_ctx = self + .session_ctx + .upgrade() + .ok_or_else(|| plan_datafusion_err!("session context has been dropped"))?; let state = session_ctx.state(); + let listing_factory = state .config() - .get_extension::() + .get_extension::() .ok_or_else(|| { - datafusion::error::DataFusionError::Execution( - "read_atlas: the listing factory is not registered on the session".to_string(), - ) + plan_datafusion_err!("read_atlas: the listing factory is not registered") })?; - - let mut listing_urls = vec![]; + let mut listing_urls = Vec::with_capacity(glob_paths.len()); for path in &glob_paths { - tracing::debug!("read_atlas processing path: {}", path); listing_urls.push(listing_factory.parse_listing_table_url(&state, path)?); } - // Build the file format from the factory registered on the session, so the - // table function shares the runtime's configured format + reader cache. - // Per-call settings (read dimensions) are passed as table options. + // Build the format from the factory registered on the session, so the + // function shares the runtime's settings and its reader cache. The + // per-call dimensions ride along as a table option. let mut format_options: HashMap = HashMap::new(); if !dimensions.is_empty() { format_options.insert("read_dimensions".to_string(), dimensions.join(",")); } - - let session_ctx = self.session_ctx.upgrade().ok_or_else(|| { - datafusion::common::plan_datafusion_err!("session context has been dropped") - })?; - let state = session_ctx.state(); let factory = state.get_file_format_factory(ATLAS_FORMAT).ok_or_else(|| { - datafusion::error::DataFusionError::Execution( - "read_atlas: the atlas file format is not registered on the session".to_string(), - ) + plan_datafusion_err!("read_atlas: the atlas file format is not registered") })?; let file_format = factory.create(&state, &format_options)?; - let fast_object_table = tokio::task::block_in_place(|| { + let table = tokio::task::block_in_place(|| { self.runtime_handle.block_on(async { FastObjectTable::try_new(&session_ctx.state(), file_format, listing_urls).await }) })?; - Ok(Arc::new(fast_object_table)) + Ok(Arc::new(table)) } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs index 31aae712..b89054e6 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs @@ -1,27 +1,69 @@ -//! `beacon-arrow-atlas` bridges Atlas array stores into Beacon's shared +//! `beacon-arrow-atlas` reads Atlas collections through Beacon's shared //! `beacon-nd-array` engine. //! -//! Atlas () is a directory-based -//! store where a single metadata file (`atlas.json` / `atlas.msgpack`, with -//! optional `.zst` / `.lz4` suffix) describes one or more named datasets, each -//! a collection of N-dimensional arrays plus per-dataset and per-array -//! attributes. As of atlas 0.14 the store is opened directly over any -//! [`object_store`] backend (local filesystem, S3, GCS, Azure) — no path -//! translation or native-filesystem root is required. -//! -//! This crate mirrors `beacon-arrow-zarr`: the DataFusion -//! integration ([`datafusion`]) discovers atlas metadata markers, opens each -//! store as an atlas *collection* straight from the query's object store, and -//! exposes every dataset as a Beacon -//! [`AnyDataset`](beacon_nd_array::dataset::AnyDataset) via lazy -//! [`NdArrayD`](beacon_nd_array::NdArrayD) backends ([`backend`]). Each atlas -//! array becomes a column; dataset-level attributes become rank-0 columns and -//! per-array attributes become `{array}.{attr}` rank-0 columns. +//! # The format +//! +//! An Atlas collection () is one +//! immutable file, `data.atlas`, with an optional `deleted.mask` beside it: +//! +//! ```text +//! my_collection/ +//! ├── data.atlas ATLS │ segment │ segment │ … │ footer │ trailer +//! └── deleted.mask optional: ordinals of deleted datasets +//! ``` +//! +//! Each dataset occupies one segment. A footer at the end records every dataset +//! name, its segment byte range, its schema, its attribute values and its +//! per-array statistics. Opening a collection reads that footer and nothing +//! else, so every metadata question is answered with no further I/O, whatever +//! the dataset count. Array data arrives chunk by chunk, on demand. +//! +//! # What this crate does with it +//! +//! [`store`] finds a collection's marker and opens it, through a reader cache. +//! [`reader`] turns one dataset into a Beacon +//! [`AnyDataset`](beacon_nd_array::dataset::AnyDataset) whose columns are lazy +//! [`NdArrayD`](beacon_nd_array::NdArrayD) values backed by [`backend`], and +//! derives the Arrow schema of a whole collection. [`compat`] holds the type +//! and column-name mapping the two share. +//! +//! # Columns +//! +//! One column per array, under the array's own name. A per-array attribute +//! becomes `{array}.{attr}`, and a dataset-level attribute becomes `.{attr}`. +//! That is the convention netCDF and Zarr use, so a query reads the same +//! whichever format holds the data. +//! +//! # What is not read +//! +//! - A `Bool` array, and a `List` or `FixedSizeList` array. `array-format` +//! stores no element of those types, so no such array can exist in a +//! collection a Rust writer produced. The mapping refuses them all the same. +//! - A list-valued *attribute*. Beacon's ND model has no rank-0 list. +//! +//! Each is dropped from the dataset with a `debug` log rather than failing the +//! scan. A collection can hold a million datasets, so a `warn` per skip would +//! be a flood. +//! +//! # No CF decoding +//! +//! Atlas has a native timestamp type, and the ingest path (`atlas create`) +//! applies `scale_factor`, `add_offset` and the CF time units *before* the +//! write. An atlas array is therefore read exactly as it is stored, unlike +//! netCDF and Zarr. A collection written by hand with packed integers and a CF +//! `units` attribute reads back as those integers. pub use atlas; pub mod backend; pub mod compat; +pub mod config; pub mod datafusion; pub mod reader; -pub mod util; +pub mod store; + +pub use config::AtlasConfig; +pub use datafusion::{AtlasFormat, AtlasFormatFactory, AtlasOptions, ReadAtlasFunc}; + +#[cfg(test)] +pub(crate) mod test_support; diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs index ffb5fb9a..fbf3c184 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs @@ -1,492 +1,702 @@ -//! High-level atlas reader that produces [`AnyDataset`] values. +//! Turning an Atlas collection into what Beacon's engine reads: one dataset as +//! an [`AnyDataset`] of lazy columns, and a whole collection as one Arrow +//! schema. //! -//! An atlas store holds one or more named datasets. [`dataset_from_atlas`] -//! returns the contents of **one** atlas dataset as a Beacon [`AnyDataset`], -//! wrapping every array in a lazy [`NdArrayD`](beacon_nd_array::NdArrayD) -//! backend and every scalar attribute in a rank-0 backend. Array *data* is read -//! on demand — only metadata and (cheap) attribute values are touched here. -//! -//! # Column naming -//! - Each array becomes a column under its own name. -//! - Dataset-level (global) attributes become rank-0 columns under their bare -//! attribute name. -//! - Per-array attributes become rank-0 columns named `{array}.{attr}`. +//! Nothing here reads array data. A dataset is built from the collection +//! footer, which the open already held, and its columns fetch their bytes when +//! the scan asks for them. +use std::collections::HashSet; use std::sync::Arc; +use arrow::datatypes::{Schema, SchemaRef}; +use atlas::{Atlas, DatasetView}; +use beacon_datafusion_ext::type_widening::{ArrowTypeWidening, LabeledSchema}; use beacon_nd_array::{ NdArrayD, - dataset::{AnyDataset, Dataset}, + arrow::schema::any_dataset_to_arrow_schema, + dataset::{AnyDataset, Dataset, resolve_read_dimensions}, + projection::DatasetProjection, }; use indexmap::IndexMap; -use object_store::{ObjectStore, path::Path as OsPath}; +use object_store::{ObjectMeta, ObjectStore}; use crate::compat; -/// Open an atlas store over `store` rooted at `prefix` (the directory holding -/// the `atlas.json` marker) and read the named dataset into an [`AnyDataset`]. +/// Build one dataset as an [`AnyDataset`] of lazy columns. /// -/// A convenience for callers holding an object store; the DataFusion code path -/// opens the [`Atlas`](atlas::Atlas) handle once per store and calls -/// [`dataset_from_atlas`] per dataset directly. -pub async fn open_dataset( - store: Arc, - prefix: OsPath, - dataset_name: &str, -) -> anyhow::Result { - let atlas = atlas::Atlas::open(store, prefix.clone()) - .await - .map_err(|e| anyhow::anyhow!("Failed to open atlas store at {:?}: {}", prefix, e))?; - dataset_from_atlas(Arc::new(atlas), dataset_name, None).await -} - -/// Build an [`AnyDataset`] from an already-open atlas handle. +/// `projected_names` is the column set the query wants: +/// +/// - `None` — every array and attribute. +/// - `Some(names)` — only the columns named. A name the dataset does not hold +/// is ignored, so a projection may name columns from any dataset of the +/// collection. +/// +/// The projection reaches the *build*, not just the result: an array outside it +/// gets no backend, and an attribute outside it is not even read out of the +/// footer. A column-subset query over a wide dataset therefore pays nothing for +/// the columns it did not ask for. /// -/// `projected_names`: -/// - `None` — include every array and attribute in the dataset. -/// - `Some(names)` — include only arrays/attributes whose column name appears -/// in `names`. Names not present in the dataset are silently ignored. This -/// lets the DataFusion source skip building backends for columns the query -/// won't use. -pub async fn dataset_from_atlas( - atlas: Arc, - dataset_name: &str, +/// A value Beacon cannot surface — a `Bool` or list array, a list attribute — +/// is dropped with a `debug` log. A collection can hold a million datasets, so +/// a louder log would be a flood. +pub async fn dataset_from_view( + view: Arc, projected_names: Option<&[String]>, ) -> anyhow::Result { - let view = atlas - .open_dataset(dataset_name) - .await - .map_err(|e| anyhow::anyhow!("Failed to open atlas dataset '{}': {}", dataset_name, e))?; - - let included = - |name: &str| projected_names.map_or(true, |names| names.iter().any(|n| n == name)); + let included = |name: &str| projected_names.is_none_or(|names| names.iter().any(|n| n == name)); + // Whether any projected column could be an attribute of this array. It + // saves building an attribute map the projection would throw away. + let wants_attrs_of = |array: &str| { + projected_names.is_none_or(|names| { + names + .iter() + .any(|name| compat::is_attr_column_of(name, array)) + }) + }; + let wants_global_attrs = + projected_names.is_none_or(|names| names.iter().any(|name| name.starts_with('.'))); - // The dataset's schema is in-memory metadata — array names/types and the - // attribute-key namespace — so we can decide what's projected *before* - // touching any `.af` file. Attribute values (`get_attribute` / - // `get_array_attribute`) and fill values are only read for columns the - // projection actually keeps: a column-subset query over a wide dataset - // never pays for the attributes it didn't ask for. let schema = view.schema(); - let mut arrays: IndexMap> = IndexMap::new(); + let mut arrays: IndexMap> = + IndexMap::with_capacity(schema.arrays.len()); - // ── Arrays and their per-array attributes ──────────────────────────── for (array_name, array_schema) in &schema.arrays { - // Per-array attributes ride alongside the array as `{array}.{attr}` - // columns, independent of whether the array data itself is projected. - if let Some(attr_keys) = schema.array_attrs.get(array_name) { - for attr_key in attr_keys.keys() { - let key = format!("{array_name}.{attr_key}"); - if !included(&key) { - continue; - } - let Some(attr_value) = - view.get_array_attribute(array_name, attr_key).await.map_err(|e| { - anyhow::anyhow!( - "Failed to read attribute '{key}' in atlas dataset '{dataset_name}': {e}" - ) - })? - else { - continue; // key declared in the namespace but unset here - }; - match compat::attribute_to_nd_array(&attr_value) { - Ok(nd) => { - arrays.insert(key, nd); - } - Err(e) => tracing::warn!( - "Skipping atlas array attribute '{key}' in dataset '{dataset_name}': {e}" - ), + if included(array_name) { + match compat::array_to_nd_array(Arc::clone(&view), array_name, array_schema) { + Ok(nd) => { + arrays.insert(array_name.clone(), nd); } + Err(e) => tracing::debug!( + dataset = %view.name(), + array = %array_name, + "atlas array left out of the dataset: {e}" + ), } } - if !included(array_name) { + if !wants_attrs_of(array_name) { continue; } - let fill_value = view.array_fill_value(array_name).await.map_err(|e| { - anyhow::anyhow!( - "Failed to read fill value for atlas array '{}' in dataset '{}': {}", - array_name, - dataset_name, - e - ) - })?; - match compat::array_to_nd_array(atlas.clone(), dataset_name, array_name, array_schema, fill_value) - { - Ok(nd) => { - arrays.insert(array_name.clone(), nd); + for (key, value) in view.array_attributes(array_name) { + let column = compat::array_attr_column(array_name, &key); + if !included(&column) { + continue; } - Err(e) => { - tracing::warn!("Skipping atlas array '{array_name}' in dataset '{dataset_name}': {e}") + match compat::attribute_to_nd_array(&value) { + Ok(nd) => { + arrays.insert(column, nd); + } + Err(e) => tracing::debug!( + dataset = %view.name(), + column = %column, + "atlas attribute left out of the dataset: {e}" + ), } } } - // ── Dataset-level (global) attributes ──────────────────────────────── - for attr_key in schema.global_attrs.keys() { - if !included(attr_key) { - continue; - } - let Some(attr_value) = view.get_attribute(attr_key).await.map_err(|e| { - anyhow::anyhow!( - "Failed to read global attribute '{attr_key}' in atlas dataset '{dataset_name}': {e}" - ) - })? - else { - continue; - }; - match compat::attribute_to_nd_array(&attr_value) { - Ok(nd) => { - arrays.insert(attr_key.clone(), nd); + if wants_global_attrs { + for (key, value) in view.attributes() { + let column = compat::global_attr_column(&key); + if !included(&column) { + continue; + } + match compat::attribute_to_nd_array(&value) { + Ok(nd) => { + arrays.insert(column, nd); + } + Err(e) => tracing::debug!( + dataset = %view.name(), + column = %column, + "atlas attribute left out of the dataset: {e}" + ), } - Err(e) => tracing::warn!( - "Skipping atlas global attribute '{attr_key}' in dataset '{dataset_name}': {e}" - ), } } arrays.sort_keys(); - let dataset = Dataset::new(dataset_name.to_string(), arrays).await; + let dataset = Dataset::new(view.name().to_string(), arrays).await; AnyDataset::try_from_dataset(dataset) .await - .map_err(|e| anyhow::anyhow!("Failed to wrap atlas dataset as AnyDataset: {}", e)) + .map_err(|e| anyhow::anyhow!("Failed to wrap atlas dataset '{}': {e}", view.name())) } -#[cfg(test)] -pub(crate) mod test_support { - //! Helpers for building atlas store fixtures in tests across the crate. - - use atlas::{Atlas, Attr, FillValue, StoreConfig}; - use std::path::Path; - - /// Build a two-dataset atlas store at `path`. - /// - /// Layout: - /// - `winter`: arrays `temperature: Float32[4]`, `cycle: Int32[4]` - /// (fill_value = -1, lets us assert fill propagation end-to-end); - /// global attributes `season: String("winter")`, `year: Int64(2024)`. - /// - `summer`: array `temperature: Float32[3]`; - /// global attribute `season: String("summer")`. - pub async fn build_two_dataset_store(path: &Path) { - build_two_dataset_store_with_config(path, StoreConfig::default()).await; +/// Narrow `dataset` to `read_dimensions`, or to a broadcast-compatible default +/// when none are given. +/// +/// Without this a `SELECT *` over a dataset whose arrays live on incompatible +/// dimension sets could not broadcast onto one grid. `log_label` names the +/// caller in the auto-selection log; pass `None` from per-dataset code, where +/// schema inference has already logged the choice. +pub fn project_read_dimensions( + dataset: AnyDataset, + read_dimensions: Option>, + log_label: Option<&str>, +) -> anyhow::Result { + match resolve_read_dimensions(&dataset, read_dimensions, log_label) { + Some(dims) => dataset + .project(&DatasetProjection::new_with_dimension_projection(dims)) + .map_err(|e| anyhow::anyhow!("Failed to project the atlas dataset by dimension: {e}")), + None => Ok(dataset), } +} - /// Same as [`build_two_dataset_store`] but lets the caller pick the - /// metadata format / compression so tests can exercise non-default atlas - /// marker filenames (e.g. `atlas.msgpack.zst`). - pub async fn build_two_dataset_store_with_config(path: &Path, config: StoreConfig) { - let mut atlas = Atlas::create_path(path, config) - .await - .expect("create atlas store"); - - // ── winter ──────────────────────────────────────────────────── - { - let mut winter = atlas.create_dataset("winter").await.expect("create winter"); - winter - .define_array::("temperature", vec!["obs".into()], vec![4], None, None) - .await - .expect("define winter.temperature"); - winter - .define_array::( - "cycle", - vec!["obs".into()], - vec![4], - None, - Some(FillValue::Int(-1)), - ) - .await - .expect("define winter.cycle"); - winter - .set_attribute("season", Attr::String("winter".into())) - .expect("set winter.season"); - winter - .set_attribute("year", Attr::Int64(2024)) - .expect("set winter.year"); - - let temps = ndarray::arr1(&[1.0f32, 2.0, 3.0, 4.0]).into_dyn(); - winter - .write_array("temperature", vec![0], temps.view()) - .await - .expect("write winter.temperature"); - let cycles = ndarray::arr1(&[10i32, 20, 30, 40]).into_dyn(); - winter - .write_array("cycle", vec![0], cycles.view()) - .await - .expect("write winter.cycle"); - } +/// Build one dataset of the collection at `marker`, over `store`. +/// +/// A convenience for a caller holding an object store. The scan opens the +/// collection once and calls [`dataset_from_view`] per dataset instead. +pub async fn open_dataset( + store: Arc, + marker: &ObjectMeta, + dataset: &str, +) -> anyhow::Result { + let atlas = crate::store::get_or_open_atlas(None, store, marker).await?; + let view = atlas + .dataset(dataset) + .map_err(|e| anyhow::anyhow!("Failed to open atlas dataset '{dataset}': {e}"))?; + dataset_from_view(Arc::new(view), None).await +} - // ── summer ──────────────────────────────────────────────────── - { - let mut summer = atlas.create_dataset("summer").await.expect("create summer"); - summer - .define_array::("temperature", vec!["obs".into()], vec![3], None, None) - .await - .expect("define summer.temperature"); - summer - .set_attribute("season", Attr::String("summer".into())) - .expect("set summer.season"); - - let temps = ndarray::arr1(&[20.0f32, 21.0, 22.0]).into_dyn(); - summer - .write_array("temperature", vec![0], temps.view()) - .await - .expect("write summer.temperature"); +/// The Arrow schema of a whole collection: every live dataset, merged. +/// +/// Atlas reconciles nothing. Two datasets may declare one array name with two +/// dtypes, so the collection's schema is the widening merge of its datasets' +/// schemas, under the rule the session carries. +/// +/// # One schema per shape, not per dataset +/// +/// Datasets that declare the same arrays share one interned schema in the +/// footer, and `atlas create` writes a fleet of files that way. Each dataset is +/// therefore reduced to a key over its interned schema and its attribute +/// namespace, and a schema is derived once per distinct key. A thousand +/// datasets of one shape cost one derivation. +/// +/// # Cost +/// +/// No I/O: everything read here came in with the footer. The pass is linear in +/// the dataset count, but resolving each dataset by name is itself a linear +/// scan of the footer, so the whole pass is quadratic in the dataset count. +/// That is a limit of the reader's API, which offers no lookup by ordinal, and +/// it is why a collection of more than a few tens of thousands of datasets +/// needs `Atlas::dataset_at` upstream. The result is cached above this crate, +/// so a table pays it once rather than once per query. +pub async fn collection_schema( + atlas: &Arc, + read_dimensions: Option<&[String]>, + label: &str, + widening: &ArrowTypeWidening, +) -> anyhow::Result { + let mut seen: HashSet = HashSet::new(); + let mut schemas: Vec = Vec::new(); + + for name in atlas.list_datasets() { + let view = atlas + .dataset(&name) + .map_err(|e| anyhow::anyhow!("Failed to open atlas dataset '{name}': {e}"))?; + + if !seen.insert(shape_key(&view)) { + continue; } - // Persist the metadata marker + array files to the store. - atlas.flush().await.expect("flush atlas store"); + let dataset = dataset_from_view(Arc::new(view), None).await?; + // The same narrowing the scan applies, so the schema states what a + // query can actually return. + let dataset = + project_read_dimensions(dataset, read_dimensions.map(<[String]>::to_vec), None)?; + let schema = any_dataset_to_arrow_schema(&dataset).map_err(|e| { + anyhow::anyhow!("Failed to derive the Arrow schema of atlas dataset '{name}': {e}") + })?; + // The dataset names the schema, so a refused column names both sides. + schemas.push(LabeledSchema::new( + Arc::new(schema), + format!("{label}#{name}"), + )); } - /// Build a store where the same array name has *different* dtypes across - /// datasets, so the collection's merged (table) type widens past either - /// dataset's own type: - /// - `a`: `value: Int16[2] = [1, 2]` - /// - `b`: `value: Float32[2] = [3.5, 4.5]` - /// - /// Merged `value` widens to `Float64`; each dataset must be read at its own - /// dtype and cast up. `a` also carries `flag: Int32[2]` that `b` lacks, to - /// exercise null-filling the missing column. - pub async fn build_widening_store(path: &Path) { - let mut atlas = Atlas::create_path(path, StoreConfig::default()) - .await - .expect("create atlas store"); - { - let mut a = atlas.create_dataset("a").await.expect("create a"); - a.define_array::("value", vec!["obs".into()], vec![2], None, None) - .await - .expect("define a.value"); - a.define_array::("flag", vec!["obs".into()], vec![2], None, None) - .await - .expect("define a.flag"); - a.write_array("value", vec![0], ndarray::arr1(&[1i16, 2]).into_dyn().view()) - .await - .expect("write a.value"); - a.write_array("flag", vec![0], ndarray::arr1(&[7i32, 8]).into_dyn().view()) - .await - .expect("write a.flag"); - } - { - let mut b = atlas.create_dataset("b").await.expect("create b"); - b.define_array::("value", vec!["obs".into()], vec![2], None, None) - .await - .expect("define b.value"); - b.write_array("value", vec![0], ndarray::arr1(&[3.5f32, 4.5]).into_dyn().view()) - .await - .expect("write b.value"); - } - atlas.flush().await.expect("flush atlas store"); + if schemas.is_empty() { + // A collection with no live dataset has no column. That is legal, and + // an empty schema is what every other reader answers with. + return Ok(Arc::new(Schema::empty())); } - /// Build a store whose two datasets give the *same* array genuinely - /// incompatible dtypes: - /// - /// - `a`: `value: String[2] = ["x", "y"]` - /// - `b`: `value: Int64[2] = [1, 2]` - /// - /// Unlike [`build_widening_store`], there is no numeric super-type here, so - /// this pins what the merged schema resolves to and whether a scan can still - /// read both datasets. `a` also carries `only_a: Int32[2]`, so the - /// "dataset declares none of the projected columns" path can be exercised by - /// projecting just that column. - pub async fn build_incompatible_store(path: &Path) { - let mut atlas = Atlas::create_path(path, StoreConfig::default()) - .await - .expect("create atlas store"); - { - let mut a = atlas.create_dataset("a").await.expect("create a"); - a.define_array::("value", vec!["obs".into()], vec![2], None, None) - .await - .expect("define a.value"); - a.define_array::("only_a", vec!["obs".into()], vec![2], None, None) - .await - .expect("define a.only_a"); - a.write_array( - "value", - vec![0], - ndarray::arr1(&["x".to_string(), "y".to_string()]) - .into_dyn() - .view(), - ) - .await - .expect("write a.value"); - a.write_array("only_a", vec![0], ndarray::arr1(&[7i32, 8]).into_dyn().view()) - .await - .expect("write a.only_a"); - } - { - let mut b = atlas.create_dataset("b").await.expect("create b"); - b.define_array::("value", vec!["obs".into()], vec![2], None, None) - .await - .expect("define b.value"); - b.write_array("value", vec![0], ndarray::arr1(&[1i64, 2]).into_dyn().view()) - .await - .expect("write b.value"); - } - atlas.flush().await.expect("flush atlas store"); - } + widening + .merge_schemas(&schemas) + .map_err(|e| anyhow::anyhow!("Failed to merge the schemas of the atlas datasets: {e}")) +} - /// Build a store of `n` datasets each holding `temperature: Float32[4]`, - /// where dataset `i` covers the disjoint range `[10*i, 10*i + 3]`. A - /// predicate like `temperature > T` then matches only the datasets whose - /// range reaches past `T`, so pruning can be checked against a known answer. - pub async fn build_ranged_store(path: &Path, n: usize) { - let mut atlas = Atlas::create_path(path, StoreConfig::default()) - .await - .expect("create atlas store"); - for i in 0..n { - let mut ds = atlas - .create_dataset(&format!("d{i}")) - .await - .expect("create dataset"); - ds.define_array::("temperature", vec!["obs".into()], vec![4], None, None) - .await - .expect("define temperature"); - let base = (10 * i) as f32; - let data = ndarray::arr1(&[base, base + 1.0, base + 2.0, base + 3.0]).into_dyn(); - ds.write_array("temperature", vec![0], data.view()) - .await - .expect("write temperature"); +/// What makes two datasets produce the same columns and types. +/// +/// The interned schema decides the arrays, and it is shared by address between +/// datasets that declare the same ones, so its pointer is the cheap half of the +/// key. Attribute *values* live outside the schema, so the keys and their types +/// are the other half. Two datasets that differ only in an attribute's value +/// share a key, which is exactly the fleet case. +fn shape_key(view: &DatasetView) -> String { + let schema = view.schema(); + let mut key = format!("{:x}", std::ptr::from_ref(schema) as usize); + + for array in schema.arrays.keys() { + for (attr, value) in view.array_attributes(array) { + key.push('|'); + key.push_str(array); + key.push('.'); + key.push_str(&attr); + key.push(':'); + key.push_str(&compat::dtype_tag(&value.dtype())); } - atlas.flush().await.expect("flush atlas store"); } + for (attr, value) in view.attributes() { + key.push_str("|."); + key.push_str(&attr); + key.push(':'); + key.push_str(&compat::dtype_tag(&value.dtype())); + } + key } #[cfg(test)] mod tests { - use super::test_support::build_two_dataset_store; use super::*; - use beacon_nd_array::NdArray; - use object_store::local::LocalFileSystem; + use crate::test_support; + use arrow::datatypes::DataType; + use beacon_nd_array::{NdArray, datatypes::TimestampNanosecond}; + + async fn view(dir: &std::path::Path, dataset: &str) -> Arc { + let atlas = test_support::open(dir).await; + Arc::new( + atlas + .dataset(dataset) + .expect("the dataset is in the collection"), + ) + } - /// Open a fixture dataset via the object-store-native path. - async fn open_fixture_dataset(dir: &std::path::Path, dataset: &str) -> AnyDataset { - let store: Arc = Arc::new(LocalFileSystem::new_with_prefix(dir).unwrap()); - open_dataset(store, OsPath::from(""), dataset) - .await - .expect("open dataset") + fn widening() -> Arc { + ArrowTypeWidening::default_extension() + } + + fn names(dataset: &AnyDataset) -> Vec { + dataset.fields().keys().cloned().collect() } + // ── one dataset ───────────────────────────────────────────────────── + #[tokio::test] - async fn open_dataset_lists_arrays_and_attributes() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; - - let winter = open_fixture_dataset(tmp.path(), "winter").await; - assert_eq!(winter.name(), "winter"); - - let ds = winter.dataset(); - let names: Vec<&str> = ds.arrays.keys().map(|s| s.as_str()).collect(); - assert!(names.contains(&"temperature"), "{names:?}"); - assert!(names.contains(&"cycle"), "{names:?}"); - assert!(names.contains(&"season"), "{names:?}"); - assert!(names.contains(&"year"), "{names:?}"); + async fn a_dataset_holds_its_arrays_and_its_attributes() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) + .await + .unwrap(); + + assert_eq!(dataset.name(), "winter"); + assert_eq!( + names(&dataset), + vec![ + ".season", + ".year", + "cycle", + "temperature", + "temperature.units", + "time", + ] + ); } + /// A dataset attribute takes a leading dot, and an array attribute takes its + /// array's name. That is what netCDF and Zarr do, and it keeps an attribute + /// from colliding with an array of the same name. #[tokio::test] - async fn open_dataset_reads_array_values() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; + async fn attributes_are_named_the_way_every_nd_format_names_them() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; - let winter = open_fixture_dataset(tmp.path(), "winter").await; - let temp = winter - .get_array("temperature") - .expect("temperature array") + let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) + .await + .unwrap(); + + let season = dataset + .get_array(".season") + .expect("the dataset attribute is a column") .as_any() - .downcast_ref::>() - .expect("downcast f32"); - assert_eq!(temp.clone_into_raw_vec().await, vec![1.0f32, 2.0, 3.0, 4.0]); + .downcast_ref::>() + .expect("a string column"); + assert!(season.shape().is_empty(), "an attribute has no axis"); + assert_eq!( + season.clone_into_raw_vec().await, + vec!["winter".to_string()] + ); - let cycle = winter - .get_array("cycle") - .expect("cycle array") + let units = dataset + .get_array("temperature.units") + .expect("the array attribute is a column") .as_any() - .downcast_ref::>() - .expect("downcast i32"); - assert_eq!(cycle.clone_into_raw_vec().await, vec![10i32, 20, 30, 40]); + .downcast_ref::>() + .expect("a string column"); + assert_eq!( + units.clone_into_raw_vec().await, + vec!["celsius".to_string()] + ); } #[tokio::test] - async fn open_dataset_reads_attributes_as_rank_zero() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; - - let winter = open_fixture_dataset(tmp.path(), "winter").await; - let season = winter - .get_array("season") - .expect("season attribute") + async fn every_column_keeps_the_type_the_footer_gave_it() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) + .await + .unwrap(); + let schema = any_dataset_to_arrow_schema(&dataset).unwrap(); + let field = |name: &str| schema.field_with_name(name).unwrap().data_type().clone(); + + assert_eq!(field("temperature"), DataType::Float32); + assert_eq!(field("cycle"), DataType::Int32); + assert_eq!( + field("time"), + DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None) + ); + assert_eq!(field(".year"), DataType::Int64); + assert_eq!(field(".season"), DataType::Utf8); + } + + #[tokio::test] + async fn array_values_read_back() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) + .await + .unwrap(); + + let temperature = dataset + .get_array("temperature") + .unwrap() .as_any() - .downcast_ref::>() - .expect("downcast string"); - assert!(season.shape().is_empty(), "attribute should be rank-0"); - assert_eq!(season.clone_into_raw_vec().await, vec!["winter".to_string()]); + .downcast_ref::>() + .unwrap(); + assert_eq!( + temperature.clone_into_raw_vec().await, + vec![1.0, 2.0, 3.0, 4.0] + ); - let year = winter - .get_array("year") - .expect("year attribute") + let time = dataset + .get_array("time") + .unwrap() .as_any() - .downcast_ref::>() - .expect("downcast i64"); - assert_eq!(year.clone_into_raw_vec().await, vec![2024i64]); + .downcast_ref::>() + .unwrap(); + assert_eq!( + time.clone_into_raw_vec().await[0], + TimestampNanosecond(test_support::EPOCH_NANOS) + ); } + /// The fill reaches the column, which is what lets the engine null an + /// unwritten cell. #[tokio::test] - async fn open_dataset_propagates_array_fill_value() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; + async fn an_arrays_fill_value_reaches_its_column() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) + .await + .unwrap(); - let winter = open_fixture_dataset(tmp.path(), "winter").await; - let cycle = winter + let cycle = dataset .get_array("cycle") - .expect("cycle array") + .unwrap() .as_any() .downcast_ref::>() - .expect("downcast i32"); - assert_eq!(cycle.fill_value().await, Some(-1i32)); + .unwrap(); + assert_eq!(cycle.fill_value().await, Some(-1)); - let temperature = winter + let temperature = dataset .get_array("temperature") - .expect("temperature array") + .unwrap() .as_any() .downcast_ref::>() - .expect("downcast f32"); - assert_eq!(temperature.fill_value().await, None); + .unwrap(); + assert_eq!(temperature.fill_value().await, None, "none was declared"); + } + + /// The chunk shape is the writer's, so a scan cuts the dataset on the grid + /// the file stores rather than on one it invents. + #[tokio::test] + async fn a_column_reports_the_stored_chunk_shape() { + let tmp = tempfile::tempdir().unwrap(); + test_support::chunked_grid(tmp.path()).await; + + let dataset = dataset_from_view(view(tmp.path(), "grid").await, None) + .await + .unwrap(); + let temperature = dataset.get_array("temperature").unwrap(); + assert_eq!(temperature.shape(), vec![4, 6]); + assert_eq!(temperature.chunk_shape(), vec![2, 3]); + assert_eq!( + temperature.dimensions(), + vec!["lat".to_string(), "lon".to_string()] + ); + } + + // ── projection ────────────────────────────────────────────────────── + + #[tokio::test] + async fn a_projection_builds_only_what_it_names() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let wanted = vec!["temperature".to_string(), ".season".to_string()]; + let dataset = dataset_from_view(view(tmp.path(), "winter").await, Some(&wanted)) + .await + .unwrap(); + + assert_eq!(names(&dataset), vec![".season", "temperature"]); } #[tokio::test] - async fn open_dataset_distinguishes_between_dataset_views() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; + async fn a_projection_may_name_a_column_this_dataset_lacks() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + // `cycle` is winter's alone; summer simply has none of it. + let wanted = vec!["temperature".to_string(), "cycle".to_string()]; + let dataset = dataset_from_view(view(tmp.path(), "summer").await, Some(&wanted)) + .await + .unwrap(); + assert_eq!(names(&dataset), vec!["temperature"]); + } - let winter = open_fixture_dataset(tmp.path(), "winter").await; - let summer = open_fixture_dataset(tmp.path(), "summer").await; + #[tokio::test] + async fn an_empty_projection_builds_nothing() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let dataset = dataset_from_view(view(tmp.path(), "winter").await, Some(&[])) + .await + .unwrap(); + assert!(names(&dataset).is_empty(), "COUNT(*) needs no column here"); + } + + // ── values Beacon cannot surface ──────────────────────────────────── + + #[tokio::test] + async fn a_list_attribute_is_dropped_and_the_rest_survives() { + let tmp = tempfile::tempdir().unwrap(); + test_support::skips(tmp.path()).await; + + let dataset = dataset_from_view(view(tmp.path(), "s").await, None) + .await + .unwrap(); + let columns = names(&dataset); + + assert!( + !columns.contains(&"value.range".to_string()), + "a list attribute has no rank-0 form: {columns:?}" + ); + assert!( + !columns.contains(&".tags".to_string()), + "nor does a list dataset attribute: {columns:?}" + ); assert_eq!( - winter.dataset().get_array("temperature").unwrap().shape(), - &[4] + columns, + vec![".title", "value", "value.units"], + "everything else is kept" ); + } + + // ── ragged datasets ───────────────────────────────────────────────── + + /// The `{array}.{attr}` naming is what the engine's ragged detection reads, + /// so a CF contiguous ragged collection is recognized without anything + /// atlas-specific. + #[tokio::test] + async fn a_plain_dataset_is_regular() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) + .await + .unwrap(); + assert!(!dataset.is_ragged()); + } + + // ── the collection schema ─────────────────────────────────────────── + + #[tokio::test] + async fn a_collections_schema_is_the_union_of_its_datasets() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let schema = collection_schema(&atlas, None, "c", &widening()) + .await + .unwrap(); + let columns: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + + for expected in [ + ".season", + ".year", + "cycle", + "temperature", + "temperature.units", + "time", + ] { + assert!( + columns.contains(&expected), + "missing {expected}: {columns:?}" + ); + } + } + + /// Two datasets that give one array two numeric types merge to the type that + /// holds both, by the rule of the session rather than one of atlas's own. + /// + /// `Int16` beside `Float32` gives `Float64`: a `Float32` holds no `Int32`, + /// so a rule that kept it for a narrow integer would make the answer depend + /// on which dataset the merge saw first. See issue #377. + #[tokio::test] + async fn a_shared_array_widens_to_a_type_that_holds_both() { + let tmp = tempfile::tempdir().unwrap(); + test_support::widening(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let schema = collection_schema(&atlas, None, "c", &widening()) + .await + .unwrap(); + assert_eq!( + schema.field_with_name("value").unwrap().data_type(), + &DataType::Float64, + "Int16 and Float32 widen to Float64" + ); + assert_eq!( + schema.field_with_name("flag").unwrap().data_type(), + &DataType::Int32, + "a column only one dataset declares keeps its own type" + ); + } + + /// A column two datasets type in two families is refused, and the error + /// names both datasets. + /// + /// Atlas reconciles nothing, so a collection can hold this. Beacon settles + /// it the way it settles two files of any other format, and the label makes + /// the offender findable in a collection of a million datasets. + #[tokio::test] + async fn types_that_do_not_widen_are_refused_by_name() { + let tmp = tempfile::tempdir().unwrap(); + test_support::incompatible(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let error = collection_schema(&atlas, None, "sensor", &widening()) + .await + .expect_err("Utf8 and Int64 are two families") + .to_string(); + + assert!(error.contains("value"), "the column: {error}"); + assert!( + error.contains("Utf8") && error.contains("Int64"), + "both types: {error}" + ); + assert!( + error.contains("sensor#a") && error.contains("sensor#b"), + "and both datasets: {error}" + ); + } + + /// A deployment that reads such a collection anyway sets `keep_first`, and + /// the column then takes the type of the first dataset in listing order. + #[tokio::test] + async fn keep_first_settles_a_conflict_with_the_first_datasets_type() { + use beacon_datafusion_ext::type_widening::{DefaultArrowTypeWidening, TypeConflict}; + + let tmp = tempfile::tempdir().unwrap(); + test_support::incompatible(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let keep_first = ArrowTypeWidening::new(Arc::new(DefaultArrowTypeWidening { + on_conflict: TypeConflict::KeepFirst, + })); + let schema = collection_schema(&atlas, None, "c", &keep_first) + .await + .unwrap(); + + // `a` is written first and states `value` as a string. assert_eq!( - summer.dataset().get_array("temperature").unwrap().shape(), - &[3] + schema.field_with_name("value").unwrap().data_type(), + &DataType::Utf8 ); - assert!(summer.dataset().get_array("cycle").is_none()); - assert!(summer.dataset().get_array("year").is_none()); } #[tokio::test] - async fn open_dataset_unknown_returns_error() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; + async fn a_collection_with_no_dataset_has_no_column() { + let tmp = tempfile::tempdir().unwrap(); + test_support::empty(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let schema = collection_schema(&atlas, None, "c", &widening()) + .await + .unwrap(); + assert!(schema.fields().is_empty()); + } + + /// Datasets that declare the same arrays share one interned schema, so the + /// derivation runs once however many of them there are. The fleet fixture + /// gives ten datasets one shape, and they differ only in an attribute + /// value. + #[tokio::test] + async fn a_fleet_of_one_shape_derives_one_schema() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 10).await; + let atlas = test_support::open(tmp.path()).await; - let store: Arc = - Arc::new(LocalFileSystem::new_with_prefix(tmp.path()).unwrap()); - let err = open_dataset(store, OsPath::from(""), "ghost") + assert_eq!(atlas.interned_schemas(), 1, "the fixture shares its schema"); + + let mut keys = HashSet::new(); + for name in atlas.list_datasets() { + keys.insert(shape_key(&atlas.dataset(&name).unwrap())); + } + assert_eq!(keys.len(), 1, "and every dataset reduces to one key"); + + let schema = collection_schema(&atlas, None, "c", &widening()) + .await + .unwrap(); + let columns: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + assert_eq!(columns, vec![".platform", "temperature"]); + } + + /// Two datasets that share arrays but not attribute *keys* produce two + /// column sets, so the key has to separate them. + #[tokio::test] + async fn a_different_attribute_namespace_is_a_different_shape() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let winter = shape_key(&atlas.dataset("winter").unwrap()); + let summer = shape_key(&atlas.dataset("summer").unwrap()); + assert_ne!(winter, summer); + } + + // ── dimensions ────────────────────────────────────────────────────── + + #[tokio::test] + async fn read_dimensions_narrow_the_schema_to_the_grid_they_name() { + let tmp = tempfile::tempdir().unwrap(); + test_support::chunked_grid(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + // Both arrays live on `lat` and `lon`; naming only `lat` leaves neither. + let dims = ["lat".to_string()]; + let schema = collection_schema(&atlas, Some(&dims), "c", &widening()) .await - .expect_err("should fail for unknown dataset"); - let msg = format!("{err:#}"); + .unwrap(); + let columns: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); assert!( - msg.contains("ghost") || msg.contains("DatasetNotFound"), - "error should mention missing dataset name: {msg}" + !columns.contains(&"temperature"), + "a 2-D array does not fit a 1-D grid: {columns:?}" ); } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/store.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/store.rs new file mode 100644 index 00000000..307800ee --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/store.rs @@ -0,0 +1,405 @@ +//! Finding an Atlas collection in a listing, opening it, and caching the +//! handle. +//! +//! A collection is a store prefix holding one required object, `data.atlas`, +//! and one optional sidecar, `deleted.mask`. The container object is the +//! *marker*: it is what a listing matches, what a plan entry carries, and what +//! the reader cache keys on. Its parent directory is the prefix +//! [`atlas::Atlas::open`] takes. + +use std::sync::Arc; + +use atlas::Atlas; +use chrono::{DateTime, Utc}; +use moka::future::Cache; +use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt, path::Path as OsPath}; + +/// The container object at the root of a collection. +pub const ATLAS_MARKER: &str = "data.atlas"; + +/// The deletion-mask sidecar beside it. Absent means nothing is deleted. +pub const ATLAS_MASK: &str = "deleted.mask"; + +/// `data.atlas` as it appears at the end of a nested path. +const MARKER_SUFFIX: &str = "/data.atlas"; + +/// Whether `path` names a collection's container object. +/// +/// The name is fixed. [`atlas::Atlas::open`] resolves `data.atlas` under the +/// prefix it is given, so a collection renamed to `sensor.atlas` cannot be +/// opened at all and is not a marker. +pub fn is_marker_path(path: &OsPath) -> bool { + let path = path.as_ref(); + path == ATLAS_MARKER || path.ends_with(MARKER_SUFFIX) +} + +/// Whether `obj` is a collection's container object. +pub fn is_atlas_marker(obj: &ObjectMeta) -> bool { + is_marker_path(&obj.location) +} + +/// The prefix a collection is opened under: the marker's parent directory. +/// +/// An empty path means the marker sits at the store root, which is what a +/// store rooted on the collection's own directory reports. +pub fn collection_prefix(marker: &OsPath) -> Option { + let path = marker.as_ref(); + if path == ATLAS_MARKER { + return Some(OsPath::default()); + } + path.strip_suffix(MARKER_SUFFIX).map(OsPath::from) +} + +/// The directory of a marker, as a string. `""` for one at the root. +fn marker_directory(marker: &OsPath) -> Option { + let path = marker.as_ref(); + if path == ATLAS_MARKER { + return Some(String::new()); + } + path.strip_suffix(MARKER_SUFFIX).map(str::to_string) +} + +/// Reduce `objects` to the unique outermost collection markers. +/// +/// Two markers at two depths of one tree keep only the ancestor: a collection +/// is one file and never contains another, so a deeper marker is a collection +/// that happens to sit inside another's directory and would be read twice. +pub fn top_level_atlas_markers(objects: &[ObjectMeta]) -> Vec { + // By directory, not by path. A path sort would put "a/b/data.atlas" before + // "a/data.atlas", because 'b' sorts under 'd', and the nested collection + // would then be the one kept. + let mut markers: Vec<(String, &ObjectMeta)> = objects + .iter() + .filter_map(|object| { + marker_directory(&object.location).map(|directory| (directory, object)) + }) + .collect(); + markers.sort_by(|(a, _), (b, _)| a.cmp(b)); + + let mut kept: Vec<(String, ObjectMeta)> = Vec::new(); + 'outer: for (directory, marker) in markers { + for (held, _) in &kept { + // A marker at the root sits above every path, but the collection + // beside it is its own, so an empty directory excludes nothing. + if !held.is_empty() && directory.starts_with(&format!("{held}/")) { + continue 'outer; + } + } + kept.push((directory, marker.clone())); + } + kept.into_iter().map(|(_, marker)| marker).collect() +} + +/// Open the collection whose container object is `marker`, over `store`. +/// +/// One `HEAD`, one tail read, and one `GET` of the deletion mask when it +/// exists. Nothing else, whatever the collection holds. +pub async fn open_collection( + store: Arc, + marker: &OsPath, +) -> anyhow::Result> { + let prefix = collection_prefix(marker).ok_or_else(|| { + anyhow::anyhow!( + "'{marker}' is not an atlas collection: the container is named '{ATLAS_MARKER}'" + ) + })?; + let atlas = Atlas::open(store, prefix.clone()) + .await + .map_err(|e| anyhow::anyhow!("Failed to open the atlas collection at '{prefix}': {e}"))?; + Ok(Arc::new(atlas)) +} + +/// What a cached handle describes, beyond the container itself. +/// +/// The container never changes after a write, so its size and modification time +/// pin its contents completely. The mask is the one part of a finished +/// collection that can change, and it decides which datasets a handle reports, +/// so it belongs in the key. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct MaskStamp { + last_modified: DateTime, + size: u64, + e_tag: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct CacheKey { + path: OsPath, + last_modified: DateTime, + size: u64, + /// `None` when the collection has no mask, which is the common case. + mask: Option, +} + +/// A cache of opened collections, sized at construction. +/// +/// Cloning shares the underlying [`moka`] cache, so the formats, sources and +/// openers a runtime hands a clone to all draw from one store. This is +/// per-runtime state; there is no process-global cache. +/// +/// Each entry owns a 256 MiB block cache and a 64 MiB I/O cache of its own, so +/// the capacity is a memory bound as much as a handle count. See +/// [`AtlasConfig::reader_cache_size`](crate::AtlasConfig::reader_cache_size). +#[derive(Clone)] +pub struct AtlasReaderCache { + cache: Cache>, +} + +impl AtlasReaderCache { + /// Build a cache holding up to `capacity` opened collections. + pub fn new(capacity: u64) -> Self { + Self { + cache: Cache::builder().max_capacity(capacity).build(), + } + } +} + +// `Atlas` is not `Debug`, and this sits inside `Debug` formats and sources. +impl std::fmt::Debug for AtlasReaderCache { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AtlasReaderCache").finish_non_exhaustive() + } +} + +/// The identity of a collection's deletion mask, or `None` when it has none. +/// +/// One `HEAD`. An error other than "not found" also reads as `None`: the mask +/// only ever *hides* datasets, so the worst a stale handle can do is report a +/// dataset a concurrent delete just hid, and the alternative is failing a query +/// over a transient head request. +async fn mask_stamp(store: &dyn ObjectStore, prefix: &OsPath) -> Option { + let path = prefix.clone().join(ATLAS_MASK); + match store.head(&path).await { + Ok(meta) => Some(MaskStamp { + last_modified: meta.last_modified, + size: meta.size, + e_tag: meta.e_tag, + }), + Err(object_store::Error::NotFound { .. }) => None, + Err(e) => { + tracing::debug!(path = %path, error = %e, "could not stat the atlas deletion mask"); + None + } + } +} + +/// A cached handle for `marker`, opening it from `store` on a miss. +/// +/// With `cache` set to `None` the collection is opened directly, with no +/// caching. Otherwise the key carries the marker's identity and the mask's, so +/// a rewritten collection or a fresh delete produces a new key and a re-open. +/// Concurrent first readers of one key coalesce inside +/// [`moka::future::Cache::try_get_with`]. +pub async fn get_or_open_atlas( + cache: Option<&AtlasReaderCache>, + store: Arc, + marker: &ObjectMeta, +) -> anyhow::Result> { + let Some(cache) = cache else { + return open_collection(store, &marker.location).await; + }; + + let prefix = collection_prefix(&marker.location).ok_or_else(|| { + anyhow::anyhow!( + "'{}' is not an atlas collection: the container is named '{ATLAS_MARKER}'", + marker.location + ) + })?; + let key = CacheKey { + path: marker.location.clone(), + last_modified: marker.last_modified, + size: marker.size, + mask: mask_stamp(store.as_ref(), &prefix).await, + }; + + let path = marker.location.clone(); + cache + .cache + .try_get_with(key, async move { open_collection(store, &path).await }) + .await + .map_err(|e: Arc| anyhow::anyhow!("{e}")) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support; + + fn object(path: &str) -> ObjectMeta { + ObjectMeta { + location: OsPath::from(path), + last_modified: DateTime::UNIX_EPOCH, + size: 0, + e_tag: None, + version: None, + } + } + + // ── markers ───────────────────────────────────────────────────────── + + #[test] + fn the_container_object_is_the_marker() { + assert!(is_atlas_marker(&object("data.atlas"))); + assert!(is_atlas_marker(&object("store/data.atlas"))); + assert!(is_atlas_marker(&object("a/b/c/data.atlas"))); + } + + #[test] + fn nothing_else_is_a_marker() { + // The mask sits beside the container and must never be read as one. + // Neither must the registry of a pre-0.16 collection: this build reads + // only the single-file format, so an old collection left on disk is + // passed over rather than misread. + for path in [ + "deleted.mask", + "store/deleted.mask", + "store/data.atlas.tmp", + "store/mydata.atlas", + "data.atlas/inner", + "store/atlas.json", + ] { + assert!(!is_atlas_marker(&object(path)), "{path}"); + } + } + + #[test] + fn the_prefix_is_the_marker_directory() { + assert_eq!( + collection_prefix(&OsPath::from("a/b/data.atlas")), + Some(OsPath::from("a/b")) + ); + assert_eq!( + collection_prefix(&OsPath::from("data.atlas")), + Some(OsPath::default()) + ); + assert_eq!(collection_prefix(&OsPath::from("a/b/other.txt")), None); + } + + #[test] + fn top_level_markers_drop_a_nested_collection() { + let objects = vec![ + object("a/data.atlas"), + object("a/b/data.atlas"), + object("c/data.atlas"), + object("c/deleted.mask"), + ]; + let kept: Vec = top_level_atlas_markers(&objects) + .iter() + .map(|m| m.location.to_string()) + .collect(); + assert_eq!(kept, vec!["a/data.atlas", "c/data.atlas"]); + } + + #[test] + fn a_sibling_directory_is_not_nested() { + // "argo2" starts with "argo", but it is not under it. + let objects = vec![object("argo/data.atlas"), object("argo2/data.atlas")]; + assert_eq!(top_level_atlas_markers(&objects).len(), 2); + } + + // ── opening ───────────────────────────────────────────────────────── + + #[tokio::test] + async fn a_collection_opens_from_its_marker() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let (store, marker) = test_support::store_and_marker(tmp.path()); + + let atlas = open_collection(store, &marker.location).await.unwrap(); + assert_eq!( + atlas.list_datasets(), + vec!["winter", "summer"], + "a collection lists in write order" + ); + } + + /// Anything but the container is refused, and the error names what a + /// collection is called. + #[tokio::test] + async fn a_path_that_is_not_the_container_is_refused() { + let tmp = tempfile::tempdir().unwrap(); + let (store, _) = test_support::store_and_marker(tmp.path()); + let error = open_collection(store, &OsPath::from("store/index.json")) + .await + .expect_err("only the container names a collection") + .to_string(); + assert!(error.contains("data.atlas"), "{error}"); + } + + // ── the reader cache ──────────────────────────────────────────────── + + #[tokio::test] + async fn one_marker_opens_once() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let (store, marker) = test_support::store_and_marker(tmp.path()); + let cache = AtlasReaderCache::new(8); + + let first = get_or_open_atlas(Some(&cache), store.clone(), &marker) + .await + .unwrap(); + let second = get_or_open_atlas(Some(&cache), store, &marker) + .await + .unwrap(); + assert!(Arc::ptr_eq(&first, &second), "the second open must hit"); + } + + #[tokio::test] + async fn a_rewritten_container_reopens() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let (store, marker) = test_support::store_and_marker(tmp.path()); + let cache = AtlasReaderCache::new(8); + + let first = get_or_open_atlas(Some(&cache), store.clone(), &marker) + .await + .unwrap(); + + let mut moved = marker.clone(); + moved.last_modified = marker.last_modified + chrono::Duration::seconds(1); + let second = get_or_open_atlas(Some(&cache), store, &moved) + .await + .unwrap(); + assert!(!Arc::ptr_eq(&first, &second), "a new mtime must miss"); + } + + /// A delete writes the mask and leaves the container alone, so the marker + /// says nothing about it. Without the mask in the key, a handle opened + /// before the delete keeps reporting the dataset it hid. + #[tokio::test] + async fn a_delete_reopens_the_collection() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let (store, marker) = test_support::store_and_marker(tmp.path()); + let cache = AtlasReaderCache::new(8); + + let before = get_or_open_atlas(Some(&cache), store.clone(), &marker) + .await + .unwrap(); + assert_eq!(before.list_datasets().len(), 2); + + before.delete_dataset("winter").await.unwrap(); + + let after = get_or_open_atlas(Some(&cache), store, &marker) + .await + .unwrap(); + assert!( + !Arc::ptr_eq(&before, &after), + "the mask changed, so the key did" + ); + assert_eq!(after.list_datasets(), vec!["summer"]); + } + + #[tokio::test] + async fn without_a_cache_every_open_is_its_own() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let (store, marker) = test_support::store_and_marker(tmp.path()); + + let first = get_or_open_atlas(None, store.clone(), &marker) + .await + .unwrap(); + let second = get_or_open_atlas(None, store, &marker).await.unwrap(); + assert!(!Arc::ptr_eq(&first, &second)); + } +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs new file mode 100644 index 00000000..19d81c4d --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs @@ -0,0 +1,353 @@ +//! Collections the tests of this crate read. +//! +//! Every fixture is written with the real [`AtlasWriter`], so what the tests +//! read is a real container: a footer, its segments, and the statistics the +//! writer recorded while it staged them. + +use std::path::Path; +use std::sync::Arc; + +use atlas::{Atlas, AtlasWriter, Attr, FillValue, TimestampNs, WriterConfig}; +use chrono::{DateTime, Utc}; +use ndarray::{ArrayD, IxDyn, arr1}; +use object_store::{ObjectMeta, ObjectStore, local::LocalFileSystem, path::Path as OsPath}; + +/// 2024-01-01T00:00:00Z, the epoch the `time` arrays count from. +pub const EPOCH_NANOS: i64 = 1_704_067_200_000_000_000; + +/// One day, in nanoseconds. +pub const DAY_NANOS: i64 = 86_400_000_000_000; + +/// A store rooted on `dir`, and the marker of the collection in it. +/// +/// The marker carries the container's real size and modification time, which +/// is what the reader cache keys on. +pub fn store_and_marker(dir: &Path) -> (Arc, ObjectMeta) { + let store: Arc = Arc::new(LocalFileSystem::new_with_prefix(dir).unwrap()); + let container = dir.join(crate::store::ATLAS_MARKER); + let (size, last_modified) = match std::fs::metadata(&container) { + Ok(meta) => ( + meta.len(), + meta.modified() + .map(DateTime::::from) + .unwrap_or(DateTime::UNIX_EPOCH), + ), + // A test that has not written a collection still needs a marker to + // hand to the code that will refuse it. + Err(_) => (0, DateTime::UNIX_EPOCH), + }; + let marker = ObjectMeta { + location: OsPath::from(crate::store::ATLAS_MARKER), + last_modified, + size, + e_tag: None, + version: None, + }; + (store, marker) +} + +/// Open a fixture collection from its directory. +pub async fn open(dir: &Path) -> Arc { + Arc::new(Atlas::open_path(dir).await.expect("open the collection")) +} + +/// Two datasets that do not share a schema. +/// +/// - `winter`: `temperature: Float32[4]`, `cycle: Int32[4]` with a fill of +/// `-1`, and `time: TimestampNs[4]`. Dataset attributes `season` and `year`; +/// `temperature` carries `units`. +/// - `summer`: `temperature: Float32[3]` alone, and the attribute `season`. +/// +/// The differing lengths make the two datasets separable in a result, and the +/// fill on `cycle` carries an unwritten cell through to a null. +pub async fn two_datasets(dir: &Path) { + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + + { + let mut winter = writer.add_dataset("winter").await.expect("add winter"); + winter + .define_array::("temperature", vec!["obs".into()], vec![4], None, None) + .await + .expect("define temperature"); + winter + .define_array::( + "cycle", + vec!["obs".into()], + vec![4], + None, + Some(FillValue::Int(-1)), + ) + .await + .expect("define cycle"); + winter + .define_array::("time", vec!["obs".into()], vec![4], None, None) + .await + .expect("define time"); + + winter + .write_array( + "temperature", + vec![0], + arr1(&[1.0f32, 2.0, 3.0, 4.0]).into_dyn().view(), + ) + .await + .expect("write temperature"); + winter + .write_array( + "cycle", + vec![0], + arr1(&[10i32, 20, 30, 40]).into_dyn().view(), + ) + .await + .expect("write cycle"); + let times: Vec = (0..4) + .map(|day| TimestampNs(EPOCH_NANOS + day * DAY_NANOS)) + .collect(); + winter + .write_array("time", vec![0], arr1(×).into_dyn().view()) + .await + .expect("write time"); + + winter.set_attribute("season", Attr::String("winter".into())); + winter.set_attribute("year", Attr::Int64(2024)); + winter + .set_array_attribute("temperature", "units", Attr::String("celsius".into())) + .expect("set units"); + winter.finish().await.expect("finish winter"); + } + + { + let mut summer = writer.add_dataset("summer").await.expect("add summer"); + summer + .define_array::("temperature", vec!["obs".into()], vec![3], None, None) + .await + .expect("define temperature"); + summer + .write_array( + "temperature", + vec![0], + arr1(&[20.0f32, 21.0, 22.0]).into_dyn().view(), + ) + .await + .expect("write temperature"); + summer.set_attribute("season", Attr::String("summer".into())); + summer.finish().await.expect("finish summer"); + } + + writer.finish().await.expect("finish the collection"); +} + +/// One dataset, `grid`, holding two chunked 2-D arrays on `lat` and `lon`. +/// +/// - `temperature: Float64[4, 6]`, chunked `[2, 3]`, written whole. Cell +/// `(row, col)` holds `row * 6 + col`, so a window states its own position. +/// - `sparse: Float64[4, 6]`, the same shape and chunking, with a fill of +/// `-999`. Only the first two rows are written, so the rest is a hole that +/// costs no bytes. +pub async fn chunked_grid(dir: &Path) { + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + + let mut grid = writer.add_dataset("grid").await.expect("add grid"); + let dims = vec!["lat".to_string(), "lon".to_string()]; + grid.define_array::( + "temperature", + dims.clone(), + vec![4, 6], + Some(vec![2, 3]), + None, + ) + .await + .expect("define temperature"); + grid.define_array::( + "sparse", + dims, + vec![4, 6], + Some(vec![2, 3]), + Some(FillValue::Float(-999.0)), + ) + .await + .expect("define sparse"); + + let values = ArrayD::from_shape_fn(IxDyn(&[4, 6]), |i| (i[0] * 6 + i[1]) as f64); + grid.write_array("temperature", vec![0, 0], values.view()) + .await + .expect("write temperature"); + + let written = ArrayD::from_shape_fn(IxDyn(&[2, 6]), |i| (i[0] * 6 + i[1]) as f64); + grid.write_array("sparse", vec![0, 0], written.view()) + .await + .expect("write sparse"); + + grid.finish().await.expect("finish grid"); + writer.finish().await.expect("finish the collection"); +} + +/// Two datasets whose shared array has two numeric types. +/// +/// - `a`: `value: Int16[2] = [1, 2]`, `flag: Int32[2] = [7, 8]`. +/// - `b`: `value: Float32[2] = [3.5, 4.5]`. +/// +/// The merged `value` widens past either dataset's own type, and `flag` is a +/// column only one dataset declares. +pub async fn widening(dir: &Path) { + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + + { + let mut a = writer.add_dataset("a").await.expect("add a"); + a.define_array::("value", vec!["obs".into()], vec![2], None, None) + .await + .expect("define value"); + a.define_array::("flag", vec!["obs".into()], vec![2], None, None) + .await + .expect("define flag"); + a.write_array("value", vec![0], arr1(&[1i16, 2]).into_dyn().view()) + .await + .expect("write value"); + a.write_array("flag", vec![0], arr1(&[7i32, 8]).into_dyn().view()) + .await + .expect("write flag"); + a.finish().await.expect("finish a"); + } + + { + let mut b = writer.add_dataset("b").await.expect("add b"); + b.define_array::("value", vec!["obs".into()], vec![2], None, None) + .await + .expect("define value"); + b.write_array("value", vec![0], arr1(&[3.5f32, 4.5]).into_dyn().view()) + .await + .expect("write value"); + b.finish().await.expect("finish b"); + } + + writer.finish().await.expect("finish the collection"); +} + +/// Two datasets whose shared array has no common numeric type. +/// +/// - `a`: `value: String[2] = ["x", "y"]`, `only_a: Int32[2] = [7, 8]`. +/// - `b`: `value: Int64[2] = [1, 2]`. +/// +/// There is no numeric super-type here, so this pins what the merge resolves +/// to and whether both datasets stay readable. +pub async fn incompatible(dir: &Path) { + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + + { + let mut a = writer.add_dataset("a").await.expect("add a"); + a.define_array::("value", vec!["obs".into()], vec![2], None, None) + .await + .expect("define value"); + a.define_array::("only_a", vec!["obs".into()], vec![2], None, None) + .await + .expect("define only_a"); + a.write_array( + "value", + vec![0], + arr1(&["x".to_string(), "y".to_string()]).into_dyn().view(), + ) + .await + .expect("write value"); + a.write_array("only_a", vec![0], arr1(&[7i32, 8]).into_dyn().view()) + .await + .expect("write only_a"); + a.finish().await.expect("finish a"); + } + + { + let mut b = writer.add_dataset("b").await.expect("add b"); + b.define_array::("value", vec!["obs".into()], vec![2], None, None) + .await + .expect("define value"); + b.write_array("value", vec![0], arr1(&[1i64, 2]).into_dyn().view()) + .await + .expect("write value"); + b.finish().await.expect("finish b"); + } + + writer.finish().await.expect("finish the collection"); +} + +/// `n` datasets named `d0..d{n-1}`, each holding `temperature: Float32[4]` +/// over the disjoint range `[10i, 10i + 3]`. +/// +/// A predicate such as `temperature > 45` then has a known answer: the +/// datasets whose range reaches past it, and no others. +pub async fn ranged(dir: &Path, n: usize) { + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + + for i in 0..n { + let mut ds = writer + .add_dataset(&format!("d{i}")) + .await + .expect("add a dataset"); + ds.define_array::("temperature", vec!["obs".into()], vec![4], None, None) + .await + .expect("define temperature"); + let base = (10 * i) as f32; + ds.write_array( + "temperature", + vec![0], + arr1(&[base, base + 1.0, base + 2.0, base + 3.0]) + .into_dyn() + .view(), + ) + .await + .expect("write temperature"); + ds.set_attribute("platform", Attr::String(format!("p{i}"))); + ds.finish().await.expect("finish a dataset"); + } + + writer.finish().await.expect("finish the collection"); +} + +/// One dataset carrying values Beacon cannot surface as columns. +/// +/// `value: Float64[2]` is readable. The list attribute `range` has no rank-0 +/// form and is dropped; the string attribute `units` beside it is kept, so a +/// test can tell "dropped" from "dropped everything". +/// +/// A `Bool` or list *array* cannot appear here: `array-format` implements no +/// element type for either, so no writer can produce one. +pub async fn skips(dir: &Path) { + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + + let mut ds = writer.add_dataset("s").await.expect("add s"); + ds.define_array::("value", vec!["obs".into()], vec![2], None, None) + .await + .expect("define value"); + ds.write_array("value", vec![0], arr1(&[1.0f64, 2.0]).into_dyn().view()) + .await + .expect("write value"); + ds.set_array_attribute("value", "units", Attr::String("metres".into())) + .expect("set units"); + ds.set_array_attribute("value", "range", Attr::Float64List(vec![0.0, 10.0])) + .expect("set range"); + ds.set_attribute("tags", Attr::StringList(vec!["a".into(), "b".into()])); + ds.set_attribute("title", Attr::String("skips".into())); + ds.finish().await.expect("finish s"); + + writer.finish().await.expect("finish the collection"); +} + +/// A collection that holds no dataset at all. +/// +/// Legal, and the writer produces one whenever a job finds nothing to ingest. +pub async fn empty(dir: &Path) { + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + writer.finish().await.expect("finish the collection"); +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/util.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/util.rs deleted file mode 100644 index 92a28283..00000000 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/util.rs +++ /dev/null @@ -1,164 +0,0 @@ -//! Atlas metadata-marker discovery helpers shared by the DataFusion -//! integration. -//! -//! An atlas store is a directory holding a single metadata marker at its root -//! — one of [`ATLAS_MARKER_NAMES`] — plus the per-array `.af` files. These -//! helpers recognize markers in an object listing and reduce a listing to the -//! outermost stores, mirroring `beacon-arrow-zarr`'s `zarr.json` discovery. - -use object_store::{ObjectMeta, path::Path as OsPath}; - -/// All marker filenames atlas may emit at the root of a store, in the same -/// priority order atlas's own `META_VARIANTS` uses (uncompressed before -/// compressed within a format; JSON before MsgPack overall). Atlas detects the -/// actual variant on `open`, so beacon only needs to recognize the names. -pub const ATLAS_MARKER_NAMES: [&str; 6] = [ - "atlas.json", - "atlas.json.zst", - "atlas.json.lz4", - "atlas.msgpack", - "atlas.msgpack.zst", - "atlas.msgpack.lz4", -]; - -/// Canonical marker filename, used as the format's listing extension and the -/// `get_ext` identity. Atlas defaults to uncompressed JSON, so this is the -/// common on-disk name. -pub const ATLAS_MARKER: &str = "atlas.json"; - -/// If `p` ends in one of the known atlas marker filenames, return that -/// filename. Used both to recognize markers and to recover the on-disk name for -/// path manipulation. -pub fn atlas_marker_filename(p: &OsPath) -> Option<&'static str> { - let s = p.as_ref(); - ATLAS_MARKER_NAMES - .iter() - .copied() - .find(|name| s == *name || s.ends_with(&format!("/{name}"))) -} - -/// `true` if `obj` is an atlas metadata marker. -pub fn is_atlas_marker(obj: &ObjectMeta) -> bool { - atlas_marker_filename(&obj.location).is_some() -} - -/// The store directory (marker's parent) as a string, or `None` if `p` is not a -/// marker. An empty string means the marker sits at the object-store root. -pub fn marker_parent(p: &OsPath) -> Option { - let s = p.as_ref(); - let name = atlas_marker_filename(p)?; - Some(s.strip_suffix(name)?.trim_end_matches('/').to_string()) -} - -/// The object-store prefix an atlas store is opened under: the directory -/// containing `marker` (its parent). Passed straight to -/// [`atlas::Atlas::open`](atlas::Atlas::open). -pub fn atlas_store_prefix(marker: &OsPath) -> Option { - marker_parent(marker).map(|dir| OsPath::from(dir.as_str())) -} - -/// Filter `objects` down to the unique top-level atlas markers. -/// -/// If two markers appear at different depths of a nested tree we keep only the -/// outermost (the ancestor) — atlas stores never contain other atlas stores, so -/// any deeper marker is spurious. Mirrors zarr's `top_level_zarr_meta_v3`. -pub fn top_level_atlas_markers(objects: &[ObjectMeta]) -> Vec { - let mut markers: Vec<&ObjectMeta> = objects.iter().filter(|o| is_atlas_marker(o)).collect(); - markers.sort_by(|a, b| a.location.as_ref().cmp(b.location.as_ref())); - - let mut kept: Vec = Vec::new(); - 'outer: for meta in &markers { - let dir = marker_parent(&meta.location).unwrap_or_default(); - for already in &kept { - let already_dir = marker_parent(&already.location).unwrap_or_default(); - if !already_dir.is_empty() && dir.starts_with(&format!("{already_dir}/")) { - continue 'outer; - } - } - kept.push((*meta).clone()); - } - kept -} - -#[cfg(test)] -mod tests { - use super::*; - - fn marker_obj(path: &str) -> ObjectMeta { - ObjectMeta { - location: OsPath::from(path), - last_modified: Default::default(), - size: 0, - e_tag: None, - version: None, - } - } - - #[test] - fn is_atlas_marker_matches_all_variants() { - for name in &ATLAS_MARKER_NAMES { - assert!(is_atlas_marker(&marker_obj(name)), "bare {name}"); - assert!( - is_atlas_marker(&marker_obj(&format!("store/{name}"))), - "store/{name}" - ); - assert!( - is_atlas_marker(&marker_obj(&format!("a/b/c/{name}"))), - "a/b/c/{name}" - ); - } - for negative in [ - "foo/data.af", - "store/atlas.json.tmp", - "store/atlas.jsona", - "atlas.json/inner", - ] { - assert!(!is_atlas_marker(&marker_obj(negative)), "{negative}"); - } - } - - #[test] - fn marker_parent_strips_each_variant() { - for name in &ATLAS_MARKER_NAMES { - assert_eq!( - marker_parent(&OsPath::from(format!("store/{name}"))), - Some("store".to_string()), - "{name}" - ); - assert_eq!( - marker_parent(&OsPath::from(*name)), - Some(String::new()), - "bare {name}" - ); - } - assert_eq!(marker_parent(&OsPath::from("foo.txt")), None); - } - - #[test] - fn store_prefix_is_marker_parent() { - assert_eq!( - atlas_store_prefix(&OsPath::from("a/b/atlas.json")), - Some(OsPath::from("a/b")) - ); - assert_eq!( - atlas_store_prefix(&OsPath::from("atlas.json")), - Some(OsPath::from("")) - ); - } - - #[test] - fn top_level_markers_drops_nested_stores() { - let objs = vec![ - marker_obj("a/atlas.json"), - marker_obj("a/b/atlas.json"), - marker_obj("c/atlas.msgpack"), - ]; - let kept: Vec = top_level_atlas_markers(&objs) - .iter() - .map(|m| m.location.to_string()) - .collect(); - assert!(kept.contains(&"a/atlas.json".to_string())); - assert!(kept.contains(&"c/atlas.msgpack".to_string())); - assert!(!kept.iter().any(|p| p == "a/b/atlas.json")); - } -} diff --git a/beacon-db/beacon-file-formats/beacon-nd-array/src/arrow/file_read.rs b/beacon-db/beacon-file-formats/beacon-nd-array/src/arrow/file_read.rs index 25b79a73..b9259620 100644 --- a/beacon-db/beacon-file-formats/beacon-nd-array/src/arrow/file_read.rs +++ b/beacon-db/beacon-file-formats/beacon-nd-array/src/arrow/file_read.rs @@ -359,6 +359,24 @@ pub struct FileRead { } impl FileRead { + /// A file the scan decided not to read at all. + /// + /// Nothing is queued and nothing is streamed, so the file costs one pop and + /// no I/O. This is what a format returns for a file it ruled out *before* + /// opening it — an Atlas dataset whose footer statistics cannot satisfy the + /// predicate, say. That decision belongs to the format, because only the + /// format knows what it can prove from its own metadata. + /// + /// This is not the same as a file that holds none of the projected columns. + /// [`plan`](Self::plan) reaches that state on its own, having opened the + /// file to find out. + pub fn skipped() -> Arc { + Arc::new(Self { + queue: None, + output: Output::Nothing, + }) + } + /// Plan `dataset` for a scan that wants `projected_schema`. /// /// Resolves the projection, fills the queue, and decides what a batch off it @@ -1213,6 +1231,25 @@ mod tests { assert!(batches.is_empty(), "it contributes no rows"); } + /// A file the scan ruled out before opening reads as nothing. + /// + /// The format decides this from its own metadata — Atlas from the + /// statistics in a collection footer — so nothing here can check the + /// decision. What this pins is the shape of the answer: no work queued, no + /// batch emitted, and a clean stream rather than an error. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn a_skipped_file_queues_nothing_and_emits_nothing() { + let skipped = FileRead::skipped(); + + assert_eq!(skipped.remaining(), 0, "nothing is queued to read"); + let batches: Vec = skipped + .stream(None) + .try_collect() + .await + .expect("the stream is clean, not an error"); + assert!(batches.is_empty(), "it contributes no rows"); + } + /// A `COUNT(*)` still counts. It projects no column *because it wants none*, /// which is the case the check above has to keep telling apart. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/beacon-db/beacon-functions/Cargo.toml b/beacon-db/beacon-functions/Cargo.toml index b972c904..67038691 100644 --- a/beacon-db/beacon-functions/Cargo.toml +++ b/beacon-db/beacon-functions/Cargo.toml @@ -46,6 +46,7 @@ beacon-delta = { path = "../beacon-file-formats/beacon-delta" } beacon-iceberg = { path = "../beacon-file-formats/beacon-iceberg" } beacon-icechunk = { path = "../beacon-file-formats/beacon-icechunk" } beacon-arrow-zarr = { path = "../beacon-file-formats/beacon-arrow-zarr" } +beacon-arrow-atlas = { path = "../beacon-file-formats/beacon-arrow-atlas" } beacon-arrow-ipc = { path = "../beacon-file-formats/beacon-arrow-ipc" } beacon-arrow-csv = { path = "../beacon-file-formats/beacon-arrow-csv" } beacon-arrow-parquet = { path = "../beacon-file-formats/beacon-arrow-parquet" } diff --git a/beacon-db/beacon-functions/src/file_formats/mod.rs b/beacon-db/beacon-functions/src/file_formats/mod.rs index fc4dbd22..020abade 100644 --- a/beacon-db/beacon-functions/src/file_formats/mod.rs +++ b/beacon-db/beacon-functions/src/file_formats/mod.rs @@ -47,6 +47,10 @@ pub fn register_table_functions( runtime_handle.clone(), Arc::downgrade(&session_ctx), )), + Arc::new(beacon_arrow_atlas::ReadAtlasFunc::new( + runtime_handle.clone(), + Arc::downgrade(&session_ctx), + )), Arc::new(beacon_arrow_netcdf::datafusion::ReadNetCDFFunc::new( runtime_handle.clone(), Arc::downgrade(&session_ctx), diff --git a/beacon-server/beacon-server-config/Cargo.toml b/beacon-server/beacon-server-config/Cargo.toml index c2f30043..a8f122f7 100644 --- a/beacon-server/beacon-server-config/Cargo.toml +++ b/beacon-server/beacon-server-config/Cargo.toml @@ -21,4 +21,5 @@ beacon-datafusion-ext = { path = "../../beacon-db/beacon-datafusion-ext" } beacon-arrow-netcdf = { path = "../../beacon-db/beacon-file-formats/beacon-arrow-netcdf" } beacon-arrow-hdf5 = { path = "../../beacon-db/beacon-file-formats/beacon-arrow-hdf5" } beacon-arrow-zarr = { path = "../../beacon-db/beacon-file-formats/beacon-arrow-zarr" } +beacon-arrow-atlas = { path = "../../beacon-db/beacon-file-formats/beacon-arrow-atlas" } beacon-arrow-bbf = { path = "../../beacon-db/beacon-file-formats/beacon-arrow-bbf" } diff --git a/beacon-server/beacon-server-config/src/lib.rs b/beacon-server/beacon-server-config/src/lib.rs index 2f8f248f..8e828974 100644 --- a/beacon-server/beacon-server-config/src/lib.rs +++ b/beacon-server/beacon-server-config/src/lib.rs @@ -10,6 +10,7 @@ use error::Result; // Per-format and storage config types are owned by their crates; beacon-config // composes them here and fills them from the environment. +pub use beacon_arrow_atlas::AtlasConfig; pub use beacon_arrow_bbf::datafusion::BbfConfig; pub use beacon_arrow_hdf5::{Hdf5Config, Hdf5Convention}; pub use beacon_arrow_netcdf::datafusion::NetcdfConfig; @@ -31,6 +32,7 @@ pub struct Config { pub netcdf: NetcdfConfig, pub hdf5: Hdf5Config, pub zarr: ZarrConfig, + pub atlas: AtlasConfig, pub bbf: BbfConfig, pub crawler: CrawlerConfig, pub file_stats: FileStatsConfig, @@ -493,6 +495,31 @@ struct RawConfig { #[envconfig(from = "BEACON_ZARR_ENABLE_STATISTICS", default = "true")] zarr_enable_statistics: bool, + /// Whether a read reuses an opened Atlas collection. + /// + /// A collection is immutable, so a cached handle stays good until its + /// deletion mask changes. Caching saves the footer read and keeps the + /// decompressed blocks of a collection between queries. + #[envconfig(from = "BEACON_ATLAS_USE_READER_CACHE", default = "true")] + atlas_use_reader_cache: bool, + + /// How many opened Atlas collections to keep. + /// + /// Each entry owns 256 MiB of decompressed blocks and 64 MiB of raw slabs, + /// so this bounds memory as well as handles. + #[envconfig(from = "BEACON_ATLAS_READER_CACHE_SIZE", default = "32")] + atlas_reader_cache_size: u64, + + /// Whether a predicate scan drops the Atlas datasets it can rule out from + /// the collection's statistics, before reading them. + #[envconfig(from = "BEACON_ATLAS_USE_PRUNING", default = "true")] + atlas_use_pruning: bool, + + /// Whether `ANALYZE FILES` measures the column ranges of an Atlas + /// collection. They come from its footer, so they cost no array read. + #[envconfig(from = "BEACON_ATLAS_ENABLE_STATISTICS", default = "true")] + atlas_enable_statistics: bool, + /// The batch size for NetCDF reads, in number of rows. This is used for both local and MPIO reads. #[envconfig(from = "BEACON_BATCH_SIZE", default = "64000")] beacon_batch_size: usize, @@ -685,6 +712,12 @@ impl From for Config { zarr: ZarrConfig { enable_statistics: raw.zarr_enable_statistics, }, + atlas: AtlasConfig { + use_reader_cache: raw.atlas_use_reader_cache, + reader_cache_size: raw.atlas_reader_cache_size, + use_pruning: raw.atlas_use_pruning, + enable_statistics: raw.atlas_enable_statistics, + }, bbf: BbfConfig { split_streams_slice: raw.bbf_split_streams_slice, }, diff --git a/beacon-server/beacon-server/src/main.rs b/beacon-server/beacon-server/src/main.rs index 88a36638..1fdd5a27 100644 --- a/beacon-server/beacon-server/src/main.rs +++ b/beacon-server/beacon-server/src/main.rs @@ -102,6 +102,8 @@ fn install_panic_hook() { /// `RUST_LOG` to see them. const NOISY_DEPENDENCIES: &[&str] = &[ "arrow", + "array_format", + "atlas", "aws_config", "aws_smithy_runtime", "datafusion", diff --git a/beacon-server/beacon-server/src/server/catalog.rs b/beacon-server/beacon-server/src/server/catalog.rs index 607f1de0..97e16124 100644 --- a/beacon-server/beacon-server/src/server/catalog.rs +++ b/beacon-server/beacon-server/src/server/catalog.rs @@ -223,6 +223,7 @@ fn read_function_for_extension(ext: &str) -> Option<&'static str> { "nc" | "cdf" | "netcdf" => "read_netcdf", "arrow" | "arrows" | "ipc" => "read_arrow", "zarr" => "read_zarr", + "atlas" => "read_atlas", "tif" | "tiff" => "read_tiff", "bbf" => "read_bbf", _ => return None, diff --git a/beacon-server/beacon-server/src/server/mod.rs b/beacon-server/beacon-server/src/server/mod.rs index e9fa8655..c3298a80 100644 --- a/beacon-server/beacon-server/src/server/mod.rs +++ b/beacon-server/beacon-server/src/server/mod.rs @@ -278,6 +278,7 @@ async fn build_runtime( .with_netcdf_config(config.netcdf.clone()) .with_hdf5_config(config.hdf5.clone()) .with_zarr_config(config.zarr.clone()) + .with_atlas_config(config.atlas.clone()) .with_sql_settings(SqlSettings { default_table: config.sql.default_table.clone(), enable_pushdown_projection: config.sql.enable_pushdown_projection, diff --git a/docs/docs/2.0.0-rc5/cf-decoding.md b/docs/docs/2.0.0-rc5/cf-decoding.md index 52b2532b..2ec4fc48 100644 --- a/docs/docs/2.0.0-rc5/cf-decoding.md +++ b/docs/docs/2.0.0-rc5/cf-decoding.md @@ -171,10 +171,20 @@ FROM read_netcdf('argo/**/*.nc') LIMIT 0; ``` -## Zarr and Atlas +## Zarr -Zarr and Atlas go through the same decoding path as NetCDF. The attribute names and the rules above -apply unchanged. +Zarr goes through the same decoding path as NetCDF. The attribute names and the rules above apply +unchanged. + +## Atlas + +Atlas decodes nothing, and needs to. Its own types include a nanosecond timestamp, and the ingest +path applies the conventions before the write: `atlas create` reads each NetCDF file with xarray, +which applies `scale_factor`, `add_offset` and the CF time units. A collection therefore holds +decoded values already, and Beacon reads its arrays exactly as they are stored. + +A `units` or `scale_factor` attribute may still sit on an array, kept from the source file. It is +a column like any other attribute — `"temperature.units"` — and it changes no value. ## Next diff --git a/docs/docs/2.0.0-rc5/data-sources/external-tables.md b/docs/docs/2.0.0-rc5/data-sources/external-tables.md index 07d72944..a7b9a2a8 100644 --- a/docs/docs/2.0.0-rc5/data-sources/external-tables.md +++ b/docs/docs/2.0.0-rc5/data-sources/external-tables.md @@ -105,24 +105,24 @@ LOCATION 'sst/*/zarr.json' ### Atlas -An Atlas table points at the `atlas.json` marker file, not at a folder. This is the same as Zarr: +An Atlas table points at the `data.atlas` file itself, not at the folder around it: ```sql CREATE EXTERNAL TABLE sensor_atlas STORED AS ATLAS -LOCATION 'collections/sensor/atlas.json' +LOCATION 'collections/sensor/data.atlas' ``` -Use a glob over the markers to put several Atlas stores in one table: +Use a glob to put several collections in one table: ```sql CREATE EXTERNAL TABLE sensor_atlas STORED AS ATLAS -LOCATION 'collections/*/atlas.json' +LOCATION 'collections/*/data.atlas' ``` -See [Atlas](/docs/2.0.0-rc5/formats/atlas) for the format details. That page -also explains how Atlas speeds up NetCDF and Zarr work. +See [Atlas](/docs/2.0.0-rc5/formats/atlas) for the format details, its `OPTIONS` and its column +naming. That page also explains how Atlas speeds up NetCDF and Zarr work. ### CSV diff --git a/docs/docs/2.0.0-rc5/data-sources/index.md b/docs/docs/2.0.0-rc5/data-sources/index.md index 666e8108..8858aeab 100644 --- a/docs/docs/2.0.0-rc5/data-sources/index.md +++ b/docs/docs/2.0.0-rc5/data-sources/index.md @@ -39,8 +39,8 @@ SELECT * FROM read_csv(['a.csv', 'b.csv']); A glob (`*`, `**`) expands across directories. One query can therefore cover thousands of files. Beacon merges their schemas. It also prunes the files that cannot match your filters. Array formats such as [Zarr](/docs/2.0.0-rc5/formats/zarr) and -[Atlas](/docs/2.0.0-rc5/formats/atlas) use a marker file. Point at `zarr.json` -or `atlas.json`, not at the chunks. +[Atlas](/docs/2.0.0-rc5/formats/atlas) are not plain files. Point at Zarr's `zarr.json` marker, +and at the `data.atlas` file of an Atlas collection, not at the chunks. Each format has its own chapter. The chapter covers the read behaviour, the attribute columns and the limitations. See [File Formats](/docs/2.0.0-rc5/formats/) for the full diff --git a/docs/docs/2.0.0-rc5/faq.md b/docs/docs/2.0.0-rc5/faq.md index f61e0a99..c7b70ac6 100644 --- a/docs/docs/2.0.0-rc5/faq.md +++ b/docs/docs/2.0.0-rc5/faq.md @@ -40,8 +40,8 @@ schema. It reads no data: SELECT * FROM read_netcdf_schema('argo/**/*.nc'); ``` -It covers `parquet`, `netcdf`, `zarr`, `arrow`, `csv`, `bbf` and `tiff`. For GeoParquet, Atlas, Delta -Lake and ODV, use a `LIMIT 0` query. [`SUMMARIZE`](/docs/2.0.0-rc5/sql/summarize) also +It covers `parquet`, `netcdf`, `zarr`, `atlas`, `arrow`, `csv`, `bbf` and `tiff`. For GeoParquet, +Delta Lake and ODV, use a `LIMIT 0` query. [`SUMMARIZE`](/docs/2.0.0-rc5/sql/summarize) also works. ### My files have different columns diff --git a/docs/docs/2.0.0-rc5/formats/atlas.md b/docs/docs/2.0.0-rc5/formats/atlas.md index 5ea86f67..085b7948 100644 --- a/docs/docs/2.0.0-rc5/formats/atlas.md +++ b/docs/docs/2.0.0-rc5/formats/atlas.md @@ -1,5 +1,5 @@ --- -description: Read Atlas array stores with read_atlas(). Dataset pruning with statistics makes range queries over large collections fast. +description: Read Atlas collections with read_atlas(). One file holds thousands of datasets, and their statistics let a range query skip whole datasets before it reads them. --- # Atlas @@ -11,20 +11,19 @@ read_atlas(glob_paths) read_atlas(glob_paths, dimensions) ``` -Beacon reads the [Atlas](/docs/2.0.0-rc5/formats/atlas) array stores that match -one or more glob patterns. Each path must point at an `atlas.json` marker file. Give an exact path -or a glob such as `**/atlas.json`. +Beacon reads the [Atlas](https://github.com/maris-development/atlas) collections that match one or +more glob patterns. A collection is one file, `data.atlas`, so a path names that file. Give an +exact path or a glob such as `**/data.atlas`. -The optional `dimensions` argument selects the arrays with the listed dimension names. Atlas holds -statistics for each column. Beacon drops whole datasets with those statistics. A range query over a -large collection therefore reads only the datasets that can match the predicate. +The optional `dimensions` argument keeps the arrays whose dimensions are all in the list. Use it to +drop the wide grids of a collection and keep its coordinates. ```sql -SELECT * FROM read_atlas('collections/sensor/atlas.json') +SELECT * FROM read_atlas('collections/sensor/data.atlas') --- Combine every Atlas store under a prefix, keeping a subset of dimensions +-- Combine every collection under a prefix, keeping a subset of dimensions SELECT time, temperature -FROM read_atlas(['collections/**/atlas.json'], ['time', 'latitude', 'longitude']) +FROM read_atlas(['collections/**/data.atlas'], ['time', 'latitude', 'longitude']) WHERE time >= '2024-01-01' ``` @@ -33,7 +32,7 @@ WHERE time >= '2024-01-01' Check the columns and the types before you write a query: ```sql -SELECT * FROM read_atlas('collections/sensor/atlas.json') LIMIT 0; +SELECT * FROM read_atlas('collections/sensor/data.atlas') LIMIT 0; ``` [Inspect a schema](/docs/2.0.0-rc5/formats/inspect-a-schema) compares the `_schema` functions, @@ -41,59 +40,110 @@ SELECT * FROM read_atlas('collections/sensor/atlas.json') LIMIT 0; ## Format details -[Atlas](https://github.com/maris-development/atlas) is an array store in a directory. It gives fast -analytical access to multi-dimensional scientific data. Atlas is a file format, like Parquet or -Zarr. Put an Atlas store in the datasets folder. Beacon then finds it and queries it automatically. -You register nothing. An Atlas store is a directory with one `atlas.json` registry. The registry -describes one or more named datasets. Each dataset holds its own arrays. - -What it does: - -- **Dataset pruning with statistics.** Atlas keeps statistics for each dataset and each column. A - query with a predicate, for example a time or latitude range, drops the datasets that cannot - match. Beacon drops them *before it reads any array data*. A range query over a large collection - therefore touches only the relevant data. -- **Column projection.** Beacon reads only the arrays that a query names. The I/O stays proportional - to the selected columns. -- **Compact, self-describing layout.** Atlas compresses the arrays with zstd. Beacon opens the - `atlas.json` registry once and caches it for the life of the process. It therefore parses the - metadata only once. -- **Object storage support.** An Atlas store lives on local disk or on S3-compatible object storage. - -Query an Atlas store with the -[`read_atlas()`](/docs/2.0.0-rc5/sql/table-functions#read-atlas) table function. Point at -the `atlas.json` marker file. Give an exact path or a glob such as `**/atlas.json`. The optional -second argument selects the arrays with the listed dimensions. +Atlas keeps thousands of N-dimensional datasets in one immutable file. A dataset holds what a +NetCDF file holds: named arrays that share dimensions, plus attributes. A collection holds many: -```sql -SELECT * FROM read_atlas(['collections/sensor/atlas.json']) +```text +my_collection/ +├── data.atlas the container: every dataset, then a footer describing them all +└── deleted.mask optional: the datasets a delete has hidden ``` -### External tables over Atlas +Two properties follow, and they are the point of the format: + +- **Metadata is one read.** Opening a collection reads its footer and nothing else. Listing the + datasets, inspecting a schema and reading an attribute are then free. Ten datasets and a million + cost the same. +- **Data arrives chunk by chunk.** Reading a region of an array fetches only the chunks that region + overlaps. + +What Beacon does with that: + +- **Dataset pruning from the footer.** Every array records its minimum, its maximum and its null + count. A query with a predicate — a time or latitude range, say — judges every dataset of a + collection in one pass and never opens the ones that cannot match. A dataset-level attribute is + exact in the footer, so `WHERE ".platform" = 'p3'` prunes on it too. +- **One dataset is one unit of work.** A collection's datasets are spread across every core, and a + worker takes the next one when it is free, so a collection of a million small datasets and one of + four large ones both divide evenly. +- **Column projection.** Only the arrays a query names get read, and only their attributes are + taken from the footer. +- **Object storage.** A collection reads from local disk, S3, GCS, Azure and HTTP alike. -For a stable table name, register the store as an -[external table](/docs/2.0.0-rc5/data-sources/external-tables#atlas). Point the `LOCATION` -at the `atlas.json` marker, as with Zarr. A glob over several markers also works: +### Columns + +| Atlas | Column | +| --- | --- | +| array `temperature` | `temperature` | +| attribute `units` of `temperature` | `temperature.units` | +| dataset attribute `platform` | `.platform` | + +The leading dot on a dataset attribute is what NetCDF and Zarr use too, and it keeps an attribute +from colliding with an array of the same name. Quote such a column: `SELECT ".platform"`. ```sql -CREATE EXTERNAL TABLE sensor_atlas -STORED AS ATLAS -LOCATION 'collections/sensor/atlas.json'; +SELECT temperature, "temperature.units", ".platform" +FROM read_atlas('collections/sensor/data.atlas') +LIMIT 1 +``` -SELECT time, temperature -FROM sensor_atlas -WHERE time >= '2024-01-01'; +### Types and decoding + +Atlas stores its own types, including a native nanosecond timestamp, so **Beacon applies no CF +decoding to a collection**. The ingest path does it instead: `atlas create` reads each NetCDF file +with xarray, which applies `scale_factor`, `add_offset` and the CF time units before the write. An +array therefore reads back exactly as it is stored. This is the one place Atlas differs from +[NetCDF](/docs/2.0.0-rc5/formats/netcdf) and [Zarr](/docs/2.0.0-rc5/formats/zarr) — see +[CF decoding](/docs/2.0.0-rc5/cf-decoding). + +A cell nobody wrote reads as the array's fill value, and the fill reads as null. Two consequences +are worth knowing: + +- A float array ingested from NetCDF carries a `NaN` fill, and `NaN` never equals itself, so a + `NaN` cell reads as `NaN` rather than as null. That is the same rule every Beacon format follows. +- A string array carries an empty-string fill, so an empty string reads as null. The ingest cannot + store a null string, so this mirrors what was written. + +Not readable as columns: a `Bool` array, a `List` or `FixedSizeList` array, and a list-valued +attribute. Each is dropped from the schema rather than failing the query. + +### Two datasets that disagree + +Atlas reconciles nothing: two datasets may declare one array name with two types. Beacon merges +them the way it merges the files of any other format. Two numeric types widen to one that holds +both. Two different families — a number and a string — refuse the table by name: + +```text +Incompatible types for field 'value': Utf8 in 'obs/data.atlas#a' vs Int64 in 'obs/data.atlas#b' ``` +Set `BEACON_TYPE_WIDENING_ON_CONFLICT=keep_first` to take the first dataset's type instead. See +[Configuration](/docs/2.0.0-rc5/server/configuration#query-engine). + +### Building a collection + +`pip install atlas-python` gives the `atlas` command. Point it at a directory of NetCDF files: + +```bash +atlas create /data/argo /collections/argo +``` + +That writes `/collections/argo/data.atlas`, one dataset per file, named after the file. It works +against a local path and against a bucket. See the +[Atlas documentation](https://github.com/maris-development/atlas). + +:::warning Collections written before Atlas 0.16 +Atlas used to be a directory of per-array files behind an `atlas.json` registry. Beacon reads the +single-file format only, so such a directory is passed over rather than read. Rewrite it with +`atlas create`, then point at the `data.atlas` it produces. +::: + ### Optimize NetCDF and Zarr with Atlas Do you query a large NetCDF or Zarr collection often? Then convert the source files into one Atlas -collection. Atlas merges many files into one store with statistics. Beacon can then drop whole -datasets with the column statistics. It reads only the arrays that you select. A spatial or time -range query is therefore much faster than a scan of the original files. - -The [Atlas repository](https://github.com/maris-development/atlas) documents the store format. It -also holds the tools that build an Atlas collection. +collection. Atlas merges many files into one container with statistics, so Beacon drops whole +datasets before it reads an array. A spatial or time range query is therefore much faster than a +scan of the original files. :::tip Cache a large, repeated aggregation with a @@ -103,24 +153,43 @@ collection and over any other table. Run `REFRESH` when the source data changes. ## As an external table -An Atlas table points at the `atlas.json` marker file, not at a folder. This is the same as Zarr: +An Atlas table points at the `data.atlas` file itself, not at the folder around it: ```sql CREATE EXTERNAL TABLE sensor_atlas STORED AS ATLAS -LOCATION 'collections/sensor/atlas.json' +LOCATION 'collections/sensor/data.atlas' ``` -Use a glob over the markers to put several Atlas stores in one table: +Use a glob to put several collections in one table: ```sql CREATE EXTERNAL TABLE sensor_atlas STORED AS ATLAS -LOCATION 'collections/*/atlas.json' +LOCATION 'collections/*/data.atlas' ``` -See [Atlas](/docs/2.0.0-rc5/formats/atlas) for the format details. That page -also explains how Atlas speeds up NetCDF and Zarr work. +See [Create External Tables](/docs/2.0.0-rc5/data-sources/external-tables) for the full DDL. See +[Data Sources](/docs/2.0.0-rc5/data-sources/) for the full read model. + +### `OPTIONS` + +`STORED AS ATLAS` reads four keys: + +| Option | Type | Default | Description | +| --- | --- | --- | --- | +| `read_dimensions` | List of dimension names | The default grid of each dataset | The dimensions the table reads. An array survives only when the list holds every one of its own. | +| `use_pruning` | Boolean | `true` (`BEACON_ATLAS_USE_PRUNING`) | Drop the datasets a predicate rules out before reading them. Turning it off only costs speed: pruning never changes an answer. | +| `use_reader_cache` | Boolean | `true` (`BEACON_ATLAS_USE_READER_CACHE`) | Reuse an opened collection across queries. | +| `enable_statistics` | Boolean | `true` (`BEACON_ATLAS_ENABLE_STATISTICS`) | Whether `ANALYZE FILES` records this collection's column ranges. A query never measures a collection, so this affects the analyzer alone. | + +```sql +CREATE EXTERNAL TABLE sensor_atlas +STORED AS ATLAS +LOCATION 'collections/*/data.atlas' +OPTIONS ('read_dimensions' 'time,lat,lon') +``` -See [Create External Tables](/docs/2.0.0-rc5/data-sources/external-tables) for the full DDL. See [Data Sources](/docs/2.0.0-rc5/data-sources/) for the -full read model. +See [`OPTIONS`](/docs/2.0.0-rc5/sql/create-external-table#options) for the rules that hold for every +key. See [Arrays to tables](/docs/2.0.0-rc5/arrays-to-tables#the-dimensions-argument) for the grid +rule. diff --git a/docs/docs/2.0.0-rc5/formats/index.md b/docs/docs/2.0.0-rc5/formats/index.md index 815ad933..dc15d1af 100644 --- a/docs/docs/2.0.0-rc5/formats/index.md +++ b/docs/docs/2.0.0-rc5/formats/index.md @@ -33,7 +33,7 @@ S3-compatible bucket, chosen at startup. See | [NetCDF](/docs/2.0.0-rc5/formats/netcdf) | `read_netcdf` | `NC` | `.nc` | | [HDF5](/docs/2.0.0-rc5/formats/hdf5) | `read_hdf5` | `HDF5`, `H5` | `.h5`, `.hdf5` | | [Zarr](/docs/2.0.0-rc5/formats/zarr) | `read_zarr` | `ZARR` | `zarr.json` marker | -| [Atlas](/docs/2.0.0-rc5/formats/atlas) | `read_atlas` | `ATLAS` | `atlas.json` marker | +| [Atlas](/docs/2.0.0-rc5/formats/atlas) | `read_atlas` | `ATLAS` | `data.atlas` file | | [GeoTIFF / COG](/docs/2.0.0-rc5/formats/geotiff) | `read_tiff` | `TIFF` | `.tif`, `.tiff` | | [BBF](/docs/2.0.0-rc5/formats/bbf) | `read_bbf` | `BBF` | `.bbf` | | [Delta Lake](/docs/2.0.0-rc5/formats/delta-lake) | `read_delta` | `DELTA` | `_delta_log/` directory | @@ -58,7 +58,7 @@ The table above says how to read each format. This one says what you get. | NetCDF | **Anonymous only**, or full with the Rust reader | Projection + dimension selection | Yes | `read_netcdf_schema` | Yes | | HDF5 | **Anonymous only**, or full with the Rust reader | Projection + dimension selection | Yes | `read_hdf5_schema` | Yes | | Zarr | Full | Projection + dimension selection, chunk pruning | No | `read_zarr_schema` | Yes | -| Atlas | Full | Predicate + projection, **file-level pruning** | Yes | `read_atlas_schema` | Yes | +| Atlas | Full | Predicate + projection, **dataset-level pruning** | Yes | `read_atlas_schema` | Yes | | GeoTIFF / COG | Full | Projection, range requests | No | `read_tiff_schema` | Yes | | BBF | Full | Predicate + projection | No | `read_bbf_schema` | No | | Delta Lake | Full | Predicate + projection, file skipping | No, but see below | `read_delta_schema` | No | @@ -77,8 +77,8 @@ Reading the columns: query — paths in SQL are relative either way. - **Pushdown** — how much of a query reaches storage instead of running after the read. *Predicate* means a `WHERE` clause prunes data. *Projection* means a narrow `SELECT` reads fewer columns. - [Atlas](/docs/2.0.0-rc5/formats/atlas) is the strongest: its collection statistics drop whole - files before any array is opened. + [Atlas](/docs/2.0.0-rc5/formats/atlas) is the strongest: the statistics in a collection's footer + drop whole datasets before any array is opened. - **Query output** — whether a query result can be written back in that format, with `COPY TO` or an `output.format` on the API. Writing rows into an existing table is a different capability: **Delta Lake** external tables accept `INSERT INTO`, and @@ -120,9 +120,9 @@ WHERE depth < 100 GROUP BY platform; ``` -Array formats such as [Zarr](/docs/2.0.0-rc5/formats/zarr) and -[Atlas](/docs/2.0.0-rc5/formats/atlas) use a marker file. Point at `zarr.json` -or `atlas.json`, not at the chunk files. +[Zarr](/docs/2.0.0-rc5/formats/zarr) is a directory behind a marker file: point at `zarr.json`, +not at the chunk files. An [Atlas](/docs/2.0.0-rc5/formats/atlas) collection is a single file: +point at `data.atlas`. Some files share a schema but have different columns. Combine those files with [`UNION BY NAME`](/docs/2.0.0-rc5/sql/union-by-name). diff --git a/docs/docs/2.0.0-rc5/guides/speed-up-queries.md b/docs/docs/2.0.0-rc5/guides/speed-up-queries.md index f3bbfbff..364eea60 100644 --- a/docs/docs/2.0.0-rc5/guides/speed-up-queries.md +++ b/docs/docs/2.0.0-rc5/guides/speed-up-queries.md @@ -65,7 +65,7 @@ them *before it reads any array data*. It then reads only the arrays that you se ```sql CREATE EXTERNAL TABLE sensor_atlas STORED AS ATLAS -LOCATION 'collections/*/atlas.json'; +LOCATION 'collections/*/data.atlas'; ``` ## 6. Compute once, query many times diff --git a/docs/docs/2.0.0-rc5/how-it-works.md b/docs/docs/2.0.0-rc5/how-it-works.md index b90ee061..e38bd386 100644 --- a/docs/docs/2.0.0-rc5/how-it-works.md +++ b/docs/docs/2.0.0-rc5/how-it-works.md @@ -45,8 +45,9 @@ moves filters and column selections as close to the data as possible: variables reads about 3 columns of bytes. - **Predicate pushdown**: Beacon turns filters into file, row group and chunk pruning. A time range filter on a [Zarr](/docs/2.0.0-rc5/formats/zarr) store fetches only the - chunks in that range. On [Atlas](/docs/2.0.0-rc5/formats/atlas), Beacon uses - the stored statistics. It can drop whole datasets before it reads any array data. + chunks in that range. On [Atlas](/docs/2.0.0-rc5/formats/atlas), every array records its own + range in the collection's footer, so Beacon judges every dataset in one pass and never opens + the ones that cannot match. - **Federated pushdown**: Beacon sends filters, projections, limits and whole aggregates to [SQL databases](/docs/2.0.0-rc5/data-sources/sql-databases) and to [remote Beacons](/docs/2.0.0-rc5/data-sources/remote-tables). Only the reduced result diff --git a/docs/docs/2.0.0-rc5/server/configuration.md b/docs/docs/2.0.0-rc5/server/configuration.md index 1320cb9f..cf86a7be 100644 --- a/docs/docs/2.0.0-rc5/server/configuration.md +++ b/docs/docs/2.0.0-rc5/server/configuration.md @@ -308,8 +308,10 @@ which values a store holds, so a store may hold values outside them. | Variable | Default | Description | | --- | --- | --- | -| `BEACON_ATLAS_USE_READER_CACHE` | `true` | Cache opened Atlas store readers in memory, avoiding re-opening the same `atlas.json` across queries. | -| `BEACON_ATLAS_READER_CACHE_SIZE` | `32` | Max Atlas reader entries to keep cached. | +| `BEACON_ATLAS_USE_READER_CACHE` | `true` | Keep opened Atlas collections in memory, so a repeated query does not read the footer of `data.atlas` again. | +| `BEACON_ATLAS_READER_CACHE_SIZE` | `32` | How many opened collections to keep. Each holds its own 256 MiB block cache and 64 MiB slab cache, so this bounds memory as well as handles. | +| `BEACON_ATLAS_USE_PRUNING` | `true` | Skip the datasets a predicate rules out from the collection's own statistics, before reading them. Off only costs speed; pruning never changes an answer. | +| `BEACON_ATLAS_ENABLE_STATISTICS` | `true` | Let `ANALYZE FILES` record a collection's column ranges. They come from the footer, so they cost no array read. | ### Beacon Binary Format (BBF) diff --git a/docs/docs/2.0.0-rc5/server/crawlers.md b/docs/docs/2.0.0-rc5/server/crawlers.md index b1831029..fe226dc9 100644 --- a/docs/docs/2.0.0-rc5/server/crawlers.md +++ b/docs/docs/2.0.0-rc5/server/crawlers.md @@ -189,18 +189,23 @@ endpoints: The crawler finds **one file per dataset** formats. The file extension must equal the format identifier exactly. The identifiers are `parquet`, `geoparquet`, `csv`, `nc` for NetCDF, `bbf`, -`arrow` and `tiff`. The crawler does **not** read an alias extension. A file must use the canonical +`arrow`, `tiff` and `atlas`. The crawler does **not** read an alias extension. A file must use the canonical extension. The crawler therefore skips `.tsv` (CSV), `.feather` (Arrow) and `.tif` (TIFF). The readers open those files directly. Register such a file with a table function or with `CREATE EXTERNAL TABLE`. -The crawler **skips** a store with a directory and a marker file. **Zarr** (`*.zarr/zarr.json`) and -**Atlas** (`atlas.json`) are such stores. The listing path does not register them as external -tables. Read a Zarr store with +The crawler **skips** a store with a directory and a marker file. **Zarr** (`*.zarr/zarr.json`) is +such a store. The listing path does not register it as an external table. Read a Zarr store with [`read_zarr`](/docs/2.0.0-rc5/sql/table-functions#read-zarr). A crawl ignores these stores and continues with the other datasets. Register them with a table function or with `CREATE EXTERNAL TABLE`. +An **Atlas** collection *is* crawled. A collection is one file, `data.atlas`, so its extension is +its format and the rule above admits it. Each collection lives in its own directory and tables +group by directory, so a crawl of several collections registers one table each. Use +`CREATE EXTERNAL TABLE ... LOCATION 'collections/*/data.atlas'` to put them in one table +instead. + The crawler matches a GeoParquet file by the `.geoparquet` extension. It creates a GeoParquet table with GeoArrow geometry decoding. The crawler reads a plain `.parquet` file as an ordinary Parquet table, also with `geo` metadata. To get geometry decoding for such a file, give it the diff --git a/docs/docs/2.0.0-rc5/server/datasets.md b/docs/docs/2.0.0-rc5/server/datasets.md index b24fea03..44c09d36 100644 --- a/docs/docs/2.0.0-rc5/server/datasets.md +++ b/docs/docs/2.0.0-rc5/server/datasets.md @@ -21,7 +21,7 @@ The default local path in the Docker container is `/beacon/data/datasets/`. | [Arrow IPC](/docs/2.0.0-rc5/formats/arrow) | `.arrow`, `.feather` | `ARROW` | `read_arrow` | yes (`ipc`) | | [NetCDF](/docs/2.0.0-rc5/formats/netcdf) | `.nc` | `NC` | `read_netcdf` | yes (+ ND-NetCDF) | | [Zarr](/docs/2.0.0-rc5/formats/zarr) | `zarr.json` marker | `ZARR` | `read_zarr` | no | -| [Atlas](/docs/2.0.0-rc5/formats/atlas) | `atlas.json` marker | `ATLAS` | `read_atlas` | no | +| [Atlas](/docs/2.0.0-rc5/formats/atlas) | `data.atlas` file | `ATLAS` | `read_atlas` | yes | | [GeoTIFF / COG](/docs/2.0.0-rc5/formats/geotiff) | `.tif`, `.tiff` | `TIFF` | `read_tiff` | no | | [BBF](/docs/2.0.0-rc5/formats/bbf) | `.bbf` | `BBF` | `read_bbf` | no | | [Delta Lake](/docs/2.0.0-rc5/formats/delta-lake) | `_delta_log/` directory | `DELTA` | `read_delta` | no | diff --git a/docs/docs/2.0.0-rc5/server/performance-tuning.md b/docs/docs/2.0.0-rc5/server/performance-tuning.md index 436d1f3b..cd1a8123 100644 --- a/docs/docs/2.0.0-rc5/server/performance-tuning.md +++ b/docs/docs/2.0.0-rc5/server/performance-tuning.md @@ -217,20 +217,32 @@ with statistics, next to the chunk pruning. It drops whole datasets before it re ## Atlas Tuning -Beacon opens an [Atlas](/docs/2.0.0-rc5/formats/atlas) store through its -`atlas.json` registry. Beacon caches the open Atlas readers. It therefore does not open the same -store for every query. +Beacon opens an [Atlas](/docs/2.0.0-rc5/formats/atlas) collection by reading the footer of its +`data.atlas` file. Beacon caches the open collections. It therefore does not read that footer again +for every query. ### Reader cache (no repeated store open) #### `BEACON_ATLAS_USE_READER_CACHE` and `BEACON_ATLAS_READER_CACHE_SIZE` -With the reader cache on, Beacon uses an open Atlas reader again. It therefore does not parse the -`atlas.json` registry for every query. +With the reader cache on, Beacon uses an open collection again. It therefore does not read the +footer for every query, and the decompressed blocks of a collection stay warm between them. + +Each cached collection holds its own block cache — 256 MiB of decompressed blocks and 64 MiB of raw +slabs — so the size is a memory bound as much as a handle count. Recommendations: - Keep `BEACON_ATLAS_USE_READER_CACHE=true`, the default, when several queries read the same Atlas collections. -- Increase `BEACON_ATLAS_READER_CACHE_SIZE`, default `32`, if you query more Atlas stores than the - cache holds. +- Increase `BEACON_ATLAS_READER_CACHE_SIZE`, default `32`, if you query more Atlas collections than + the cache holds. Remember the memory bound above before raising it far. + +#### `BEACON_ATLAS_USE_PRUNING` + +On by default. A query with a predicate judges every dataset of a collection against the statistics +in its footer, in one pass, and never opens the ones that cannot match. Turning it off only costs +speed: pruning never changes an answer, and the filter above the scan still decides every row. + +`EXPLAIN ANALYZE` reports what it did as `atlas_datasets_pruned`, `atlas_datasets_scanned` and +`atlas_index_rows`. diff --git a/docs/docs/2.0.0-rc5/sql/create-external-table.md b/docs/docs/2.0.0-rc5/sql/create-external-table.md index d04e3149..218ffdb5 100644 --- a/docs/docs/2.0.0-rc5/sql/create-external-table.md +++ b/docs/docs/2.0.0-rc5/sql/create-external-table.md @@ -46,7 +46,7 @@ CREATE EXTERNAL TABLE argo STORED AS NC LOCATION 'argo/**/*.nc' | `NC` | `.nc` | [NetCDF](/docs/2.0.0-rc5/formats/netcdf) | | `HDF5`, `H5` | `.h5`, `.hdf5` | [HDF5](/docs/2.0.0-rc5/formats/hdf5) | | `ZARR` | Zarr v3 (`zarr.json`) | [Zarr](/docs/2.0.0-rc5/formats/zarr) | -| `ATLAS` | Atlas array store (`atlas.json`) | [Atlas](/docs/2.0.0-rc5/formats/atlas) | +| `ATLAS` | Atlas collection (`data.atlas`) | [Atlas](/docs/2.0.0-rc5/formats/atlas) | | `CSV` | `.csv`, `.tsv` | [CSV](/docs/2.0.0-rc5/formats/csv) | | `ARROW` | Arrow IPC (`.arrow`, `.feather`) | [Arrow IPC](/docs/2.0.0-rc5/formats/arrow) | | `TIFF` | GeoTIFF / Cloud-Optimized GeoTIFF | [GeoTIFF](/docs/2.0.0-rc5/formats/geotiff) | @@ -69,13 +69,13 @@ federate a table in an external SQL database. See [SQL Databases](/docs/2.0.0-rc5/data-sources/sql-databases). Their `LOCATION` is the remote table name. The connection details go in `OPTIONS`, with an encrypted `password`. -A Zarr table must point at a `zarr.json` entry file. An Atlas table must point at an `atlas.json` -marker: +A Zarr table must point at a `zarr.json` entry file. An Atlas collection is a single file, so an +Atlas table points at that file: ```sql CREATE EXTERNAL TABLE sst STORED AS ZARR LOCATION 'sst/*/zarr.json' -CREATE EXTERNAL TABLE sensor STORED AS ATLAS LOCATION 'sensor/atlas.json' +CREATE EXTERNAL TABLE sensor STORED AS ATLAS LOCATION 'sensor/data.atlas' ``` `GEOPARQUET` reads Parquet files. Beacon decodes their geometry columns to native GeoArrow. See @@ -165,6 +165,7 @@ validates it and then reads the server setting alone. See the format page of eac | `NC` | `read_dimensions`, `use_rust_reader`, `enable_statistics` | [NetCDF](/docs/2.0.0-rc5/formats/netcdf#options) | | `HDF5`, `H5` | `read_dimensions`, `use_rust_reader`, `enable_statistics`, `unify_phony_dimensions`, `convention` | [HDF5](/docs/2.0.0-rc5/formats/hdf5#options) | | `ZARR` | `read_dimensions`, `enable_statistics` | [Zarr](/docs/2.0.0-rc5/formats/zarr#options) | +| `ATLAS` | `read_dimensions`, `use_pruning`, `use_reader_cache`, `enable_statistics` | [Atlas](/docs/2.0.0-rc5/formats/atlas#options) | | `CSV` | `delimiter`, `infer_records` | [CSV](/docs/2.0.0-rc5/formats/csv#options) | | `BBF` | `split_streams_slice` | [BBF](/docs/2.0.0-rc5/formats/bbf#options) | | `DELTA` | `version`, `timestamp` | [Delta Lake](/docs/2.0.0-rc5/formats/delta-lake#options) | diff --git a/docs/docs/2.0.0-rc5/sql/table-functions.md b/docs/docs/2.0.0-rc5/sql/table-functions.md index 49d49ae1..dd9f21be 100644 --- a/docs/docs/2.0.0-rc5/sql/table-functions.md +++ b/docs/docs/2.0.0-rc5/sql/table-functions.md @@ -115,20 +115,20 @@ read_atlas(glob_paths) read_atlas(glob_paths, dimensions) ``` -Beacon reads the [Atlas](/docs/2.0.0-rc5/formats/atlas) array stores that match -one or more glob patterns. Each path must point at an `atlas.json` marker file. Give an exact path -or a glob such as `**/atlas.json`. +Beacon reads the [Atlas](/docs/2.0.0-rc5/formats/atlas) collections that match one or more glob +patterns. A collection is one file, `data.atlas`, so a path names that file. Give an exact path or +a glob such as `**/data.atlas`. -The optional `dimensions` argument selects the arrays with the listed dimension names. Atlas holds -statistics for each column. Beacon drops whole datasets with those statistics. A range query over a -large collection therefore reads only the datasets that can match the predicate. +The optional `dimensions` argument keeps the arrays whose dimensions are all in the list. Every +array of a collection records its own minimum and maximum, so a range query judges every dataset +in one pass and reads only the ones that can match. ```sql -SELECT * FROM read_atlas('collections/sensor/atlas.json') +SELECT * FROM read_atlas('collections/sensor/data.atlas') --- Combine every Atlas store under a prefix, keeping a subset of dimensions +-- Combine every collection under a prefix, keeping a subset of dimensions SELECT time, temperature -FROM read_atlas(['collections/**/atlas.json'], ['time', 'latitude', 'longitude']) +FROM read_atlas(['collections/**/data.atlas'], ['time', 'latitude', 'longitude']) WHERE time >= '2024-01-01' ``` diff --git a/integration-tests/formats/README.md b/integration-tests/formats/README.md index afba726a..d6f73e95 100644 --- a/integration-tests/formats/README.md +++ b/integration-tests/formats/README.md @@ -15,6 +15,7 @@ pytest formats/test_netcdf.py -v # one | `test_parquet.py` | `pyarrow` | row groups, zstd and snappy, statistics on and off, nulls, a struct column, a directory | | `test_csv.py` | text | the delimiter argument, quoting, nulls, type inference, gzip, a glob | | `test_zarr.py` | `zarr` | v3, consolidated metadata, the zstd codec, a rank-3 grid, an unwritten chunk | +| `test_atlas.py` | `atlas-python` | many datasets in one file, attributes, dataset pruning, a glob, `OPTIONS` | | `test_hdf5.py` | `h5py` | plain HDF5, dimension scales, a nested group, a compound dataset, strings | | `test_arrow.py` | `pyarrow` | one batch and many, a dictionary column, nulls, every Arrow type | | `test_geoparquet.py` | `geopandas` | WKB and GeoArrow, a line, a polygon, a covering bbox, `ST_X`/`ST_Y` | diff --git a/integration-tests/formats/test_atlas.py b/integration-tests/formats/test_atlas.py new file mode 100644 index 00000000..8aae9291 --- /dev/null +++ b/integration-tests/formats/test_atlas.py @@ -0,0 +1,241 @@ +"""Atlas, end to end, in one file. + +Writes its own Atlas collections, opens an embedded Beacon over them, queries them, creates an +external table, reopens the database, and checks the table survived. + + pytest formats/test_atlas.py -v + +Needs `atlas-python`, `netCDF4` and the `beacondb` extension. All are skipped cleanly if absent. + +An Atlas collection is one write-once file, `data.atlas`, holding many datasets. `atlas.create` +builds one from a directory of netCDF files: one dataset per file, named after the file. A +`LOCATION` therefore names the container, not the directory around it. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +beacondb = pytest.importorskip("beacondb", reason="build it with maturin") +atlas = pytest.importorskip("atlas", reason="pip install atlas-python") +pytest.importorskip("netCDF4", reason="pip install netCDF4") + +import numpy as np # noqa: E402 +from netCDF4 import Dataset # noqa: E402 + +#: netCDF4 1.7.4 sets `.shape` on a numpy array, which numpy 2.5 deprecated. It comes from +#: inside the library on every write and there is no writer call that avoids it. +pytestmark = pytest.mark.filterwarnings( + "ignore:Setting the shape on a NumPy array has been deprecated" +) + +ROWS = 8 +#: Datasets per collection, one per source file. +FILES = 5 +TOTAL_ROWS = ROWS * FILES + + +def _source_file(path: Path, index: int) -> None: + """One netCDF file: `temperature` and `depth` over `obs`, plus attributes. + + File `i` covers temperatures `[10i, 10i + 7]`, so each dataset lands in a range of its own + and a threshold has an answer that can be written down. + """ + with Dataset(path, "w", format="NETCDF4") as ds: + ds.createDimension("obs", ROWS) + temperature = ds.createVariable("temperature", "f4", ("obs",)) + temperature[:] = np.arange(ROWS, dtype="float32") + index * 10 + temperature.units = "celsius" + depth = ds.createVariable("depth", "f4", ("obs",)) + depth[:] = np.arange(ROWS, dtype="float32") * 10.0 + ds.platform = f"p{index}" + + +@pytest.fixture(scope="module") +def datasets(tmp_path_factory) -> Path: + """Build every collection this module queries.""" + root = tmp_path_factory.mktemp("atlas") + + # The netCDF files `atlas.create` ingests. They are not queried themselves. + source = tmp_path_factory.mktemp("atlas-source") + for index in range(FILES): + _source_file(source / f"d{index}.nc", index) + + atlas.create(source, root / "obs") + + # A second collection under a nested prefix, so a glob has something to cover. + nested = tmp_path_factory.mktemp("atlas-source-nested") + _source_file(nested / "extra.nc", FILES) + atlas.create(nested, root / "more" / "obs") + + return root + + +@pytest.fixture +def con(datasets, tmp_path): + with beacondb.connect(str(tmp_path / "beacon.db"), datasets=str(datasets)) as connection: + yield connection + + +# --- reading ------------------------------------------------------------------ + + +def test_a_collection_reads_every_dataset(con): + """One dataset per source file, and every row of each.""" + assert con.sql( + "SELECT count(*) AS n FROM read_atlas('obs/data.atlas')" + ).fetchall() == [(TOTAL_ROWS,)] + + +def test_the_schema_is_reported(con): + relation = con.sql("SELECT temperature, depth FROM read_atlas('obs/data.atlas')") + assert relation.columns == ["temperature", "depth"] + + +def test_values_read_back(con): + rows = con.sql( + "SELECT temperature, depth FROM read_atlas('obs/data.atlas') " + "ORDER BY temperature LIMIT 2" + ).fetchall() + assert rows == [(0.0, 0.0), (1.0, 10.0)] + + +def test_a_filter_and_an_aggregate(con): + """File `i` covers `[10i, 10i + 7]`, so `>= 20` keeps the last three files whole.""" + got = con.sql( + "SELECT count(*) n, min(temperature) lo, max(temperature) hi " + "FROM read_atlas('obs/data.atlas') WHERE temperature >= 20" + ).fetchall()[0] + assert got == (ROWS * 3, 20.0, 47.0) + + +def test_a_predicate_nothing_meets_returns_nothing(con): + """Every dataset is ruled out by its own statistics, and none is opened.""" + assert con.sql( + "SELECT count(*) AS n FROM read_atlas('obs/data.atlas') WHERE temperature > 100000" + ).fetchall() == [(0,)] + + +def test_two_reads_return_the_same_rows(con): + """Datasets are read in parallel and land in completion order, so compare with ORDER BY.""" + query = "SELECT temperature FROM read_atlas('obs/data.atlas') ORDER BY temperature" + assert con.sql(query).fetchall() == con.sql(query).fetchall() + + +def test_a_glob_covers_several_collections(con): + assert con.sql( + "SELECT count(*) AS n FROM read_atlas('**/data.atlas')" + ).fetchall() == [(TOTAL_ROWS + ROWS,)] + + +# --- attributes --------------------------------------------------------------- + + +def test_a_variable_attribute_is_a_column(con): + """A per-array attribute is `{array}.{attr}`, as it is for netCDF and Zarr.""" + units = con.sql( + 'SELECT DISTINCT "temperature.units" AS u FROM read_atlas(\'obs/data.atlas\')' + ).fetchall() + assert units == [("celsius",)] + + +def test_a_dataset_attribute_is_a_column_under_a_dot(con): + """A collection-level attribute of the source file becomes `.{attr}`. + + The leading dot is what keeps an attribute from colliding with an array of the same name. + """ + platforms = con.sql( + 'SELECT DISTINCT ".platform" AS p FROM read_atlas(\'obs/data.atlas\') ORDER BY p' + ).fetchall() + assert platforms == [(f"p{i}",) for i in range(FILES)] + + +def test_an_attribute_predicate_selects_one_dataset(con): + """An attribute is exact in the footer, so a predicate on one reaches its dataset alone.""" + got = con.sql( + "SELECT count(*) n, min(temperature) lo, max(temperature) hi " + "FROM read_atlas('obs/data.atlas') WHERE \".platform\" = 'p3'" + ).fetchall()[0] + assert got == (ROWS, 30.0, 37.0) + + +# --- the schema function ------------------------------------------------------ + + +def test_the_schema_function_reports_the_columns(con): + columns = { + row[0] + for row in con.sql( + "SELECT column_name FROM read_atlas_schema('obs/data.atlas')" + ).fetchall() + } + assert {"temperature", "depth", "temperature.units", ".platform"} <= columns + + +# --- the external table ------------------------------------------------------- + + +def test_an_external_table_names_the_container(con): + """A `LOCATION` points at `data.atlas` itself, not at the directory holding it.""" + con.execute("CREATE EXTERNAL TABLE obs STORED AS ATLAS LOCATION 'obs/data.atlas'") + assert con.sql("SELECT count(*) AS n FROM obs").fetchall() == [(TOTAL_ROWS,)] + assert "obs" in con.list_tables() + + +def test_an_external_table_takes_a_glob(con): + con.execute("CREATE EXTERNAL TABLE every STORED AS ATLAS LOCATION '**/data.atlas'") + assert con.sql("SELECT count(*) AS n FROM every").fetchall() == [(TOTAL_ROWS + ROWS,)] + + +def test_an_external_table_survives_a_restart(datasets, tmp_path): + path = str(tmp_path / "restart.db") + + with beacondb.connect(path, datasets=str(datasets)) as con: + con.execute("CREATE EXTERNAL TABLE obs STORED AS ATLAS LOCATION 'obs/data.atlas'") + + with beacondb.connect(path, datasets=str(datasets)) as con: + assert con.sql("SELECT count(*) AS n FROM obs").fetchall() == [(TOTAL_ROWS,)] + assert "obs" in con.list_tables() + rows = con.sql("SELECT temperature FROM obs ORDER BY temperature LIMIT 1").fetchall() + assert rows == [(0.0,)] + + +def test_pruning_can_be_turned_off_per_table(datasets, tmp_path): + """The switch changes what is read, never what is returned.""" + with beacondb.connect(str(tmp_path / "options.db"), datasets=str(datasets)) as con: + con.execute( + "CREATE EXTERNAL TABLE pruned STORED AS ATLAS LOCATION 'obs/data.atlas' " + "OPTIONS ('use_pruning' 'true')" + ) + con.execute( + "CREATE EXTERNAL TABLE whole STORED AS ATLAS LOCATION 'obs/data.atlas' " + "OPTIONS ('use_pruning' 'false')" + ) + query = "SELECT temperature FROM {} WHERE temperature >= 20 ORDER BY temperature" + assert con.sql(query.format("pruned")).fetchall() == con.sql( + query.format("whole") + ).fetchall() + + +# --- what is not supported ---------------------------------------------------- + + +def test_a_pre_0_16_collection_is_passed_over(con, datasets): + """Atlas before 0.16 was a directory behind an `atlas.json` registry. + + This build reads the single-file format alone. Such a directory holds no container, so + nothing recognises it and no row of it reaches a query. Whether that surfaces as an error + or as an empty result is not the point and not asserted — what matters is that a stale + collection is never half-read as if it were a current one. + """ + legacy = datasets / "legacy" + legacy.mkdir(exist_ok=True) + (legacy / "atlas.json").write_text("{}") + + try: + rows = con.sql("SELECT * FROM read_atlas('legacy/atlas.json')").fetchall() + except Exception: + return + assert rows == [], "a pre-0.16 directory holds no container, so it contributes no rows" diff --git a/integration-tests/requirements.txt b/integration-tests/requirements.txt index dc8f050a..10a2bb05 100644 --- a/integration-tests/requirements.txt +++ b/integration-tests/requirements.txt @@ -8,6 +8,8 @@ pyarrow==21.0.0 # upgrade changes the bytes of a file, and a new failure after one is a signal worth keeping. netCDF4==1.7.4 zarr==3.3.0 +# Builds an Atlas collection from netCDF files; the Rust side reads it. +atlas-python==0.16.4 h5py==3.14.0 geopandas==1.1.4 rasterio==1.5.1 From 846d22446f9e18e8be924e99de749cf321db1b65 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 2 Sep 2026 12:43:52 +0200 Subject: [PATCH 02/16] remove the atlas rebuild plan The plan served the rebuild and is finished. What it decided is in the code and its doc comments, what it changed for a user is in the changelog and the format page, and the four requests it made of atlas-rust belong upstream rather than in this repository. --- .../beacon-arrow-atlas/REBUILD_PLAN.md | 906 ------------------ 1 file changed, 906 deletions(-) delete mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/REBUILD_PLAN.md diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/REBUILD_PLAN.md b/beacon-db/beacon-file-formats/beacon-arrow-atlas/REBUILD_PLAN.md deleted file mode 100644 index ea536617..00000000 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/REBUILD_PLAN.md +++ /dev/null @@ -1,906 +0,0 @@ -# Rebuild `beacon-arrow-atlas` on atlas-rust 0.16 - -Status: complete, phases 0 to 5. Date: 2026-09-02. -Scope: this crate, its wiring in core, server, clients, tests and docs. - -| Phase | State | -| --- | --- | -| 0 — prerequisites | done: the crate is a workspace member and pins `atlas-rust = "=0.16.4"` | -| 1 — the reader | done: `config`, `store`, `compat`, `backend`, `reader`, `test_support` | -| 2 — the scan | done: `datafusion/{mod,source,options,metrics,table_function}`, `FileRead::skipped`; 74 atlas + 243 nd-array tests pass, clippy clean | -| 3 — pruning and statistics | done: `pruning`, `statistics`; 105 tests pass, clippy clean | -| 4 — wiring | done: core, functions, server, config, clients, changelog; 8 end-to-end tests in `beacon-core/tests/atlas.rs` | -| 5 — integration and docs | done: `formats/test_atlas.py`, the format page rewritten and 13 other pages updated; the docs site builds | - -Sections 11 to 15 record what each phase settled that this plan had wrong. - -## 1. Why a rebuild - -- This crate targets atlas-rust 0.14. Version 0.14 stores a collection as a - directory: an `atlas.json` marker plus one `.af` file per array. -- atlas-rust 0.16.4 stores a collection as one immutable file. The file is - `/data.atlas`. An optional `/deleted.mask` sits beside it. -- Every atlas call this crate makes is gone in 0.16: `Atlas::open_dataset`, - `Atlas::merged_schema`, `Atlas::pruning_index`, `StoreConfig`, - `MergedSchema`, `ColumnKey`, `StatVal`, `PruningIndex`, `ArraySchema::codec`, - `DatasetSchema::array_attrs` and `DatasetSchema::global_attrs`. -- The workspace excludes the crate since #371 (commit `cebecfda`). Core does - not register it. `STORED AS ATLAS` and `read_atlas` fail today. -- The old crate predates the morsel scan. It dealt dataset names round-robin at - plan time and streamed flat batches outside the nd spine. NetCDF, HDF5, Zarr - and TIFF now read through `MorselSource`, `FileRead`, `NdSourceExec` and - `NdBroadcastExec`. - -Decision: rewrite the crate from scratch. Keep the crate name and path. Keep -three ideas from the old crate: the reader cache, the `atlas_*` metrics and -the per-table `OPTIONS`. - -## 2. The new format, checked against the 0.16.4 source - -Source: `~/.cargo/registry/src/*/atlas-rust-0.16.4/`. - -| Fact | Where | -| --- | --- | -| Container: `ATLS` header, segments back to back, `zstd(msgpack(footer))`, 16-byte trailer. | `src/format/mod.rs` | -| One dataset is one segment. A segment is a complete array-format 0.12.0 file. `SegmentStore` presents it as `seg{ordinal}.af`. | `src/format/segment_store.rs` | -| An open costs one `HEAD`, one 64 KiB tail read, one more read when the footer is larger, and one `GET` of `deleted.mask`. A missing mask is fine. | `src/reader/mod.rs:57-158` | -| The footer holds every dataset name, segment range, schema, attribute value and per-array statistic. Metadata costs no I/O after the open. | `src/format/footer.rs` | -| A schema holds `dtype`, `shape`, `chunk_shape`, `dimension_names` and `fill_value` per array. Datasets with equal schemas share one pool entry. | `src/schema/array.rs` | -| Statistics per array: `min`, `max`, `null_count`, `row_count`. `null_count` counts the elements equal to the fill value. A never-written array has no entry. Lists have no `min` or `max`. | `src/format/footer.rs:305-320` | -| Attributes sit at dataset scope and at array scope. A value is a scalar, a list, or `TimestampNanoseconds`. | `src/schema/attr.rs` | -| Dtypes: `Bool`, `Int8..Int64`, `UInt8..UInt64`, `Float32`, `Float64`, `String`, `Binary`, `TimestampNs`, `List`, `FixedSizeList`. | `array-format/src/dtype.rs` | -| `read_array::` needs `T: ArrayElement`. That is the numeric types, `String`, `Vec` and `TimestampNs`. There is no `bool` and no list element. | `array-format/src/array.rs` | -| `read_array` checks the dtype and the shape against the footer, opens the segment once per handle, and fetches only the chunks the region overlaps. Unwritten cells come from the fill value. | `src/reader/mod.rs:521-591` | -| Atlas merges no schema. Two datasets can declare one array name with two dtypes. | `README.md`, "Types" | -| The collection is immutable. A delete writes ordinals to `deleted.mask`. `list_datasets()` hides them. `dataset(name)` refuses them. | `src/format/mask.rs` | -| `Atlas::dataset(name)` scans the footer linearly. There is no lookup by ordinal. | `src/reader/mod.rs:394-402` | -| Each `Atlas` handle owns one `DeltaCache`: a 256 MiB block budget and a 64 MiB I/O budget. | `src/reader/mod.rs:87-90` | -| `Atlas` and `DatasetView` are `Send + Sync` and implement `Debug`. | `src/lib.rs`, `src/reader/mod.rs` | - -The Python ingest (`atlas create`) shapes the data Beacon meets in practice: - -- One dataset per NetCDF file. The dataset name is the file name. -- xarray opens the file with `mask_and_scale=True`. Scale and offset are - applied. Time is `datetime64[ns]`, so it lands as `TimestampNs`. -- `_FillValue` becomes the fill value. The defaults are `NaN` for floats, - `NaT` (`i64::MIN`) for timestamps, `""` for strings and none for integers. -- Attributes such as `units` and `calendar` stay as plain attributes. -- Python adds `_pyatlas_coords` (a JSON string), `_pyatlas_timedelta` and - `json:`-prefixed string attributes. Beacon shows them as strings. - -Dependency facts: - -- atlas-rust 0.16.4 was published on 2026-09-01. That is today. The repo rule - says: do not take a package that is less than one week old. This is the - user's own crate, so the rule is a prompt to confirm, not a block. -- The local checkout `~/git/atlas` is at 0.15.0. It lacks the 0.16 commits. - Fetch before any upstream work. -- Every transitive dependency fits the workspace: `object_store 0.13`, - `ndarray 0.17`, `rkyv 0.8` (the lock pins 0.8.10), `moka 0.12`, - `lz4_flex 0.11`, `zerocopy 0.8`, `zstd 0.13`, `rmp-serde 1`, `tempfile 3`. - MSRV 1.85 sits under the workspace floor of 1.94. - -## 3. Design - -### 3.1 One dataset is one morsel - -```text -CREATE EXTERNAL TABLE t STORED AS ATLAS LOCATION 'obs/**/data.atlas' - -listing obs/a/data.atlas obs/b/data.atlas markers - │ │ -create_physical_plan │ open (cached) │ list_datasets() - ▼ ▼ -entries [a#d0 a#d1 ... a#dN] [b#d0 ... b#dM] one PartitionedFile per dataset - │ -repartitioned MorselSource queue ──▶ one standing entry per partition - │ -OpenFile::open entry ─▶ Atlas (cached) ─▶ first open of a collection: build its pruning index - │ ─▶ index row of the entry: kept, or skipped - │ ─▶ DatasetView ─▶ AnyDataset (lazy backends) - ▼ ─▶ FileRead::plan (chunk grid, predicate masks) -workers pop chunk ─▶ view.read_array(start, shape) ─▶ nd-encoded batch - │ -plan DataSourceExec ─▶ NdSourceExec ─▶ NdBroadcastExec -``` - -Level 1 of the morsel scan is a dataset. Level 2 is a stored chunk of it. The -backend reports the array's `chunk_shape`, so `FileRead` cuts the dataset on -the chunk grid the writer chose. One pop then reads one stored chunk per -projected array. array-format packs chunks into 8 MiB blocks and caches the -decompressed block, so arrays that share a block cost one fetch. - -### 3.2 Discovery - -- The marker is `data.atlas`. The collection prefix is its parent directory. -- The factory `get_ext()` returns `atlas`. The format `get_ext()` returns - `atlas` too. The listing filter matches by suffix, so `data.atlas` matches - and `deleted.mask` does not. `STORED AS ATLAS` already falls back to the - glob `*.atlas` (see `listing_table_factory_ext.rs`). -- `is_atlas_marker(obj)`: the location is `data.atlas` or ends in - `/data.atlas`. Another `*.atlas` name is skipped with a debug log, because - `Atlas::open` hardcodes the object name. A rename adapter is out of scope. -- `top_level_atlas_markers` keeps one marker per directory. -- `discover_datasets` emits one `DatasetMetadata` per marker with format - `atlas`. The crawler rule "extension equals format" now admits - `data.atlas`, so a crawler can build an atlas table. Update its test. -- `schema_units` uses `units_over_stores`. `schema_options_fingerprint` - returns `SchemaOptions::new("atlas").finish()` when no `read_dimensions` is - set, and `None` otherwise. That matches Zarr. -- Only the 0.16 format is read. A pre-0.16 collection is a directory of `.af` - files with an `atlas.json` registry; nothing here recognizes one, so its - marker is simply not a marker and a listing passes over it. There is no - compatibility path and no migration hint. - -### 3.3 Open and cache - -- `open_collection(store, marker) -> Atlas` calls `Atlas::open(store, prefix)`. -- `AtlasReaderCache` is a `moka::future::Cache>` sized by - `reader_cache_size`. The key is the marker path, its `last_modified`, its - `size`, and a mask stamp. -- The mask stamp is one `HEAD` of `/deleted.mask`: `None` when absent, - else `(last_modified, size, e_tag)`. The container never changes, so this - stamp is the only thing that can retire a cached handle. One `HEAD` per - open-through-cache is the cost. A query pays one or two. -- `get_or_open_atlas(cache: Option<&AtlasReaderCache>, store, marker)` opens - directly when the cache is `None`. -- Memory bound: each handle owns 320 MiB of cache budget. A cache of 32 - handles can hold 10 GiB. Keep the default at 32 and document the bound. See - section 10 for the upstream fix. - -### 3.4 Dataset to `AnyDataset` - -Build one `Arc` per dataset and share it across its backends. -That avoids one linear name lookup per array read. - -Column model. Use the netCDF and Zarr convention: - -| Atlas | Column | -| --- | --- | -| array `a` | `a` | -| array attribute `k` of `a` | `a.k` | -| dataset attribute `k` | `.k` | - -The old crate used the bare key for a dataset attribute. The new convention -avoids a collision between an attribute and an array of one name. Ragged -detection reads `a.sample_dimension`, which is unchanged. - -Type mapping. Put it in `compat.rs` and pin it with tests: - -| Atlas dtype | `NdArrayDataType` | Note | -| --- | --- | --- | -| `Int8..Int64`, `UInt8..UInt64` | the same width | | -| `Float32`, `Float64` | `F32`, `F64` | | -| `String` | `String` | | -| `Binary` | `Binary` | | -| `TimestampNs` | `Timestamp` | `TimestampNanosecond` wraps the same `i64` | -| `Bool` array | skipped | array-format reads no `bool` | -| `List`, `FixedSizeList` | skipped | the nd model has no list | -| `Attr` scalar | rank-0 array | `Bool` attributes are kept | -| `Attr` list | skipped | | - -Log a skip at `debug`, not `warn`. A collection of a million datasets would -flood the log. - -Fill values. `view.array_fill_value(a)` gives a `FillValue`. Convert it with -`::fill_element` and report it from the backend. The nd -engine nulls the cells equal to it. Three consequences to document: - -- A `NaT` timestamp fill reads as null. -- A `""` string fill reads as null. An empty string in the data reads as null - too. That mirrors the Python ingest, which cannot store a null string. -- A `NaN` float fill leaves `NaN` cells as `NaN`, because `NaN != NaN`. That - is the engine's rule for every format. - -No CF decoding. The Python ingest applies scale and offset and decodes time -before the write. The format has a native timestamp type. Beacon therefore -reads every atlas array as stored. Update `cf-decoding.md`, which says the -opposite today. A Rust-written collection with packed integers and CF `units` -reads as integers. Document that. - -Projection. `dataset_from_view(view, projected: Option<&[String]>)` builds a -backend only for a projected column. A wide dataset then pays nothing for the -columns a query never names. - -Dimensions. Apply `resolve_read_dimensions` and `DatasetProjection` as Zarr's -`project_read_dimensions` does. Every array has one dimension name per axis, -so nothing is invented. - -### 3.5 Array backend - -```rust -pub struct AtlasArrayBackend { - view: Arc, - array: String, - shape: Vec, - dimensions: Vec, - chunk_shape: Vec, - fill_value: Option, -} -``` - -- `read_subset(subset)` calls `view.read_array::(array, start, shape)` and - converts the `ArcArray` with `into_owned()`. -- `AtlasElement` bridges `NdArrayType` and `ArrayElement`. Numeric types, - `String` and `Vec` pass through. `TimestampNanosecond` maps element-wise - from `TimestampNs`. -- `chunk_shape()` returns the stored chunk shape. This is what aligns level 2 - of the morsel scan with the file. -- `AttributeBackend` holds one value at rank 0. Copy it from the Zarr crate. - -### 3.6 Schema inference - -Input: the markers of the listing. - -1. Open each collection through the cache. -2. For each live dataset, compute a dedupe key: the address of - `view.schema()` (the interned pool entry) plus the sorted list of - `(attribute key, Attr::dtype())` for the dataset and its arrays. -3. For each new key, build the lazy `AnyDataset`, apply the dimension - narrowing, and derive its Arrow schema with `any_dataset_to_arrow_schema`. - Label it with `marker#dataset`. -4. Merge every labeled schema with `session_widening(state).merge_schemas`. -5. Return an empty schema for a collection with no live dataset. - -A fleet of similar datasets then costs one schema per distinct shape, not one -per dataset. The pass over the footer stays O(datasets) but does no I/O. - -### 3.7 Physical plan - -`AtlasFormat::create_physical_plan`: - -1. `reject_partition_columns("Atlas", &conf)`. A dataset has no path, so no - partition value can come from one. -2. Collect the markers from the file groups. Reduce them to top-level markers. -3. For each marker, open the collection and list its datasets. Build one - `PartitionedFile` per dataset, in `list_datasets()` order. Keep the - marker's `ObjectMeta` verbatim, so the file-statistics pruner still - recognises the store. Put the dataset name and its position in that - listing in `extensions` as `AtlasEntry { dataset: String, position: usize }`. - The position is the row of the dataset in the pruning index (section 3.8). -4. Encode the file schema with `beacon_datafusion_ext::nd::encoded_schema`. -5. Build `AtlasSource` with the read dimensions, the pushed projection, the - cache and the pruning switch. Rebuild the config with the new groups. -6. Wrap the scan: `DataSourceExec` under `NdSourceExec` under - `NdBroadcastExec`. Copy Zarr's `nd_scan_plan`. - -`AtlasSource` implements `FileSource` as `ZarrSource` does: - -- `repartitioned` puts every entry in one `MorselSource` through - `morsel_scan`, unless the scan is ordered or has one partition. -- `try_pushdown_filters` folds the filters into one predicate. -- `try_pushdown_projection` merges projections. -- `create_file_opener` builds an `AtlasOpener` that holds the queue and an - `Arc`. - -`AtlasDatasets` implements `OpenFile`: - -1. Read the `AtlasEntry`. A missing entry is an internal error. -2. Open the collection through the cache. -3. Get the collection's `CandidateFilter` from the `PruneCache`. The first - open of a collection in a scan builds the pruning index (section 3.8). - Every later open reads it. A pruned entry returns `FileRead::skipped()` - and counts in `atlas_datasets_pruned`. -4. Build the `DatasetView`. -5. Build the projected `AnyDataset`, narrow the dimensions, and call - `FileRead::plan` with the projected schema, the batch size, the predicate, - `FilePartitions::none()` and the read metrics. - -`FileRead::skipped()` does not exist yet. Add it to `beacon-nd-array` as a -public constructor for "a file the scan decided not to read": no queue and -`Output::Nothing`. It is four lines. - -The single-partition path reads each entry through the same `open`, then -streams it, as Zarr's opener does. - -The EXPLAIN size of a partition is the marker size times the dataset count. -That is cosmetic. Document it. - -### 3.8 Pruning with a collection index - -A collection can hold millions of datasets. A predicate evaluation per -dataset would cost millions of evaluations. Build one index per collection -instead, and evaluate the predicate once, over every dataset in one -vectorised pass. - -**When.** The first `OpenFile::open` for an entry of a collection builds the -index. Every later open of that collection reads the result. A `PruneCache` -on `AtlasSource` memoises the result per marker path for the life of the -scan. Its `moka::future::Cache::get_with` coalesces the partitions: the first -one builds, the rest await the same future. Each partition's opener holds a -clone of the cache, and the clones share one store. - -**What.** `PruningIndex` holds, in `list_datasets()` order: - -- `names: Vec`. The row order, and the guard of section "How an - entry uses it". -- One `StatColumn` per referenced column: `min: ArrayRef`, `max: ArrayRef`, - `null_count: UInt64Array`, `row_count: UInt64Array`. Every array has length - N, one row per live dataset. - -**How it is built.** - -1. Take the pushed predicate and the logical projected schema. Derive that - schema from the encoded projected schema through - `nd::encoding::nd_value_type`, because the scan schema is nd-encoded. - Build one `PruningPredicate`. A predicate the engine refuses gives - `CandidateFilter::KeepAll`. -2. `collect_columns` names the referenced columns. Resolve each one: - - An array name. Call `atlas.array_stats_by_dataset(name)`. That is one - linear pass over the footer per column, with no view and no name lookup. - Align its `(dataset, stats)` pairs to the rows with a - `HashMap<&str, usize>` built once from `names`. A dataset without an - entry stays unknown in that row. - - An attribute, `.k` or `a.k`. The value is exact, so `min = max = value` - and `null_count = 0`. This prunes `WHERE ".platform" = 'X'` from the - footer alone. Atlas has no bulk attribute accessor. A value needs one - `DatasetView`, and `Atlas::dataset(name)` is a linear scan, so the pass - is quadratic. Build an attribute column only while N is at most - `ATTRIBUTE_INDEX_LIMIT`, 100 000. Above it, leave the column unknown - until the upstream lookup of section 10 lands. - - Anything else. Unknown. The column gets no `StatColumn`, and the - predicate cannot prune on it. -3. Pack each column into typed Arrow arrays in the table type. Use a typed - builder per target type: `Float64Builder`, `Int64Builder`, - `TimestampNanosecondBuilder`, `StringBuilder` and so on. Keep a - `ScalarValue::cast_to` fallback for a type the fast path lacks. Rules per - value: - - Cast the dataset's native dtype to the table type. A cast failure is - null. - - A `NaN` bound is null. `total_cmp` sorts `NaN` last, so a `NaN` max says - nothing about the other values. - - `Bytes` becomes `Utf8` for a `Utf8` column when it is valid UTF-8, and - `Binary` for a `Binary` column. Otherwise null. - - `TimestampNs` becomes `TimestampNanosecond`. - - A missing entry is null for `min` and `max`, and `None` for both counts. - Run the pack on `spawn_blocking`. A million rows is CPU work, not I/O. -4. `AtlasPruningStatistics` implements `PruningStatistics` over the index. - `num_containers` is N. Each accessor returns the column's array, or `None` - for a column the index lacks. -5. `pruning_predicate.prune(&stats)` returns one `bool` per row. Keep it as - `CandidateFilter::Rows { kept: BooleanBuffer, names }`. - -**How an entry uses it.** `AtlasEntry` carries `position`, the dataset's row -in the plan-time listing. `CandidateFilter::keeps(position, name)` reads the -bit at `position` when `names[position] == name`. A mismatch means the -listing changed between plan and open, so the entry is kept. `KeepAll` keeps -every entry. No string is hashed per entry. - -**Counts.** `row_count` is the element count of one array, not the row count -of the broadcast. The predicate uses the counts to decide "every value is -null" and "no value is null". Both hold per array before and after a -broadcast, so the counts are exact for that purpose. - -**Cost.** One footer pass per referenced column, one typed pack, one -vectorised evaluation. For a million datasets and one `Float64` column the -index holds 32 MB and builds in well under a second. A pruned entry then -costs one cache lookup and one empty `FileRead` at its pop. - -**Fail open.** Any error in the build gives `KeepAll`. A row the index cannot -judge stays in. The predicate runs again above the scan, so a kept dataset -that matches nothing costs a read and never a wrong row. - -The pushed predicate also reaches `FileRead::plan`, which prunes chunks on the -coordinate arrays. The two levels compose. - -**Later options, out of scope.** - -- Keep the packed `StatColumn`s on the reader-cache entry, keyed by column - name and table type. The collection is immutable, so a second query pays no - pack. -- When `repartitioned` already holds the predicate, build the index there and - queue only the candidates. That saves one pop and one prefetch task per - pruned dataset. -- Implement `contained()` for attribute columns. Their values are exact, so an - `IN` list prunes too. - -### 3.9 Statistics for the analyzer - -`FileFormat::infer_stats` folds the footer per marker, as Zarr's -`StoreRanges` does: - -- An array column: the lowest `min` and the highest `max` over the live - datasets that hold statistics for it, cast to the table type. Unknown when - any dataset's bound is missing, `NaN`, or fails to cast. -- An attribute column: the same fold over the values. -- A dataset without the column adds nothing. - -It costs no I/O beyond the open. Gate it like Zarr all the same: only -`create_for_analysis` enables it, and `enable_statistics` decides. - -### 3.10 Configuration and `OPTIONS` - -```rust -pub struct AtlasConfig { - pub use_reader_cache: bool, // true - pub reader_cache_size: u64, // 32 - pub use_pruning: bool, // true - pub enable_statistics: bool, // true -} -``` - -| `OPTIONS` key | Env | Effect | -| --- | --- | --- | -| `read_dimensions` | | The dimensions the table reads | -| `use_reader_cache` | `BEACON_ATLAS_USE_READER_CACHE` | Consult the reader cache | -| `use_pruning` | `BEACON_ATLAS_USE_PRUNING` | Prune datasets on footer statistics | -| `enable_statistics` | `BEACON_ATLAS_ENABLE_STATISTICS` | Let the analyzer measure a collection | - -Read the keys with `format_option()`. A key arrives `format.`-prefixed and -lowercased. Reject a bad boolean at `CREATE EXTERNAL TABLE` time. - -### 3.11 Table functions - -`read_atlas(glob_paths)` and `read_atlas(glob_paths, dimensions)`. A path -names a `data.atlas` or a glob such as `**/data.atlas`. The function builds -the format from the session factory with `read_dimensions`, then a -`FastObjectTable`. The `read_atlas_schema` wrapper comes for free. - -### 3.12 Metrics - -Keep `atlas_open_time`, `atlas_prune_time`, `atlas_dataset_build_time`, -`atlas_datasets_scanned` and `atlas_datasets_pruned`. Add -`atlas_index_builds`, the number of pruning indexes a scan built, and -`atlas_index_rows`, their row total. The partition that builds an index -records its build time in `atlas_prune_time`. Register one `ReadMetrics` per -partition, as the other nd formats do. - -### 3.13 Behaviour changes versus the old crate - -1. The marker is `data.atlas`, not `atlas.json`. Every `LOCATION` changes. -2. A dataset attribute column is `.k`, not `k`. -3. A dataset that lacks every projected column contributes no rows. The old - crate null-filled its rows. NetCDF and Zarr already behave this way. -4. The scan goes through the nd spine. `NdBroadcastExec` sits above it, and - the nd projection pushdown rule applies. -5. A partitioned atlas table is refused with a clear error. -6. The pruning index is built from the footer at the first open of a scan. - Nothing is persisted, and nothing is read from disk to build it. -7. A pre-0.16 collection is not read at all. Its `atlas.json` is not a marker, - so a listing passes over it rather than failing a query. -8. A column two datasets type in two families — `String` in one and `Int64` in - another — fails schema inference, and the error names both datasets. The old - crate took atlas's own merge, which made every such column text. Beacon now - settles it the way it settles two files of any other format, and - `BEACON_TYPE_WIDENING_ON_CONFLICT=keep_first` takes the first dataset's type - instead. See section 11. - -## 4. Crate layout - -```text -beacon-arrow-atlas/ - Cargo.toml atlas-rust = "=0.16.4" via the workspace - src/lib.rs crate docs, module list, re-export of `atlas` - src/config.rs AtlasConfig - src/store.rs markers, prefix, open, AtlasReaderCache - src/compat.rs dtype, Attr and FillValue mapping; column names - src/backend.rs AtlasArrayBackend, AttributeBackend, AtlasElement - src/reader.rs dataset_from_view, collection_schema, project_read_dimensions - src/datafusion/mod.rs AtlasFormatFactory, AtlasFormat, nd_scan_plan - src/datafusion/source.rs AtlasSource, AtlasOpener, AtlasDatasets, AtlasEntry - src/datafusion/pruning.rs PruningIndex, StatColumn, CandidateFilter, PruneCache - src/datafusion/statistics.rs the infer_stats fold - src/datafusion/options.rs AtlasOptions - src/datafusion/metrics.rs AtlasScanMetrics - src/datafusion/table_function.rs ReadAtlasFunc - src/test_support.rs #[cfg(test)] fixtures built with AtlasWriter -``` - -Every test lives beside the code it covers, including the end-to-end ones in -`datafusion/mod.rs`. The fixtures are `#[cfg(test)]`, so an integration target -under `tests/` could not reach them. - -Delete every file of the old `src/` first. Nothing in it compiles against -0.16. - -## 5. Work plan - -Environment for every step: - -```bash -export PATH="$HOME/.cargo/bin:$PATH" -source ~/.config/beacon/build-env.sh -``` - -The active toolchain is stable 1.98. CI also builds at 1.94. Use no feature -newer than 1.94. `cargo fmt --check` is not clean repo-wide, so format the -new crate alone. - -### Phase 0: prerequisites - -1. Confirm the dependency rule for atlas-rust 0.16.4 (published today). -2. Add `atlas-rust = "=0.16.4"` to `[workspace.dependencies]`. -3. Remove the `exclude` line from the workspace `Cargo.toml`. Add the crate - to `members`. -4. Run `cargo tree -p beacon-arrow-atlas -e normal | head` after step 5 of - phase 1 to confirm one `rkyv`, one `object_store` and one `ndarray`. - -### Phase 1: the reader - -1. Write `config.rs`, `store.rs`, `compat.rs`, `backend.rs`, `reader.rs`. -2. Write `test_support.rs`. Build collections in a `tempdir` with - `AtlasWriter`: two datasets with attributes and a fill; a widening pair - (`Int16` and `Float32`); an incompatible pair (`String` and `Int64`); a - ranged fleet of `n` datasets; a chunked 2-D grid; a dataset with a list - attribute; an empty collection. A `Bool` or list *array* cannot be a - fixture: no Rust writer can produce one, so that mapping is unit-tested - alone. A deleted dataset is made in the test that wants one, by calling - `delete_dataset` on the open collection. -3. Unit tests: marker recognition, prefix, cache hit and miss on the mask - stamp, every dtype mapping, fill conversion, column names, a full read, a - window read that spans chunks, the timestamp path, the skips. -4. `cargo test -p beacon-arrow-atlas`. - -### Phase 2: the scan - -1. Add `FileRead::skipped()` to `beacon-nd-array`. -2. Write `datafusion/mod.rs`, `source.rs`, `options.rs`, `metrics.rs`, - `table_function.rs`. -3. End-to-end tests in `datafusion/mod.rs`, through `ListingTable` and - `FastObjectTable`: every row once at 1, 4 and 8 partitions; `COUNT(*)`; - projection; the widening cast; the null fill of a missing column; the - incompatible pair as `Utf8`; the deleted dataset absent; `read_dimensions` - narrows the schema; the plan shape `NdBroadcastExec` over `NdSourceExec` - over `DataSourceExec`; a chunk-pruned scan reads fewer encoded batches; - `EXPLAIN` does not open a segment. -4. `cargo test -p beacon-arrow-atlas` and - `cargo test -p beacon-nd-array --lib`. - -### Phase 3: pruning and statistics - -1. Write `pruning.rs` and `statistics.rs`. -2. Tests: the index over the ranged fleet has one row per live dataset in - listing order; `> 45` keeps `d5..d9`; an impossible predicate prunes - everything; a permissive one keeps everything; the mixed-dtype pair casts - before it compares; an attribute predicate prunes from the footer; a `NaN` - bound is null and fails open; an unknown column fails open; a deleted - dataset has no row; a position whose name differs is kept; eight - partitions build the index once (`atlas_index_builds` is 1); results match - with pruning on and off; the metrics report the counts; a synthetic index - of 200 000 rows builds and prunes in one test without a timeout; - `infer_stats` folds a fleet and goes unknown on a mixed dtype. -3. `cargo test -p beacon-arrow-atlas`. - -### Phase 4: wiring - -Section 6 lists the files. Then: - -```bash -cargo clippy --workspace --lib --bins --tests -cargo test --workspace --no-fail-fast --lib --bins --tests -cargo fmt -p beacon-arrow-atlas -``` - -`beacon-datafusion-ext` does not test standalone. Use the workspace run. - -### Phase 5: integration and docs - -1. `integration-tests/formats/test_atlas.py`. It needs `atlas-python`. Add it - to `requirements-optional.txt` and skip when absent. Build a collection - from `test_file.nc` with `atlas.create`, query it, create an external - table, restart, check the table survives. -2. Update the docs of section 8. -3. Add a CHANGELOG entry. - -## 6. Wiring outside the crate - -| File | Change | -| --- | --- | -| `Cargo.toml` (workspace) | Drop the `exclude`. Add the member. Add `atlas-rust = "=0.16.4"`. | -| `beacon-db/beacon-core/Cargo.toml` | Add the crate. | -| `beacon-db/beacon-core/src/runtime_builder.rs` | `pub atlas: AtlasConfig`, `with_atlas_config`, and `AtlasFormatFactory::new(AtlasOptions::default(), builder.atlas.clone())` in `register_file_formats`. | -| `beacon-db/beacon-core/src/crawler/discovery.rs` | Update the marker test: `d/x/data.atlas` is crawlable. | -| `beacon-db/beacon-core/tests/schema_functions.rs` | Add `read_atlas_schema`. | -| `beacon-db/beacon-functions/Cargo.toml`, `src/file_formats/mod.rs` | Register `ReadAtlasFunc`. | -| `beacon-db/beacon-file-formats/beacon-nd-array/src/arrow/file_read.rs` | Add `FileRead::skipped()`. | -| `beacon-db/beacon-db-py/src/connection.rs`, `python/beacondb/_beacondb.pyi` | Add `read_atlas` and `read_atlas_schema`. | -| `beacon-server/beacon-server-config/src/lib.rs` | Re-export `AtlasConfig`. Add the four `BEACON_ATLAS_*` fields. Fill `atlas`. | -| `beacon-server/beacon-server/src/server/mod.rs` | `.with_atlas_config(config.atlas.clone())`. | -| `beacon-server/beacon-server/src/server/catalog.rs` | `"atlas" => "read_atlas"`. | -| `beacon-server/beacon-server/src/main.rs` | Add `atlas` and `array_format` to the quiet log list. | -| `beacon-clients/beacon-web/src/components/external-table-dialog.tsx` | Hint: "the `data.atlas` file". | -| `beacon-clients/beacon-web/src/pages/crawlers.tsx` | Add `{ value: "atlas", label: "Atlas" }`. | -| `integration-tests/formats/test_atlas.py`, `requirements-optional.txt` | New suite. | -| `CHANGELOG.md` | Entry. | - -`beacon-file-stats` and `fast_object` mention Atlas in comments only. They -need no code change. - -## 7. Tests to keep from the old crate - -Port these assertions. Rewrite the fixtures with `AtlasWriter`. - -- `reads_all_datasets_through_datafusion` and the `FastObjectTable` twin. -- `widened_array_dtype_is_cast_from_each_dataset`. -- `missing_column_is_null_filled_per_dataset`. -- The incompatible-dtype case, with its answer corrected: the merge is - refused and the error names both datasets, and `keep_first` resolves it to the - first dataset's type. See section 11. -- `pruning_matches_unpruned_results`, `pruning_on_mixed_dtype_column_end_to_end`, - `pruning_across_many_partitions_is_correct`, - `scan_metrics_report_pruned_and_scanned_counts`, and the `pack_column` - tests of the old `pruning.rs`. The old index came from atlas; the new one is - built here, but the pack and the `PruningStatistics` adapter are the same - shape. -- `partitioned_scan_reads_every_dataset_row`. -- `cache_returns_same_arc_for_identical_marker` and - `cache_reopens_when_last_modified_changes`. Add a mask-change case. -- `discover_datasets_emits_one_entry_per_store`. - -## 8. Docs - -Rewrite `docs/docs/2.0.0-rc5/formats/atlas.md`: the `data.atlas` file, the -mask, the column model, the `OPTIONS` table, footer pruning, the fill rules, -the skips, and "not readable: 0.14 collections". Then update every page that -names `atlas.json`: - -- `formats/index.md` (two rows and the marker note) -- `data-sources/external-tables.md#atlas` -- `server/datasets.md`, `server/configuration.md`, `server/performance-tuning.md` -- `server/crawlers.md` (atlas is crawlable now) -- `cf-decoding.md` ("Zarr and Atlas": Atlas decodes nothing) -- `sql/table-functions.md`, `sql/table-functions-utility.md` -- `guides/speed-up-queries.md`, `guides/query-a-collection.md`, `guides/query-s3.md` -- `faq.md`, `how-it-works.md` - -## 9. Risks and open points - -1. Freshness. atlas-rust 0.16.4 is one day old. A 0.16.5 with an API change - would land on this crate first. Pin exactly and accept. -2. Name lookup. `Atlas::dataset(name)` is O(datasets). A full scan of a - million datasets makes a million lookups. That is quadratic. The pruning - index avoids it for array columns, because `array_stats_by_dataset` is one - linear pass. A dataset that survives pruning still pays one lookup at its - open, and an attribute column in the index pays one per dataset. The - rebuild therefore works as is for a selective query over a large - collection, and for any query over a collection up to the tens of - thousands. Above that the upstream lookup of section 10 is required. -3. Memory. Every cached handle owns 320 MiB of cache budget. Document the - bound. The upstream shared cache removes it. -4. Mask freshness. The reader cache pays one `HEAD` per open. The schema - cache keys on the listed objects, and a `.atlas` listing omits the mask. A - delete can therefore leave a stale merged schema until the next listing - change. The stale schema can only hold an extra column, which reads as - null. Accept. -5. Chunk reads. `read_array` walks every chunk coordinate of the array per - call to find the overlap. A long array read chunk by chunk pays - O(chunks²) comparisons. Acceptable for now. See section 10. -6. Pruning schema. `PruningPredicate` needs the logical column types, and the - scan's projected schema is nd-encoded. Derive the logical schema through - `nd_value_type` and pin it with a test that a pushed `>` prunes. -7. Bool arrays and list values are invisible. The Python ingest refuses bool - by default too. Document. -8. `EXPLAIN` sizes. See section 3.7. - -## 10. Requests to atlas-rust - -Not blockers. Each one lifts a limit above. - -1. `Atlas::dataset_at(ordinal)` and a `HashMap` name index built - at open. Removes the quadratic scan, and lets attribute columns join the - pruning index at any size. -2. `Atlas::attribute_by_dataset(key)` and `Atlas::array_attribute_by_dataset(array, key)`, - the attribute twins of `array_stats_by_dataset`. One footer pass per - attribute column, with no view at all. -3. `Atlas::open_with_cache(store, prefix, Arc)`. Lets Beacon share - one block budget across every open collection. -4. Expose the schema pool, or `Atlas::array_dtypes()`: per array name, the - set of dtypes the live datasets declare. Makes schema inference - O(pool + attributes). -5. In array-format `assemble_nd`, iterate the chunk coordinates that overlap - the slice, not every coordinate of the array. - -## 11. What building it settled - -Phases 0 and 1 answered five questions this plan had guessed at. Each is now -pinned by a test. - -**1. Two families of type are refused, not stringified.** The old crate took -atlas's own merge, where `String` absorbed everything, so a collection whose -datasets typed one array as `String` and `Int64` read back as text. The -schema of a collection now merges through the session's -`ArrowTypeWidening`, exactly as the files of every other format do, and its -default refuses that pair: - -```text -Incompatible types for field 'value': Utf8 in 'sensor#a' vs Int64 in 'sensor#b' -``` - -The label is `{collection}#{dataset}`, so the offending dataset is named -rather than searched for. `BEACON_TYPE_WIDENING_ON_CONFLICT=keep_first` takes -the first dataset's type instead and casts the rest to it. This is a -behaviour change for a collection that holds such a column, and section 3.13 -records it. - -**2. An integer beside a `Float32` widens to `Float64`.** Not to `Float32`, as -the old crate's test asserted. A `Float32` holds no `Int32`, so a rule that -kept it would make the merged type depend on which dataset came first. That is -issue #377, and the session rule is what settles it. - -**3. `list_datasets()` reports write order.** Not sorted order. The pruning -index of section 3.8 keys its rows on that order, and `AtlasEntry::position` -indexes into it, so both are read from the same call at plan time. - -**4. A `Bool` or list array cannot be a fixture.** `array-format` implements -no element type for either, so no Rust writer can produce one and no -collection can hold one. The refusal is unit-tested at the mapping instead. A -list *attribute* is writable, and the reader drops it as designed. - -**5. The marker list must sort by directory, not by path.** A path sort puts -`a/b/data.atlas` before `a/data.atlas`, because `b` sorts under `d`, so a -nested collection would be kept and its parent dropped. The old crate sorted -by path and was correct only by accident: its marker was `atlas.json`, and -`a/atlas.json` sorts before `a/b/atlas.json`. Renaming the marker broke the -assumption. `top_level_atlas_markers` now sorts on the directory. - -Two Phase 0 results worth keeping: - -- The dependency graph is clean. `atlas-rust 0.16.4` brings - `array-format 0.12.0`, and the subtree resolves to one `rkyv 0.8.10` (shared - with `beacon-file-stats`), one `object_store 0.13.2`, one `ndarray 0.17.2` - and one `zstd 0.13.3`. Nothing was duplicated by adding the crate. -- The crate compiles clean at the workspace toolchain with no warnings of its - own, and `cargo check --workspace --lib` still passes. - -## 12. What the scan settled - -Phase 2 answered four more questions, all of them about the DataFusion -surface rather than about atlas. - -**1. The schema-adapter hooks are gone.** DataFusion 53 deprecates -`SchemaAdapterFactory` and gives `FileSource::with_schema_adapter_factory` -and `schema_adapter_factory` defaults that refuse and return `None`. The old -crate implemented both. The nd read adapts its own batches through -`batch_adapter_factory` inside `FileRead::plan`, so the source leaves the -deprecated pair alone and carries no adapter field. - -**2. A `COUNT(*)` still has to read one array.** The projection reaches the -dataset build, so projecting nothing would build nothing, and the read would -have no grid to take a row count from — `FileRead` would then index into an -empty dataset and fail. The opener picks the widest readable array out of the -footer instead and builds that one column. A dataset with no readable array -at all falls back to building its attributes, and contributes the single row -its scalars define. - -**3. A predicate column is built even when the projection omits it.** The -filter that stays above the scan forces its columns into the projection -today, so this changes nothing now. But the chunk pruning inside the read -matches columns by name, and it would stop pruning silently if that ever -stopped holding. - -**4. Partition columns reach a scan through the source's `TableSchema`.** -`FileScanConfigBuilder` has no `with_table_partition_cols`; the config reads -them off the source. Only the refusal test needed to know this — an Atlas -dataset lives inside a container rather than at a path, so -`reject_partition_columns` turns such a table away before a scan is built. - -The legacy format is gone from the crate entirely, by request: there is no -`atlas.json` constant, no detection, and no migration hint. A pre-0.16 -collection's marker is simply not a marker, so a listing passes over it. - -## 13. What pruning settled - -**1. A scan is offered its filters even when it has none.** DataFusion calls -`try_pushdown_filters` with an empty list, and `conjunction` over nothing is -the literal `true`. A source that stores that becomes a source with a -predicate, and this one would then build a pruning index for every scan, -including `SELECT *`. The fix is to treat a predicate that names no column as -no predicate. The other nd formats store the same literal; it costs them -nothing, because their chunk pruning finds no column ranges in it and stops. - -**2. Statistics need a count guard, and pruning does not.** A dataset that -declares an array and never writes it has no footer entry, and its cells read -back as the array's fill — or as zeros, when it declares none. In the pruning -index that dataset gets null bounds, and DataFusion keeps a container whose -statistics are null, so the answer stays right. In the statistics *fold* the -same dataset would simply be skipped, and the range would then exclude values -the collection holds. So `infer_stats` claims a bound only when every live -dataset reported one. A uniform collection — what `atlas create` writes — is -unaffected; a heterogeneous one reports unknown and is read in full. - -**3. The index is what makes attribute pruning affordable at all.** An array -column costs one `array_stats_by_dataset` pass. An attribute has no bulk -accessor, so it costs one `DatasetView` per dataset and each is a linear -scan. That is capped by `ATTRIBUTE_INDEX_LIMIT`, and it is the strongest case -for `Atlas::attribute_by_dataset` upstream (section 10, item 2). - -**4. `infer_stats` measures nothing during a query.** It follows Zarr: a -format built by `create()` reports unknown whatever the option says, and only -`create_for_analysis` turns the fold on. Folding a footer is cheap, but a -listing of thousands of collections would still open every one of them while -planning. - -Two costs are deliberate and worth revisiting if a profile asks: - -- The pack goes through `ScalarValue` per value rather than a typed builder. - It runs on a blocking thread, once per collection per scan. A typed builder - per target type is the optimization, and the packing is one function. -- The index is not kept between queries. The collection is immutable, so the - packed columns could live on the reader-cache entry, keyed by column and - target type. - -## 14. What the wiring settled - -**1. An Atlas collection is crawlable, and Zarr is not.** The crawler's rule is -that a file's extension must equal its format name. A Zarr store is a -directory behind a `zarr.json`, so `json != zarr` and it is skipped. A -collection is one file named `data.atlas`, so `atlas == atlas` and a crawler -builds a table over it. Tables group by directory and each collection has its -own, so a crawl of many collections makes one table each; several in one table -is what an external table over a glob is for. The discovery test that asserted -atlas was skipped is now the test that asserts it is not. - -**2. `beacon-core` needed `ndarray` as a dev-dependency.** Its integration -tests write real collections, and the atlas writer takes `ndarray` views. -Beacon itself never writes a collection, so the dependency belongs in -`dev-dependencies` alone. - -**3. The four settings reach a table two ways.** `BEACON_ATLAS_USE_READER_CACHE`, -`BEACON_ATLAS_READER_CACHE_SIZE`, `BEACON_ATLAS_USE_PRUNING` and -`BEACON_ATLAS_ENABLE_STATISTICS` set the runtime defaults, and every one but -the cache size is overridable per table through `OPTIONS`. An embedded caller -sets them with `RuntimeBuilder::with_atlas_config`. - -### What phase 4 verified, and what it did not - -Verified: `beacon-arrow-atlas` (105 tests), the whole of `beacon-core` -(45 suites, including 8 new end-to-end tests over a real runtime — the table -function, its `_schema` counterpart, a glob, `STORED AS ATLAS`, recovery -across a restart, the dimensions argument, and pruning against an unpruned -control). `cargo check` passes for `beacon-functions`, `beacon-server-config` -and `beacon-server` (lib and bins). Clippy is clean in every touched crate's -own sources. - -Not run: the test suites of `beacon-server-config` and `beacon-server`, and -any build of the web client. The changes there are a config struct, three -one-line wirings and two string edits, but they are unproven. - -### A note on `cargo fmt` - -Do not run `cargo fmt -p ` on the crates this touches. The repository is -not format-clean, so formatting a whole package rewrites files the change never -went near — one pass here reformatted 62 unrelated files across `beacon-core` -and `beacon-functions`, and they had to be reverted one by one. Format the new -crate, whose every file is new, and leave the rest alone. - -## 15. What the ingest settled - -The one thing no Rust test could answer: does the reader handle what `atlas -create` actually writes? It was checked directly — a collection built by -atlas-python 0.16.4 from five netCDF files, opened with this crate. - -**1. Real collections carry the statistics pruning needs.** Each dataset -reported its own `min`, `max`, `null_count` and `row_count` for every array, -with the ranges the source files held. Dataset pruning therefore works on -collections built the normal way, which was the open question behind the whole -design. - -**2. `atlas create` writes `{destination}/data.atlas`.** One dataset per source -file, named after the file *with its suffix* — `d0.nc`, not `d0`. A `LOCATION` -names the container inside that directory. - -**3. xarray's conventions arrive intact and are handled.** A float array gets a -`NaN` fill, because xarray reads with `mask_and_scale=True`; the reader carries -it through, and `NaN` never equals itself, so such a cell reads as `NaN` rather -than as null. Python's own marker attribute lands as the column -`._pyatlas_coords`, beside `.platform` and `temperature.units`. The collection -schema came out as `._pyatlas_coords`, `.platform`, `depth`, `temperature`, -`temperature.units` — exactly the column model this plan describes. - -**4. The reference writers live in `requirements.txt`.** Not -`requirements-optional.txt`, which this plan said: that file is for the Flight -SQL driver alone. The `formats/` convention is one pinned writer per format in -the main file, with each test skipping itself when its writer is absent. -`atlas-python==0.16.4` is pinned there. - -### What phase 5 verified, and what it did not - -Verified: the fixture half of `formats/test_atlas.py` runs against real -atlas-python; the reader handles its output; the test file imports, collects -and skips cleanly when `beacondb` is absent; the documentation site builds, -which is also a dead-link check because VitePress fails a build on one. - -Not run: the body of `formats/test_atlas.py`, which needs the `beacondb` -extension built with maturin. Its SQL and its use of the embedded API follow -`formats/test_zarr.py`, and the same ground is covered in Rust by -`beacon-core/tests/atlas.rs`. From ffb12cbdb14b07d6bf70aa8ab14fab55a11d0b96 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Fri, 4 Sep 2026 16:46:31 +0200 Subject: [PATCH 03/16] feat: read Atlas on the 0.17 container, one segment per variable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Atlas 0.17 changes where the bytes live. A container used to hold one segment per dataset, with the footer carrying every dataset's shapes, attribute values and statistics. It now holds one segment per *variable*: a segment holds one array name across the whole collection, and each dataset's copy sits inside it under the dataset's own name. The footer therefore names things and nothing else. An array's layout (shape, chunking, dimension names, fill value), its statistics, and every attribute value moved into the variable's segment, and reading one is async. One open answers for every dataset of the collection, so a column costs one request whether the collection holds ten datasets or a million. What that changes here: - The dataset build asks for a layout per array instead of reading it off the footer, and skips a dtype Beacon cannot surface before it opens anything. - The pruning index gathers a column with one call — `array_stats_by_dataset` for an array, `attributes_by_dataset` for an attribute — and pivots what comes back. That drops the 100k-dataset limit on indexing an attribute, which existed only because an attribute had needed a view per dataset. - Schema inference keys a dataset on its arrays, its attribute keys and its dimension names. The interned schema is coarser than it was — it names types and no longer implies a grid — so the names come from the segments to keep two datasets with different grids apart. - `COUNT(*)` picks its driving array by element count from the layouts. - An attribute can no longer be a timestamp: atlas stores none, because one would go to disk as an i64 and could not come back. Verified against the fixtures atlas-rust 0.17 ships, including the one its Python layer writes: schema, values, timestamps, string arrays and both attribute scopes all read back. --- CHANGELOG.md | 29 ++- Cargo.lock | 24 +- Cargo.toml | 2 +- .../beacon-arrow-atlas/src/compat.rs | 42 ++-- .../beacon-arrow-atlas/src/datafusion/mod.rs | 2 +- .../src/datafusion/pruning.rs | 227 +++++++++--------- .../src/datafusion/source.rs | 67 ++++-- .../src/datafusion/statistics.rs | 46 ++-- .../beacon-arrow-atlas/src/lib.rs | 32 ++- .../beacon-arrow-atlas/src/reader.rs | 194 +++++++++++---- docs/docs/2.0.0-rc5/formats/atlas.md | 49 ++-- integration-tests/formats/test_atlas.py | 12 +- integration-tests/requirements.txt | 2 +- 13 files changed, 458 insertions(+), 270 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ce4dfd17..bc7edf2a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,16 +12,19 @@ tag. Releases before 2.0.0 are recorded in the ### Added -- **Atlas is readable again, on the single-file format.** Atlas 0.16 replaced the directory of - per-array files with one write-once container, `data.atlas`, holding every dataset and a footer - that describes them all. Beacon's reader was written against the old layout and had been - excluded from the build, so `STORED AS ATLAS` and `read_atlas` failed. It is rebuilt on the new - format and registered again. A `LOCATION` now names the container rather than a marker beside +- **Atlas is readable again, on the single-file format.** Atlas replaced the directory of + per-array files with one write-once container, `data.atlas`: one segment per variable, holding + that array for every dataset, and a footer that names them all. Beacon's reader was written + against the old layout and had been excluded from the build, so `STORED AS ATLAS` and + `read_atlas` failed. It is rebuilt on the new format — container version 8, which Atlas 0.17 + writes — and registered again. A `LOCATION` now names the container rather than a marker beside it: `LOCATION 'obs/data.atlas'`, or a glob such as `'obs/**/data.atlas'`. A collection written - before 0.16 is not read at all — rewrite it with `atlas create`. Three behaviour changes come + before 0.17 is not read at all — rewrite it with `atlas create`. Three behaviour changes come with the rebuild. A dataset-level attribute is a column under a leading dot, `".platform"`, matching NetCDF and Zarr instead of the bare key. A scan reads through the shared nd pipeline, - so one dataset is one unit of work and a collection divides over every core. And a column two + so one dataset is one unit of work and a collection divides over every core — and a dataset + stored in several chunks divides further, so one large dataset still uses all of them. And a + column two datasets type in two families now refuses the merge by name rather than silently becoming text; `BEACON_TYPE_WIDENING_ON_CONFLICT=keep_first` settles it the other way. Atlas collections are also crawlable now, because a collection is one file whose extension is its format. @@ -29,13 +32,15 @@ tag. Releases before 2.0.0 are recorded in the and `BEACON_ATLAS_ENABLE_STATISTICS` configure it, and the same keys work per table through `OPTIONS`. -- **A predicate over an Atlas collection skips whole datasets.** The footer records the minimum, - the maximum and the null count of every array, so a collection can be judged before it is read. +- **A predicate over an Atlas collection skips whole datasets.** Atlas records the minimum, the + maximum and the null count of every array, so a collection can be judged before it is read. The first scan of a collection pivots those statistics into one index — one row per dataset, one typed Arrow column per column the predicate names — and evaluates the predicate over all of - them in a single vectorised pass. A collection of a million datasets therefore costs one pass - rather than a million decisions, and a dataset ruled out is never opened. A dataset-level - attribute is exact in the footer, so `WHERE ".platform" = 'p3'` prunes on it too. Pruning only + them in a single vectorised pass. Because a variable lives in one segment, gathering a column's + statistics is one request whatever the dataset count, so a collection of a million datasets + costs one request per column and one pass rather than a million decisions; a dataset ruled out + is never opened. A dataset-level attribute is exact, so `WHERE ".platform" = 'p3'` prunes on it + too, and at the same cost. Pruning only ever removes datasets that hold no matching row: every path falls back to reading everything, and the filter above the scan still decides each row. `EXPLAIN ANALYZE` reports it as `atlas_datasets_pruned`, `atlas_index_builds` and `atlas_index_rows`. diff --git a/Cargo.lock b/Cargo.lock index 4c119418..4753c51a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -289,11 +289,12 @@ dependencies = [ [[package]] name = "array-format" -version = "0.12.0" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7379c3303a9b9f0693ccd040d85da31e302a601efeafae8a6cb69341e42de41" +checksum = "672b40e04f7c403e9322972c476414f8847e5e0aed0006b80128977d39138ab9" dependencies = [ "bytes", + "ecow", "futures", "indexmap 2.14.0", "lz4_flex 0.11.6", @@ -1192,9 +1193,9 @@ dependencies = [ [[package]] name = "atlas-rust" -version = "0.16.4" +version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f7ba85f9c7eee0f991be08eefbbe9eef18fc9daf0992fde1a4215b6befd2fec" +checksum = "4dfadaa65a76f424f053c764cd439ab61c22ffade15ade54fe35432df2561e5b" dependencies = [ "array-format", "async-trait", @@ -1207,6 +1208,8 @@ dependencies = [ "parking_lot", "rmp-serde", "serde", + "smallvec", + "smol_str 0.3.6", "tempfile", "thiserror 2.0.20", "tokio", @@ -5136,6 +5139,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "ecow" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78e4f79b296fbaab6ce2e22d52cb4c7f010fe0ebe7a32e34fa25885fd797bd02" + [[package]] name = "either" version = "1.17.0" @@ -11143,6 +11152,9 @@ name = "smallvec" version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" +dependencies = [ + "serde", +] [[package]] name = "smol_str" @@ -11158,6 +11170,10 @@ name = "smol_str" version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4aaa7368fcf4852a4c2dd92df0cace6a71f2091ca0a23391ce7f3a31833f1523" +dependencies = [ + "borsh", + "serde_core", +] [[package]] name = "snafu" diff --git a/Cargo.toml b/Cargo.toml index 72d16def..c4dcaaeb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,7 +83,7 @@ object_store = { version = "0.13.1", features = ["aws", "gcp", "azure", "http"] # The atlas container embeds array-format files verbatim, so a change to the # crate's bytes is a change to the format Beacon reads. Pin it exactly. -atlas-rust = "=0.16.4" +atlas-rust = "=0.17.0" oxcdf = { git = "https://github.com/robinskil/oxcdf.git", version = "0.4.0", features = ["async", "object-store", "ndarray"] } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs index 65131d77..d8377bbb 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use arrow::datatypes::DataType; -use atlas::{ArraySchema, Attr, DType, DatasetView, FillValue}; +use atlas::{ArrayLayout, Attr, DType, DatasetView, FillValue}; use beacon_nd_array::{ NdArray, NdArrayD, datatypes::NdArrayDataType, datatypes::TimestampNanosecond, }; @@ -104,9 +104,10 @@ pub(crate) fn dtype_tag(dtype: &DType) -> String { /// Wrap one atlas array as a lazy [`NdArrayD`] over `view`. /// -/// Nothing is read here. The shape, the dimension names, the chunk shape and -/// the fill value all come from the collection footer, which the open already -/// held; the values arrive when the engine asks the backend for a subset. +/// No array data is read here. `dtype` comes from the collection footer, and +/// `layout` from the variable's segment, which one open serves for the whole +/// collection. The values themselves arrive when the engine asks the backend +/// for a subset. /// /// The chunk shape is the one the writer chose. It is what lets the scan cut a /// dataset on the grid the file actually stores, so one unit of work is one @@ -114,9 +115,10 @@ pub(crate) fn dtype_tag(dtype: &DType) -> String { pub fn array_to_nd_array( view: Arc, array_name: &str, - schema: &ArraySchema, + dtype: &DType, + layout: &ArrayLayout, ) -> anyhow::Result> { - let fill: Option = schema.fill_value.clone().map(Into::into); + let fill: Option = layout.fill_value().cloned(); macro_rules! lazy { ($ty:ty) => {{ @@ -126,16 +128,20 @@ pub fn array_to_nd_array( let backend = AtlasArrayBackend::<$ty>::new( view, array_name.to_string(), - schema.shape.clone(), - schema.dimension_names.clone(), - schema.chunk_shape.clone(), + layout.shape().to_vec(), + layout + .dimension_names() + .into_iter() + .map(str::to_string) + .collect(), + layout.chunk_shape().to_vec(), fill, ); Ok(Arc::new(NdArray::new_with_backend(backend)?) as Arc) }}; } - match &schema.dtype { + match dtype { DType::Int8 => lazy!(i8), DType::Int16 => lazy!(i16), DType::Int32 => lazy!(i32), @@ -190,7 +196,6 @@ pub fn attribute_to_nd_array(attr: &Attr) -> anyhow::Result> { Attr::Float64(v) => scalar!(*v), Attr::String(v) => scalar!(v.clone()), Attr::Binary(v) => scalar!(v.clone()), - Attr::TimestampNanoseconds(v) => scalar!(TimestampNanosecond(*v)), other => Err(anyhow::anyhow!( "attribute is a {} list, which has no rank-0 form in Beacon", dtype_tag(&other.dtype()) @@ -301,21 +306,6 @@ mod tests { assert_eq!(typed.clone_into_raw_vec().await, vec![2024]); } - #[tokio::test] - async fn a_timestamp_attribute_keeps_its_type() { - let nanos = 1_700_000_000_000_000_000; - let nd = attribute_to_nd_array(&Attr::TimestampNanoseconds(nanos)).unwrap(); - assert_eq!(nd.datatype(), NdArrayDataType::Timestamp); - let typed = nd - .as_any() - .downcast_ref::>() - .unwrap(); - assert_eq!( - typed.clone_into_raw_vec().await, - vec![TimestampNanosecond(nanos)] - ); - } - #[tokio::test] async fn a_bool_attribute_is_a_column() { let nd = attribute_to_nd_array(&Attr::Bool(true)).unwrap(); diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs index 3b7695ba..f03a2623 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs @@ -391,7 +391,7 @@ impl FileFormat for AtlasFormat { } match get_or_open_atlas(self.cache.as_ref(), Arc::clone(store), object).await { - Ok(atlas) => Ok(statistics::collection_statistics(&atlas, &table_schema)), + Ok(atlas) => Ok(statistics::collection_statistics(&atlas, &table_schema).await), Err(e) => { tracing::debug!(object = %object.location, "not measuring this collection: {e}"); Ok(Statistics::new_unknown(&table_schema)) diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs index b88aff90..1909e73a 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs @@ -3,9 +3,8 @@ //! # One index, not a decision per dataset //! //! A collection can hold millions of datasets. Evaluating a predicate against -//! each one in turn would cost millions of evaluations, and each would need its -//! own `DatasetView` — a linear scan of the footer — so the pass would be -//! quadratic before it did any work. +//! each one in turn would cost millions of evaluations, and each would open the +//! dataset to get its numbers. //! //! Instead the first opener that reaches a collection builds one //! [`PruningIndex`] over it: one row per live dataset, and one column of typed @@ -13,9 +12,11 @@ //! [`PruningPredicate`] then evaluates the whole collection in one vectorised //! pass, and the result is a bit per dataset that every partition reads. //! -//! The statistics come from the footer the open already held, so the build -//! costs no I/O. An array column costs one linear pass through -//! [`Atlas::array_stats_by_dataset`], with no view and no name lookup at all. +//! One column is one request. Atlas stores a variable in one segment, so +//! [`Atlas::array_stats_by_dataset`] and [`Atlas::attributes_by_dataset`] each +//! return every live dataset's value from a single open — however many datasets +//! there are. The index costs one open per column the predicate names, and no +//! array data at all. //! //! # Pruning is only ever an optimization //! @@ -29,24 +30,14 @@ use std::sync::Arc; use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use atlas::{Atlas, Attr, StatValue}; +use atlas::{ArrayStats, Atlas, Attr, StatValue}; use datafusion::common::Column; use datafusion::common::pruning::PruningStatistics; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::utils::collect_columns; use datafusion::physical_optimizer::pruning::PruningPredicate; use datafusion::scalar::ScalarValue; - -/// Above this many datasets, an attribute column stays out of the index. -/// -/// An array column costs one footer pass. An attribute has no bulk accessor in -/// the reader, so its value needs one `DatasetView` per dataset and each of -/// those is a linear scan — the pass is quadratic. It is worth paying on a -/// collection of thousands and not on one of millions, and the only cost of -/// skipping it is that a predicate on that attribute prunes nothing. -/// -/// `Atlas::attribute_by_dataset` upstream would remove the limit. -const ATTRIBUTE_INDEX_LIMIT: usize = 100_000; +use indexmap::IndexMap; // ─── What a scan does with the answer ──────────────────────────────────────── @@ -244,17 +235,46 @@ pub async fn candidate_filter( return CandidateFilter::KeepAll; } - // The pivot is pure CPU over data already in memory. A million rows is real - // work, so it does not run on the async runtime. - let atlas = Arc::clone(atlas); - let schema = Arc::clone(logical_schema); - let wanted: Vec = referenced + // Fetching is one request per column, whatever the dataset count. The + // pivot after it is pure CPU over what is then in memory, and a million + // rows is real work, so it does not run on the async runtime. + let wanted: Vec<(String, DataType)> = referenced .iter() - .map(|column| column.name().to_string()) + .filter_map(|column| { + let field = schema_field(logical_schema, column.name())?; + Some((column.name().to_string(), field)) + }) .collect(); + let mut fetched: Vec<(String, DataType, Measured)> = Vec::with_capacity(wanted.len()); + let arrays = atlas.list_arrays(); + for (column, target) in wanted { + let measured = if arrays.iter().any(|array| array == &column) { + match atlas.array_stats_by_dataset(&column).await { + Ok(stats) => Some(Measured::Array(stats)), + Err(e) => { + tracing::debug!(column, "atlas statistics unavailable for pruning: {e}"); + None + } + } + } else { + attribute_values(atlas, &column) + .await + .map(Measured::Attribute) + }; + if let Some(measured) = measured { + fetched.push((column, target, measured)); + } + } + + if fetched.is_empty() { + // Nothing the predicate names has statistics, so nothing can be ruled + // out. + return CandidateFilter::KeepAll; + } + let built = tokio::task::spawn_blocking(move || { - let index = build_index(&atlas, &names, &wanted, &schema); + let index = build_index(&names, fetched); (names, index) }) .await; @@ -262,11 +282,6 @@ pub async fn candidate_filter( let Ok((names, index)) = built else { return CandidateFilter::KeepAll; }; - if index.columns.is_empty() { - // Nothing the predicate names has statistics, so nothing can be ruled - // out. - return CandidateFilter::KeepAll; - } match pruning.prune(&index) { Ok(kept) => CandidateFilter::Rows { @@ -280,42 +295,72 @@ pub async fn candidate_filter( } } -/// Pivot the footer into one [`StatColumn`] per column that has statistics. -fn build_index( - atlas: &Atlas, - names: &[String], - wanted: &[String], - schema: &SchemaRef, -) -> PruningIndex { - // Where each dataset sits, so a footer pass in write order can be scattered - // into rows without a search. +/// The type a column carries in the scan's own schema, or `None` when it holds +/// no such column and nothing can be typed against it. +fn schema_field(schema: &SchemaRef, column: &str) -> Option { + schema + .field_with_name(column) + .ok() + .map(|field| field.data_type().clone()) +} + +/// What one request measured about one column, across every live dataset. +enum Measured { + /// An array's statistics. A dataset that wrote the array has an entry, and + /// [`ArrayStats::name`] names it. + Array(Vec), + /// An attribute's value, keyed by dataset. + Attribute(IndexMap), +} + +/// Every live dataset's value for one attribute column, or `None` when the +/// collection carries no such attribute. +/// +/// A column is `.key` at dataset scope, or `array.key` at array scope. An array +/// name and an attribute key may both hold dots, so every split of the latter +/// is a candidate and the first that finds a value wins. +async fn attribute_values(atlas: &Atlas, column: &str) -> Option> { + let found = |values: IndexMap| (!values.is_empty()).then_some(values); + + if let Some(key) = column.strip_prefix('.') { + return found(atlas.attributes_by_dataset(None, key).await.ok()?); + } + for (index, character) in column.char_indices() { + if character != '.' { + continue; + } + let (array, rest) = column.split_at(index); + if let Ok(values) = atlas.attributes_by_dataset(Some(array), &rest[1..]).await + && let Some(values) = found(values) + { + return Some(values); + } + } + None +} + +/// Pivot what was fetched into one [`StatColumn`] per column. +fn build_index(names: &[String], fetched: Vec<(String, DataType, Measured)>) -> PruningIndex { + // Where each dataset sits, so a pass in write order can be scattered into + // rows without a search. let row_of: HashMap<&str, usize> = names .iter() .enumerate() .map(|(row, name)| (name.as_str(), row)) .collect(); - let arrays = atlas.list_arrays(); let mut columns = HashMap::new(); - for column in wanted { - let Ok(field) = schema.field_with_name(column) else { - continue; - }; - let target = field.data_type(); - - let packed = if arrays.iter().any(|array| array == column) { - Some(pack_array_column( - atlas, - names.len(), - &row_of, - column, - target, - )) - } else { - pack_attribute_column(atlas, names, &row_of, column, target) + for (column, target, measured) in fetched { + let packed = match measured { + Measured::Array(stats) => { + Some(pack_array_column(&stats, names.len(), &row_of, &target)) + } + Measured::Attribute(values) => { + pack_attribute_column(&values, names.len(), &row_of, &target) + } }; if let Some(packed) = packed { - columns.insert(column.clone(), packed); + columns.insert(column, packed); } } @@ -325,17 +370,16 @@ fn build_index( } } -/// One array column, from one linear pass over the footer. +/// One array column. /// /// A dataset with no entry for the array keeps a null bound and unknown counts. /// That is what a dataset which does not declare the array looks like, and it /// is also what one that declared it and never wrote it looks like — both must /// stay in, and a null does exactly that. fn pack_array_column( - atlas: &Atlas, + stats: &[ArrayStats], rows: usize, row_of: &HashMap<&str, usize>, - column: &str, target: &DataType, ) -> StatColumn { let null = ScalarValue::try_from(target).unwrap_or(ScalarValue::Null); @@ -344,14 +388,15 @@ fn pack_array_column( let mut null_counts: Vec> = vec![None; rows]; let mut row_counts: Vec> = vec![None; rows]; - for (dataset, stats) in atlas.array_stats_by_dataset(column) { - let Some(&row) = row_of.get(dataset.as_str()) else { + for entry in stats { + // A per-dataset entry names its dataset, not its array. + let Some(&row) = row_of.get(entry.name.as_str()) else { continue; }; - mins[row] = stat_to_scalar(stats.min.as_ref(), target, &null); - maxes[row] = stat_to_scalar(stats.max.as_ref(), target, &null); - null_counts[row] = Some(stats.null_count); - row_counts[row] = Some(stats.row_count); + mins[row] = stat_to_scalar(entry.min.as_ref(), target, &null); + maxes[row] = stat_to_scalar(entry.max.as_ref(), target, &null); + null_counts[row] = Some(entry.null_count); + row_counts[row] = Some(entry.row_count); } StatColumn { @@ -362,44 +407,27 @@ fn pack_array_column( } } -/// One attribute column, or `None` when it is not worth the pass. +/// One attribute column, or `None` when no value of it casts. /// /// An attribute's value is exact, so it is both the minimum and the maximum of /// its dataset. That prunes an equality on a dataset-level attribute — the -/// platform a file came from, say — out of the footer alone. +/// platform a file came from, say — from one request. fn pack_attribute_column( - atlas: &Atlas, - names: &[String], + found: &IndexMap, + rows: usize, row_of: &HashMap<&str, usize>, - column: &str, target: &DataType, ) -> Option { - if names.len() > ATTRIBUTE_INDEX_LIMIT { - tracing::debug!( - datasets = names.len(), - column, - "not indexing an attribute over a collection this large; see ATTRIBUTE_INDEX_LIMIT" - ); - return None; - } - - let rows = names.len(); let null = ScalarValue::try_from(target).unwrap_or(ScalarValue::Null); let mut values = vec![null.clone(); rows]; let mut null_counts: Vec> = vec![None; rows]; let mut seen = false; - for name in names { - let Some(&row) = row_of.get(name.as_str()) else { - continue; - }; - let Ok(view) = atlas.dataset(name) else { - continue; - }; - let Some(value) = attribute_of(&view, column) else { + for (dataset, attr) in found { + let Some(&row) = row_of.get(dataset.as_str()) else { continue; }; - let Some(scalar) = attr_to_scalar(&value) else { + let Some(scalar) = attr_to_scalar(attr) else { continue; }; values[row] = scalar.cast_to(target).unwrap_or_else(|_| null.clone()); @@ -424,24 +452,6 @@ fn pack_attribute_column( }) } -/// The attribute a column name refers to, dataset-level or per-array. -fn attribute_of(view: &atlas::DatasetView, column: &str) -> Option { - if let Some(key) = column.strip_prefix('.') { - return view.get_attribute(key); - } - // An array name and an attribute key may both hold dots, so every split is - // a candidate. - for (index, character) in column.char_indices() { - if character == '.' { - let (array, rest) = column.split_at(index); - if let Some(value) = view.get_array_attribute(array, &rest[1..]) { - return Some(value); - } - } - } - None -} - /// Pack scalars into one typed array, or a column of nulls when they will not. fn scalars_to_array(values: Vec, rows: usize, target: &DataType) -> ArrayRef { ScalarValue::iter_to_array(values) @@ -487,7 +497,6 @@ fn attr_to_scalar(attr: &Attr) -> Option { Attr::Float64(v) => ScalarValue::Float64(Some(*v)), Attr::String(v) => ScalarValue::Utf8(Some(v.clone())), Attr::Binary(v) => ScalarValue::Binary(Some(v.clone())), - Attr::TimestampNanoseconds(v) => ScalarValue::TimestampNanosecond(Some(*v), None), _ => return None, }) } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs index 4063e181..ec760635 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs @@ -14,6 +14,11 @@ //! completion, so a collection of a million small datasets and a collection of //! four large ones both divide over every core. //! +//! The second level of the queue is the dataset's own chunk grid, the one the +//! writer chose. A dataset stored as a single chunk is one unit of work; a +//! chunked one is as many as it has chunks, and several partitions drain it +//! together. Each unit becomes one nd record batch. +//! //! [`AtlasFormat`]: super::AtlasFormat use std::any::Any; @@ -256,11 +261,11 @@ impl FileSource for AtlasSource { /// Take the filters as a hint, and leave them above the scan. /// - /// The scan uses a predicate twice: to skip a chunk whose coordinates - /// cannot hold a matching row, and (once pruning lands) to skip a whole - /// dataset whose footer statistics cannot. Neither is exact — both work in - /// whole chunks and whole datasets — so the filter above the scan still - /// decides each row, and `PushedDown::No` is what says so. + /// The scan uses a predicate twice: to skip a whole dataset whose recorded + /// statistics cannot hold a matching row, and to skip a chunk whose + /// coordinates cannot. Neither is exact — both work in whole datasets and + /// whole chunks — so the filter above the scan still decides each row, and + /// `PushedDown::No` is what says so. fn try_pushdown_filters( &self, filters: Vec>, @@ -398,9 +403,9 @@ impl AtlasDatasets { /// column into the projection today, but the chunk pruning inside the read /// matches columns by name and would silently stop pruning if that ever /// changed. - fn projected_names(&self, view: &DatasetView) -> Option> { + async fn projected_names(&self, view: &DatasetView) -> Result>> { if self.projected_schema.fields().is_empty() { - return count_driver(view).map(|driver| vec![driver]); + return Ok(count_driver(view).await?.map(|driver| vec![driver])); } let mut names: Vec = self @@ -417,23 +422,43 @@ impl AtlasDatasets { } } } - Some(names) + Ok(Some(names)) } } /// The array a `COUNT(*)` reads to establish a dataset's row count: the widest /// one Beacon can read. /// -/// From the footer alone, so choosing it costs no I/O and no backend. `None` -/// for a dataset with no readable array, and the caller then builds what there -/// is — an attribute-only dataset contributes the one row its scalars define. -fn count_driver(view: &DatasetView) -> Option { - view.schema() - .arrays +/// The footer names the candidates, and their sizes come from the segments +/// those arrays live in — one open each for the whole collection, however many +/// datasets a `COUNT(*)` walks. No array data is read, and the driver gets the +/// only backend the dataset builds. +/// +/// `None` for a dataset with no readable array, and the caller then builds what +/// there is — an attribute-only dataset contributes the one row its scalars +/// define. +async fn count_driver(view: &DatasetView) -> Result> { + let readable: Vec = view + .schema() .iter() - .filter(|(_, schema)| compat::array_dtype_to_nd(&schema.dtype).is_some()) - .max_by_key(|(_, schema)| schema.shape.iter().product::()) - .map(|(name, _)| name.clone()) + .filter(|meta| compat::array_dtype_to_nd(meta.dtype()).is_some()) + .map(|meta| meta.name().to_string()) + .collect(); + + let mut widest: Option<(String, usize)> = None; + for array in readable { + let layout = view.array_layout(&array).await.map_err(|e| { + DataFusionError::Execution(format!( + "Failed to read the layout of atlas array '{array}' of dataset '{}': {e}", + view.name() + )) + })?; + let cells = layout.element_count(); + if widest.as_ref().is_none_or(|(_, held)| cells > *held) { + widest = Some((array, cells)); + } + } + Ok(widest.map(|(array, _)| array)) } #[async_trait::async_trait] @@ -480,7 +505,7 @@ impl OpenFile for AtlasDatasets { )) })?); - let projected = self.projected_names(&view); + let projected = self.projected_names(&view).await?; let dataset = dataset_from_view(view, projected.as_deref()) .await .map_err(|e| DataFusionError::Execution(format!("{e}")))?; @@ -680,6 +705,8 @@ mod tests { let datasets = datasets_wanting(vec!["temperature"], None); let names = datasets .projected_names(&view(tmp.path(), "winter").await) + .await + .expect("the layouts resolve") .expect("a projection"); assert_eq!(names, vec!["temperature".to_string()]); } @@ -704,6 +731,8 @@ mod tests { let datasets = datasets_wanting(vec!["temperature"], Some(predicate)); let names = datasets .projected_names(&view(tmp.path(), "winter").await) + .await + .expect("the layouts resolve") .expect("a projection"); assert_eq!( names, @@ -723,6 +752,8 @@ mod tests { let datasets = datasets_wanting(vec![], None); let names = datasets .projected_names(&view(tmp.path(), "grid").await) + .await + .expect("the layouts resolve") .expect("a driver"); assert_eq!(names.len(), 1); assert!( diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs index 42507449..d6699c52 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs @@ -1,9 +1,9 @@ //! The column ranges of a whole collection, for the file analyzer. //! -//! A collection reports one range per column, folded over its live datasets out -//! of the footer. It costs no array read: the writer computed each dataset's -//! minimum and maximum while it staged the data, and the open already holds -//! them. +//! A collection reports one range per column, folded over its live datasets. +//! It costs no array read: the writer computed each dataset's minimum and +//! maximum while it staged the data, and stored them in that variable's +//! segment. One open per column answers for the whole collection. //! //! # Why a wrong answer here is worse than no answer //! @@ -18,19 +18,18 @@ use datafusion::common::{ColumnStatistics, Statistics, stats::Precision}; use datafusion::scalar::ScalarValue; /// The statistics of one collection, in `table_schema` order. -pub fn collection_statistics(atlas: &Atlas, table_schema: &Schema) -> Statistics { +pub async fn collection_statistics(atlas: &Atlas, table_schema: &Schema) -> Statistics { let arrays = atlas.list_arrays(); let live = atlas.dataset_count(); let mut statistics = Statistics::default(); for field in table_schema.fields() { let range = if arrays.iter().any(|array| array == field.name()) { - column_range(atlas, field.name(), field.data_type(), live) + column_range(atlas, field.name(), field.data_type(), live).await } else { - // An attribute column would need one `DatasetView` per dataset, and - // each of those is a linear scan of the footer. The query-time - // pruning index pays that where it is bounded and worth it; a - // background pass over every collection is not the place. + // An attribute is exact rather than a range, and bounding a column + // by it says nothing a scan can use. The query-time pruning index + // reads attributes; the analyzer has no use for them. None }; @@ -54,26 +53,25 @@ pub fn collection_statistics(atlas: &Atlas, table_schema: &Schema) -> Statistics /// report would then produce a range those zeros sit outside of, and pruning /// would drop the collection for a query that matches them. /// -/// The footer cannot tell "declares it and never wrote it" from "does not -/// declare it" without a view per dataset, which is a linear scan each. So the +/// A segment holds an entry only for a dataset that wrote the array, so the /// count is the proof: a bound is claimed only when every live dataset reported /// one. A uniform collection — which is what `atlas create` writes, and what /// this format exists for — satisfies that; a heterogeneous one goes unknown /// and is read in full. -fn column_range( +async fn column_range( atlas: &Atlas, column: &str, target: &DataType, live: usize, ) -> Option<(ScalarValue, ScalarValue)> { - let per_dataset = atlas.array_stats_by_dataset(column); + let per_dataset = atlas.array_stats_by_dataset(column).await.ok()?; if per_dataset.is_empty() || per_dataset.len() != live { return None; } let mut low: Option = None; let mut high: Option = None; - for (_, stats) in per_dataset { + for stats in per_dataset { // One dataset without a bound leaves the column unbounded: its values // may lie anywhere. let min = bound(stats.min.as_ref(), target)?; @@ -155,7 +153,8 @@ mod tests { let statistics = collection_statistics( &atlas, &schema(vec![Field::new("temperature", DataType::Float32, true)]), - ); + ) + .await; let (min, max) = range(&statistics, 0); // d0 starts at 0 and d9 ends at 93. assert_eq!(min, Precision::Exact(ScalarValue::Float32(Some(0.0)))); @@ -179,7 +178,8 @@ mod tests { // both declare `temperature`. Field::new("temperature", DataType::Float32, true), ]), - ); + ) + .await; assert_eq!(range(&statistics, 0).0, Precision::Absent); assert_eq!( range(&statistics, 1).0, @@ -202,7 +202,8 @@ mod tests { let statistics = collection_statistics( &atlas, &schema(vec![Field::new("value", DataType::Float64, true)]), - ); + ) + .await; let (min, max) = range(&statistics, 0); // a holds [1, 2] as Int16 and b holds [3.5, 4.5] as Float32. assert_eq!(min, Precision::Exact(ScalarValue::Float64(Some(1.0)))); @@ -223,7 +224,8 @@ mod tests { Field::new("ghost", DataType::Float32, true), Field::new(".platform", DataType::Utf8, true), ]), - ); + ) + .await; assert_eq!(range(&statistics, 0).0, Precision::Absent); assert_eq!( range(&statistics, 1).0, @@ -241,7 +243,8 @@ mod tests { let statistics = collection_statistics( &atlas, &schema(vec![Field::new("temperature", DataType::Float32, true)]), - ); + ) + .await; assert_eq!(range(&statistics, 0).0, Precision::Absent); } @@ -259,7 +262,8 @@ mod tests { let statistics = collection_statistics( &atlas, &schema(vec![Field::new("temperature", DataType::Float32, true)]), - ); + ) + .await; assert_eq!( range(&statistics, 0).1, Precision::Exact(ScalarValue::Float32(Some(83.0))), diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs index b89054e6..3aaf667e 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs @@ -8,15 +8,23 @@ //! //! ```text //! my_collection/ -//! ├── data.atlas ATLS │ segment │ segment │ … │ footer │ trailer +//! ├── data.atlas ATLS │ temperature │ salinity │ … │ footer │ trailer //! └── deleted.mask optional: ordinals of deleted datasets //! ``` //! -//! Each dataset occupies one segment. A footer at the end records every dataset -//! name, its segment byte range, its schema, its attribute values and its -//! per-array statistics. Opening a collection reads that footer and nothing -//! else, so every metadata question is answered with no further I/O, whatever -//! the dataset count. Array data arrives chunk by chunk, on demand. +//! **One segment is one variable, not one dataset.** A segment holds one array +//! name across the whole collection, and each dataset's copy sits inside it +//! under the dataset's own name. A footer at the end records every dataset +//! name, every variable's byte range, and the arrays and attribute keys each +//! dataset declares — with their element types, and nothing more. +//! +//! Opening a collection reads that footer, so listing the datasets and asking +//! what one declares cost no further I/O, whatever the dataset count. Three +//! things are *not* in the footer, and each comes from the variable's own +//! segment: an array's layout (shape, chunking, dimension names, fill value), +//! its statistics, and every attribute value. One open answers each of those +//! for the whole collection, so reading `temperature` across a million datasets +//! opens one segment. Array data then arrives block by block, on demand. //! //! # What this crate does with it //! @@ -27,6 +35,14 @@ //! derives the Arrow schema of a whole collection. [`compat`] holds the type //! and column-name mapping the two share. //! +//! # One dataset is one unit of work +//! +//! The scan plans one entry per dataset and puts them all in one morsel queue, +//! as the netCDF and Zarr scans do for their files. A partition takes the next +//! dataset when it is free, and each dataset is cut on the chunk grid the +//! writer actually chose, so a dataset stored as one chunk yields one unit and +//! a chunked one yields many. See [`datafusion::source`]. +//! //! # Columns //! //! One column per array, under the array's own name. A per-array attribute @@ -41,6 +57,10 @@ //! collection a Rust writer produced. The mapping refuses them all the same. //! - A list-valued *attribute*. Beacon's ND model has no rank-0 list. //! +//! A timestamp *attribute* does not arise: atlas stores none, because an +//! attribute would have to go to disk as a plain `i64` and could not come back +//! as a timestamp. An array element type still has its own timestamp. +//! //! Each is dropped from the dataset with a `debug` log rather than failing the //! scan. A collection can hold a million datasets, so a `warn` per skip would //! be a flood. diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs index fbf3c184..2718b27f 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs @@ -10,7 +10,7 @@ use std::collections::HashSet; use std::sync::Arc; use arrow::datatypes::{Schema, SchemaRef}; -use atlas::{Atlas, DatasetView}; +use atlas::{Atlas, Attr, DType, DatasetView}; use beacon_datafusion_ext::type_widening::{ArrowTypeWidening, LabeledSchema}; use beacon_nd_array::{ NdArrayD, @@ -39,7 +39,15 @@ use crate::compat; /// /// A value Beacon cannot surface — a `Bool` or list array, a list attribute — /// is dropped with a `debug` log. A collection can hold a million datasets, so -/// a louder log would be a flood. +/// a louder log would be a flood. Such an array is settled from the footer +/// alone, so its segment is never opened. +/// +/// # What this reads +/// +/// Names and element types come from the footer the open already held. A +/// layout and an attribute value live in the variable's own segment, so the +/// first dataset to want one opens it. A segment covers that variable across +/// the whole collection, so every later dataset reuses the same handle. pub async fn dataset_from_view( view: Arc, projected_names: Option<&[String]>, @@ -57,20 +65,40 @@ pub async fn dataset_from_view( let wants_global_attrs = projected_names.is_none_or(|names| names.iter().any(|name| name.starts_with('.'))); - let schema = view.schema(); - let mut arrays: IndexMap> = - IndexMap::with_capacity(schema.arrays.len()); + // The schema borrows the footer, so it is resolved before the first await + // rather than held across one. + let declared = declared_arrays(&view); + let mut arrays: IndexMap> = IndexMap::with_capacity(declared.len()); - for (array_name, array_schema) in &schema.arrays { + for (array_name, dtype) in &declared { if included(array_name) { - match compat::array_to_nd_array(Arc::clone(&view), array_name, array_schema) { - Ok(nd) => { - arrays.insert(array_name.clone(), nd); + match compat::array_dtype_to_nd(dtype) { + // The layout is in the variable's segment, so it is asked for + // only once the dtype says the column can exist at all. + Some(_) => { + let layout = view.array_layout(array_name).await.map_err(|e| { + anyhow::anyhow!( + "Failed to read the layout of atlas array '{array_name}' \ + of dataset '{}': {e}", + view.name() + ) + })?; + match compat::array_to_nd_array(Arc::clone(&view), array_name, dtype, &layout) { + Ok(nd) => { + arrays.insert(array_name.clone(), nd); + } + Err(e) => tracing::debug!( + dataset = %view.name(), + array = %array_name, + "atlas array left out of the dataset: {e}" + ), + } } - Err(e) => tracing::debug!( + None => tracing::debug!( dataset = %view.name(), array = %array_name, - "atlas array left out of the dataset: {e}" + "atlas array left out of the dataset: {} is no Beacon column", + compat::dtype_tag(dtype) ), } } @@ -78,7 +106,7 @@ pub async fn dataset_from_view( if !wants_attrs_of(array_name) { continue; } - for (key, value) in view.array_attributes(array_name) { + for (key, value) in attributes_of(&view, Some(array_name)).await? { let column = compat::array_attr_column(array_name, &key); if !included(&column) { continue; @@ -97,7 +125,7 @@ pub async fn dataset_from_view( } if wants_global_attrs { - for (key, value) in view.attributes() { + for (key, value) in attributes_of(&view, None).await? { let column = compat::global_attr_column(&key); if !included(&column) { continue; @@ -123,6 +151,43 @@ pub async fn dataset_from_view( .map_err(|e| anyhow::anyhow!("Failed to wrap atlas dataset '{}': {e}", view.name())) } +/// Every array the dataset declares, as owned name and element type. +/// +/// The schema borrows the collection footer. Resolving it up front keeps that +/// borrow off the async path, where the same view is also cloned into a +/// backend. +fn declared_arrays(view: &DatasetView) -> Vec<(String, DType)> { + view.schema() + .iter() + .map(|meta| (meta.name().to_string(), meta.dtype().clone())) + .collect() +} + +/// The attribute values of one scope: `Some(array)` for an array's own, +/// `None` for the dataset's. +/// +/// Values live in a segment, not in the footer, so this reads one. An array +/// with no attribute costs nothing, because the schema settles it first. +async fn attributes_of( + view: &DatasetView, + array: Option<&str>, +) -> anyhow::Result> { + let scope = match array { + Some(array) => view.array_attributes(array).await, + None => view.attributes().await, + }; + scope.map_err(|e| { + let what = array.map_or_else( + || "the dataset attributes".to_string(), + |array| format!("the attributes of array '{array}'"), + ); + anyhow::anyhow!( + "Failed to read {what} of atlas dataset '{}': {e}", + view.name() + ) + }) +} + /// Narrow `dataset` to `read_dimensions`, or to a broadcast-compatible default /// when none are given. /// @@ -175,13 +240,11 @@ pub async fn open_dataset( /// /// # Cost /// -/// No I/O: everything read here came in with the footer. The pass is linear in -/// the dataset count, but resolving each dataset by name is itself a linear -/// scan of the footer, so the whole pass is quadratic in the dataset count. -/// That is a limit of the reader's API, which offers no lookup by ordinal, and -/// it is why a collection of more than a few tens of thousands of datasets -/// needs `Atlas::dataset_at` upstream. The result is cached above this crate, -/// so a table pays it once rather than once per query. +/// Linear in the dataset count, and every step is in memory: the footer keys +/// its datasets by name, so resolving one is a single hash lookup, and a key is +/// built from the footer and from segments that one open serves collection +/// wide. Only the distinct keys are derived into a schema. The result is cached +/// above this crate, so a table pays even that once rather than once per query. pub async fn collection_schema( atlas: &Arc, read_dimensions: Option<&[String]>, @@ -196,7 +259,7 @@ pub async fn collection_schema( .dataset(&name) .map_err(|e| anyhow::anyhow!("Failed to open atlas dataset '{name}': {e}"))?; - if !seen.insert(shape_key(&view)) { + if !seen.insert(shape_key(&view).await?) { continue; } @@ -228,32 +291,71 @@ pub async fn collection_schema( /// What makes two datasets produce the same columns and types. /// -/// The interned schema decides the arrays, and it is shared by address between -/// datasets that declare the same ones, so its pointer is the cheap half of the -/// key. Attribute *values* live outside the schema, so the keys and their types -/// are the other half. Two datasets that differ only in an attribute's value -/// share a key, which is exactly the fleet case. -fn shape_key(view: &DatasetView) -> String { - let schema = view.schema(); - let mut key = format!("{:x}", std::ptr::from_ref(schema) as usize); - - for array in schema.arrays.keys() { - for (attr, value) in view.array_attributes(array) { - key.push('|'); - key.push_str(array); - key.push('.'); - key.push_str(&attr); - key.push(':'); - key.push_str(&compat::dtype_tag(&value.dtype())); +/// Three things decide a dataset's Arrow schema, and the key holds all three: +/// +/// - The arrays it declares, with their element types. Datasets that declare +/// the same ones share one interned schema in the footer, and `atlas create` +/// writes a fleet of files that way. +/// - Its attribute keys and their types, at both scopes. Those are named in the +/// interned schema too, so two datasets that differ only in an attribute's +/// *value* share a key. That is exactly the fleet case. +/// - Each array's dimension names, which the interned schema does **not** hold. +/// They pick the default grid, and a different grid keeps different columns, +/// so two datasets that agree on everything else can still differ here. +/// +/// A shape is deliberately left out: an array of a different length is the same +/// column. +/// +/// # Cost +/// +/// The names and the types come from the footer. A dimension name comes from +/// its variable's segment, which one open serves for every dataset of the +/// collection, so the lookup is in memory after the first. +async fn shape_key(view: &DatasetView) -> anyhow::Result { + let mut key = String::new(); + + for (array, dtype) in declared_arrays(view) { + key.push('|'); + key.push_str(&array); + key.push(':'); + key.push_str(&compat::dtype_tag(&dtype)); + + // An array Beacon cannot read is no column, so its grid decides + // nothing and its segment stays shut. + if compat::array_dtype_to_nd(&dtype).is_some() { + let layout = view.array_layout(&array).await.map_err(|e| { + anyhow::anyhow!( + "Failed to read the layout of atlas array '{array}' of dataset '{}': {e}", + view.name() + ) + })?; + key.push('@'); + key.push_str(&layout.dimension_names().join(",")); } } - for (attr, value) in view.attributes() { - key.push_str("|."); - key.push_str(&attr); + + // Attribute keys and types are in the interned schema, so this reads + // nothing. The values are not, and they do not belong in the key. + fn push(key: &mut String, array: &str, attr: &str, dtype: &DType) { + key.push('|'); + key.push_str(array); + key.push('.'); + key.push_str(attr); key.push(':'); - key.push_str(&compat::dtype_tag(&value.dtype())); + key.push_str(&compat::dtype_tag(dtype)); + } + let schema = view.schema(); + for meta in schema.iter() { + let array = meta.name(); + for (attr, dtype) in meta.attribute_pairs() { + push(&mut key, array, attr, dtype); + } + } + for (attr, dtype) in schema.attribute_pairs() { + push(&mut key, "", attr, dtype); } - key + + Ok(key) } #[cfg(test)] @@ -656,7 +758,7 @@ mod tests { let mut keys = HashSet::new(); for name in atlas.list_datasets() { - keys.insert(shape_key(&atlas.dataset(&name).unwrap())); + keys.insert(shape_key(&atlas.dataset(&name).unwrap()).await.unwrap()); } assert_eq!(keys.len(), 1, "and every dataset reduces to one key"); @@ -675,8 +777,8 @@ mod tests { test_support::two_datasets(tmp.path()).await; let atlas = test_support::open(tmp.path()).await; - let winter = shape_key(&atlas.dataset("winter").unwrap()); - let summer = shape_key(&atlas.dataset("summer").unwrap()); + let winter = shape_key(&atlas.dataset("winter").unwrap()).await.unwrap(); + let summer = shape_key(&atlas.dataset("summer").unwrap()).await.unwrap(); assert_ne!(winter, summer); } diff --git a/docs/docs/2.0.0-rc5/formats/atlas.md b/docs/docs/2.0.0-rc5/formats/atlas.md index 085b7948..404a09f8 100644 --- a/docs/docs/2.0.0-rc5/formats/atlas.md +++ b/docs/docs/2.0.0-rc5/formats/atlas.md @@ -45,29 +45,36 @@ NetCDF file holds: named arrays that share dimensions, plus attributes. A collec ```text my_collection/ -├── data.atlas the container: every dataset, then a footer describing them all +├── data.atlas one segment per variable, then a footer describing them all └── deleted.mask optional: the datasets a delete has hidden ``` -Two properties follow, and they are the point of the format: +The file stores one segment per **variable**, not one per dataset. A segment holds one array name +across the whole collection, and each dataset's copy sits inside it. Three properties follow, and +they are the point of the format: -- **Metadata is one read.** Opening a collection reads its footer and nothing else. Listing the - datasets, inspecting a schema and reading an attribute are then free. Ten datasets and a million - cost the same. -- **Data arrives chunk by chunk.** Reading a region of an array fetches only the chunks that region - overlaps. +- **The catalogue is one read.** Opening a collection reads its footer and nothing else. Listing + the datasets and asking what each declares are then free. Ten datasets and a million cost the + same. +- **One variable is one read.** Everything else about a column — its shape, its statistics, its + attribute values — sits in that variable's segment, and one open answers for every dataset in the + collection. +- **Data arrives block by block.** Reading a region of an array fetches only the blocks that region + overlaps, and a block holds one type for a run of neighbouring datasets, so it compresses well. What Beacon does with that: -- **Dataset pruning from the footer.** Every array records its minimum, its maximum and its null - count. A query with a predicate — a time or latitude range, say — judges every dataset of a - collection in one pass and never opens the ones that cannot match. A dataset-level attribute is - exact in the footer, so `WHERE ".platform" = 'p3'` prunes on it too. +- **Dataset pruning.** Every array records its minimum, its maximum and its null count. A query + with a predicate — a time or latitude range, say — judges every dataset of a collection in one + vectorised pass and never opens the ones that cannot match. Judging a column costs one request, + whatever the dataset count. A dataset-level attribute is exact, so `WHERE ".platform" = 'p3'` + prunes on it too. - **One dataset is one unit of work.** A collection's datasets are spread across every core, and a worker takes the next one when it is free, so a collection of a million small datasets and one of - four large ones both divide evenly. + four large ones both divide evenly. A dataset stored in several chunks divides further, so a + single large dataset still uses every core. - **Column projection.** Only the arrays a query names get read, and only their attributes are - taken from the footer. + fetched. - **Object storage.** A collection reads from local disk, S3, GCS, Azure and HTTP alike. ### Columns @@ -78,6 +85,9 @@ What Beacon does with that: | attribute `units` of `temperature` | `temperature.units` | | dataset attribute `platform` | `.platform` | +An attribute holds a number, a string or a boolean. Atlas stores no timestamp attribute, so a date +kept as an attribute arrives as the number or the string it was written as. + The leading dot on a dataset attribute is what NetCDF and Zarr use too, and it keeps an attribute from colliding with an array of the same name. Quote such a column: `SELECT ".platform"`. @@ -109,8 +119,8 @@ attribute. Each is dropped from the schema rather than failing the query. ### Two datasets that disagree -Atlas reconciles nothing: two datasets may declare one array name with two types. Beacon merges -them the way it merges the files of any other format. Two numeric types widen to one that holds +Atlas reconciles nothing: two datasets may declare one array name with two types, and it stores +each as declared. Beacon merges them the way it merges the files of any other format. Two numeric types widen to one that holds both. Two different families — a number and a string — refuse the table by name: ```text @@ -132,10 +142,11 @@ That writes `/collections/argo/data.atlas`, one dataset per file, named after th against a local path and against a bucket. See the [Atlas documentation](https://github.com/maris-development/atlas). -:::warning Collections written before Atlas 0.16 -Atlas used to be a directory of per-array files behind an `atlas.json` registry. Beacon reads the -single-file format only, so such a directory is passed over rather than read. Rewrite it with -`atlas create`, then point at the `data.atlas` it produces. +:::warning Collections written before Atlas 0.17 +Beacon reads container format version 8, which Atlas 0.17 writes. An older collection — a v1 +container from Atlas 0.16, or the directory of per-array files behind an `atlas.json` registry that +came before it — is not read. Rewrite it with `atlas create`, then point at the `data.atlas` it +produces. ::: ### Optimize NetCDF and Zarr with Atlas diff --git a/integration-tests/formats/test_atlas.py b/integration-tests/formats/test_atlas.py index 8aae9291..67dd58c7 100644 --- a/integration-tests/formats/test_atlas.py +++ b/integration-tests/formats/test_atlas.py @@ -222,13 +222,13 @@ def test_pruning_can_be_turned_off_per_table(datasets, tmp_path): # --- what is not supported ---------------------------------------------------- -def test_a_pre_0_16_collection_is_passed_over(con, datasets): +def test_a_stale_collection_is_passed_over(con, datasets): """Atlas before 0.16 was a directory behind an `atlas.json` registry. - This build reads the single-file format alone. Such a directory holds no container, so - nothing recognises it and no row of it reaches a query. Whether that surfaces as an error - or as an empty result is not the point and not asserted — what matters is that a stale - collection is never half-read as if it were a current one. + This build reads container version 8 alone, which Atlas 0.17 writes. Such a directory + holds no container at all, so nothing recognises it and no row of it reaches a query. + Whether that surfaces as an error or as an empty result is not the point and not asserted — + what matters is that a stale collection is never half-read as if it were a current one. """ legacy = datasets / "legacy" legacy.mkdir(exist_ok=True) @@ -238,4 +238,4 @@ def test_a_pre_0_16_collection_is_passed_over(con, datasets): rows = con.sql("SELECT * FROM read_atlas('legacy/atlas.json')").fetchall() except Exception: return - assert rows == [], "a pre-0.16 directory holds no container, so it contributes no rows" + assert rows == [], "a legacy directory holds no container, so it contributes no rows" diff --git a/integration-tests/requirements.txt b/integration-tests/requirements.txt index 10a2bb05..7f8cec72 100644 --- a/integration-tests/requirements.txt +++ b/integration-tests/requirements.txt @@ -9,7 +9,7 @@ pyarrow==21.0.0 netCDF4==1.7.4 zarr==3.3.0 # Builds an Atlas collection from netCDF files; the Rust side reads it. -atlas-python==0.16.4 +atlas-python==0.17.0 h5py==3.14.0 geopandas==1.1.4 rasterio==1.5.1 From 77b3330a2237e59e8a691d9c1b0032ec9c92b8ab Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Mon, 7 Sep 2026 17:17:32 +0200 Subject: [PATCH 04/16] wip --- .../beacon-arrow-atlas/src/datafusion/mod.rs | 522 ++++++++++++++++++ .../src/datafusion/source.rs | 34 +- .../beacon-arrow-atlas/src/reader.rs | 150 ++++- .../beacon-arrow-atlas/src/test_support.rs | 118 ++++ 4 files changed, 818 insertions(+), 6 deletions(-) diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs index f03a2623..635d8b06 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs @@ -1216,4 +1216,526 @@ mod tests { .is_none() ); } + // ── a wide collection on two dimensions, through SQL ───────────────── + + /// The shapes every test in this section reads: three datasets whose + /// profile and level counts all differ. + const WIDE: &[(usize, usize)] = &[(3, 4), (5, 4), (2, 6)]; + + /// Rows a full-grid read returns: `profiles * levels`, summed per dataset. + const WIDE_GRID_ROWS: usize = 44; + + /// Rows a read narrowed to `profile` returns: the profiles themselves. + const WIDE_PROFILE_ROWS: usize = 10; + + /// Columns the merged collection carries: eight arrays, four attributes + /// each, and two dataset attributes. + const WIDE_COLUMNS: usize = 42; + + /// Columns that survive a narrowing to `profile`. The four grid arrays go; + /// their attributes stay. + const WIDE_PROFILE_COLUMNS: usize = 38; + + /// A wide collection in a temporary directory, and a table over it. + /// + /// The caller holds the returned directory. It deletes the collection when + /// it drops. + async fn wide_table(ctx: &SessionContext, format: Arc) -> tempfile::TempDir { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + register_with(ctx, tmp.path(), "wide", format).await; + tmp + } + + /// One scalar, as its rendered text. Enough to pin a value without a + /// downcast per Arrow type. + async fn scalar(ctx: &SessionContext, sql: &str) -> String { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string() + .lines() + .nth(3) + .expect("a one-row result") + .trim() + .trim_matches('|') + .trim() + .to_string() + } + + /// The row count is the grid every dataset contributes, and not the + /// profile count. + #[tokio::test] + async fn a_wide_collection_reads_every_row() { + let ctx = context(4); + let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; + + assert_eq!( + count(&ctx, "SELECT COUNT(*) FROM wide").await as usize, + WIDE_GRID_ROWS + ); + } + + #[tokio::test] + async fn a_wide_table_carries_every_column() { + let ctx = context(1); + let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; + + let df = ctx.sql("SELECT * FROM wide").await.unwrap(); + assert_eq!(df.schema().fields().len(), WIDE_COLUMNS); + } + + /// A per-profile array beside a per-level one, in one result. + #[tokio::test] + async fn a_wide_row_carries_both_of_its_grids() { + let ctx = context(1); + let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; + + let batches = ctx + .sql( + "SELECT platform, latitude, pressure, temperature, salinity \ + FROM wide WHERE platform = 'set1' AND pressure = 2.0 \ + ORDER BY latitude LIMIT 1", + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let rendered = arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string(); + // set1 holds temperature = 100 + level, and level 2 is pressure 2. + assert!(rendered.contains("102.0"), "{rendered}"); + assert!(rendered.contains("32.0"), "{rendered}"); + } + + /// An attribute rides along as a constant column. + #[tokio::test] + async fn an_attribute_of_a_wide_collection_reads_as_a_column() { + let ctx = context(1); + let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; + + assert_eq!( + scalar(&ctx, "SELECT DISTINCT \".title\" FROM wide").await, + "wide profiles" + ); + assert_eq!( + scalar(&ctx, "SELECT DISTINCT \"temperature.long_name\" FROM wide").await, + "the temperature" + ); + } + + /// A predicate over a real column, with a known answer. Each dataset owns a + /// disjoint `temperature` range, so this is the pruning arithmetic too. + #[tokio::test] + async fn a_predicate_over_a_wide_collection_selects_the_rows_that_match() { + let ctx = context(4); + let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; + + // Only set2 reaches past 150: its 2 * 6 cells hold 200 to 205. + assert_eq!( + count(&ctx, "SELECT COUNT(*) FROM wide WHERE temperature > 150.0").await, + 12 + ); + assert_eq!( + count(&ctx, "SELECT COUNT(DISTINCT platform) FROM wide").await, + 3 + ); + } + + /// Pruning reads the footer's per-dataset ranges and drops what cannot + /// match. It must not change an answer, whichever way the switch is set. + #[tokio::test] + async fn pruning_does_not_change_the_answers_over_a_wide_collection() { + let queries = [ + "SELECT COUNT(*) FROM wide WHERE temperature > 150.0", + "SELECT COUNT(*) FROM wide WHERE latitude > 12.0", + "SELECT COUNT(*) FROM wide WHERE platform = 'set1'", + "SELECT COUNT(*) FROM wide WHERE salinity < 32.0", + ]; + + let pruned = context(4); + let _a = wide_table(&pruned, Arc::new(AtlasFormat::default())).await; + let whole = context(4); + let _b = wide_table(&whole, Arc::new(AtlasFormat::default().with_pruning(false))).await; + + for sql in queries { + assert_eq!( + count(&pruned, sql).await, + count(&whole, sql).await, + "pruning changed the answer to: {sql}" + ); + } + } + + /// A projected scan over a pruned collection. + /// + /// The predicate is indexed against the table schema, and the pruning + /// engine reads it against the projected one. `temperature` sits at column + /// 32 of 42, so a five-column projection puts that index out of range. This + /// is the case where the two must be brought into step. + #[tokio::test] + async fn a_projected_scan_prunes_without_losing_its_columns() { + let ctx = context(4); + let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; + + let rows = rows( + &ctx, + "SELECT platform, latitude, longitude, pressure, temperature \ + FROM wide WHERE temperature IS NOT NULL \ + ORDER BY platform, latitude, pressure LIMIT 5", + ) + .await; + assert_eq!(rows, 5); + } + + /// A narrowing to `profile` reads the profiles and not their levels. + /// + /// `COUNT(*)` projects nothing, so the scan picks the widest array of the + /// dataset to count. That array lives on both dimensions, and this + /// narrowing drops it. So the driver has to respect the dimensions too. + #[tokio::test] + async fn narrowing_a_wide_table_to_one_dimension_reads_one_row_per_profile() { + let ctx = context(4); + let _tmp = wide_table( + &ctx, + Arc::new(AtlasFormat::new(AtlasOptions { + read_dimensions: Some(vec!["profile".to_string()]), + })), + ) + .await; + + assert_eq!( + count(&ctx, "SELECT COUNT(*) FROM wide").await as usize, + WIDE_PROFILE_ROWS + ); + // The same number through a column, rather than the count driver. + assert_eq!( + count(&ctx, "SELECT COUNT(latitude) FROM wide").await as usize, + WIDE_PROFILE_ROWS + ); + assert_eq!( + ctx.sql("SELECT * FROM wide") + .await + .unwrap() + .schema() + .fields() + .len(), + WIDE_PROFILE_COLUMNS + ); + } + + /// A narrowed read whose projection names only attributes builds no + /// dimensioned array. The narrowing has nothing to drop and must still + /// succeed. + #[tokio::test] + async fn a_narrowed_read_of_attributes_alone_succeeds() { + let ctx = context(1); + let _tmp = wide_table( + &ctx, + Arc::new(AtlasFormat::new(AtlasOptions { + read_dimensions: Some(vec!["profile".to_string()]), + })), + ) + .await; + + assert_eq!( + scalar(&ctx, "SELECT DISTINCT \"temperature.units\" FROM wide").await, + "1" + ); + assert_eq!( + scalar(&ctx, "SELECT DISTINCT \".institution\" FROM wide").await, + "test" + ); + } + + /// Every dataset is one unit of work, and a partitioned scan divides them + /// without reading one twice. + #[tokio::test] + async fn a_partitioned_scan_of_a_wide_collection_reads_every_row_once() { + for partitions in [1, 2, 8] { + let ctx = context(partitions); + let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; + assert_eq!( + count(&ctx, "SELECT COUNT(*) FROM wide").await as usize, + WIDE_GRID_ROWS, + "over {partitions} partitions" + ); + } + } + + // ── STORED AS ATLAS ────────────────────────────────────────────────── + + /// A session that answers `STORED AS ATLAS`, as the runtime builds one. + /// + /// Two lookups sit behind that clause, under two spellings of one name. + /// DataFusion resolves the `STORED AS` word in `table_factories`, upper + /// cased, and Beacon registers [`ListingTableFactoryExt`] there. That + /// factory then resolves the *file format* by the same word lower cased, + /// which is where [`ATLAS_FORMAT`] answers. + fn ddl_context(partitions: usize) -> SessionContext { + use beacon_datafusion_ext::listing_table_factory_ext::ListingTableFactoryExt; + use datafusion::execution::session_state::SessionStateBuilder; + + let mut config = SessionConfig::new() + .with_target_partitions(partitions) + .with_extension(Arc::new(ListingFactory::dynamic())) + .with_extension(Arc::new(ListingTableFactoryExt)); + // What the runtime sets. DataFusion's default matches a glob against + // the file name alone, and a collection is always one directory down, + // so `**/data.atlas` would list nothing and the table would be empty. + config + .options_mut() + .execution + .listing_table_ignore_subdirectory = false; + let state = SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .with_table_factory( + ATLAS_FORMAT.to_uppercase(), + Arc::new(ListingTableFactoryExt), + ) + .build(); + let ctx = SessionContext::new_with_state(state); + ctx.state_ref() + .write() + .register_file_format( + Arc::new(AtlasFormatFactory::new( + Default::default(), + Default::default(), + )), + true, + ) + .expect("the atlas format registers under its own name"); + ctx + } + + /// A directory as SQL takes it: forward slashes, whatever the platform. + fn location(dir: &Path) -> String { + dir.to_string_lossy() + .replace(std::path::MAIN_SEPARATOR, "/") + } + + /// Run one DDL statement, and name the statement if it is refused. + async fn run_ddl(ctx: &SessionContext, sql: &str) { + ctx.sql(sql) + .await + .unwrap_or_else(|e| panic!("refused `{sql}`: {e}")) + .collect() + .await + .unwrap(); + } + + /// The clause a user writes, over the directory that holds the container. + #[tokio::test] + async fn stored_as_atlas_reads_a_collection_directory() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + let ctx = ddl_context(4); + + run_ddl( + &ctx, + &format!( + "CREATE EXTERNAL TABLE t STORED AS ATLAS LOCATION '{}/'", + location(tmp.path()) + ), + ) + .await; + + assert_eq!( + count(&ctx, "SELECT COUNT(*) FROM t").await as usize, + WIDE_GRID_ROWS + ); + assert_eq!( + ctx.sql("SELECT * FROM t") + .await + .unwrap() + .schema() + .fields() + .len(), + WIDE_COLUMNS + ); + } + + /// The container object names the collection just as well as its directory + /// does. That is what a `LOCATION` copied out of a listing looks like. + #[tokio::test] + async fn stored_as_atlas_reads_the_container_object() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + let ctx = ddl_context(4); + + run_ddl( + &ctx, + &format!( + "CREATE EXTERNAL TABLE t STORED AS ATLAS LOCATION '{}/{ATLAS_MARKER}'", + location(tmp.path()) + ), + ) + .await; + + assert_eq!( + count(&ctx, "SELECT COUNT(*) FROM t").await as usize, + WIDE_GRID_ROWS + ); + } + + /// SQL folds the `STORED AS` word, so either spelling reaches the format. + #[tokio::test] + async fn stored_as_atlas_is_case_insensitive() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + + for spelling in ["ATLAS", "atlas", "Atlas"] { + let ctx = ddl_context(1); + run_ddl( + &ctx, + &format!( + "CREATE EXTERNAL TABLE t STORED AS {spelling} LOCATION '{}/'", + location(tmp.path()) + ), + ) + .await; + + assert_eq!( + count(&ctx, "SELECT COUNT(*) FROM t").await as usize, + WIDE_GRID_ROWS, + "STORED AS {spelling}" + ); + } + } + + /// `OPTIONS` reaches the format, so a table can name its dimensions + /// without the `read_atlas` function. + #[tokio::test] + async fn stored_as_atlas_takes_its_options() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + let ctx = ddl_context(4); + + run_ddl( + &ctx, + &format!( + "CREATE EXTERNAL TABLE t STORED AS ATLAS \ + OPTIONS ('read_dimensions' 'profile') LOCATION '{}/'", + location(tmp.path()) + ), + ) + .await; + + assert_eq!( + count(&ctx, "SELECT COUNT(*) FROM t").await as usize, + WIDE_PROFILE_ROWS + ); + } + + /// An `OPTIONS` value is a string, so a list of dimensions is one comma + /// separated string. `read_atlas` takes a real SQL list and joins it into + /// the same option, so both spellings reach one place. + /// + /// Both dimensions of this collection reads it whole, which is what the + /// default already does. The point is the parse, not the answer. + #[tokio::test] + async fn stored_as_atlas_takes_a_comma_separated_dimension_list() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + + // A space after a comma, and a trailing comma, are both tolerated. + for list in ["profile,level", "profile, level", "level,profile,"] { + let ctx = ddl_context(4); + run_ddl( + &ctx, + &format!( + "CREATE EXTERNAL TABLE t STORED AS ATLAS \ + OPTIONS ('read_dimensions' '{list}') LOCATION '{}/'", + location(tmp.path()) + ), + ) + .await; + + assert_eq!( + count(&ctx, "SELECT COUNT(*) FROM t").await as usize, + WIDE_GRID_ROWS, + "'{list}' did not name the whole grid" + ); + assert_eq!( + ctx.sql("SELECT * FROM t") + .await + .unwrap() + .schema() + .fields() + .len(), + WIDE_COLUMNS, + "'{list}'" + ); + } + } + + /// A glob puts several collections in one table, and the rows are their + /// union. + /// + /// This is the form the docs give for a data lake: + /// `LOCATION 'collections/**/data.atlas'`. It rests on + /// `listing_table_ignore_subdirectory` being off, because a collection's + /// container always sits one directory below the glob's prefix. + #[tokio::test] + async fn stored_as_atlas_globs_several_collections_into_one_table() { + // Two collections under one root, so the union is a number neither one + // could produce alone. + let tmp = tempfile::tempdir().unwrap(); + for name in ["one", "two"] { + test_support::wide_profiles(&tmp.path().join(name), WIDE).await; + } + let root = location(tmp.path()); + + for glob in ["**/data.atlas", "*/data.atlas"] { + let ctx = ddl_context(4); + run_ddl( + &ctx, + &format!("CREATE EXTERNAL TABLE t STORED AS ATLAS LOCATION '{root}/{glob}'"), + ) + .await; + + assert_eq!( + count(&ctx, "SELECT COUNT(*) FROM t").await as usize, + 2 * WIDE_GRID_ROWS, + "'{glob}' did not read both collections" + ); + // One schema over both, merged under the session's widening rule. + assert_eq!( + ctx.sql("SELECT * FROM t") + .await + .unwrap() + .schema() + .fields() + .len(), + WIDE_COLUMNS, + "'{glob}'" + ); + } + } + + /// A bad option is an error at `CREATE EXTERNAL TABLE`, not at the first + /// query against the table. + #[tokio::test] + async fn stored_as_atlas_refuses_a_bad_option_at_ddl() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + let ctx = ddl_context(1); + + let error = ctx + .sql(&format!( + "CREATE EXTERNAL TABLE t STORED AS ATLAS \ + OPTIONS ('use_pruning' 'maybe') LOCATION '{}/'", + location(tmp.path()) + )) + .await + .expect_err("'maybe' is no boolean") + .to_string(); + + assert!(error.contains("use_pruning"), "{error}"); + assert!(error.contains("maybe"), "{error}"); + } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs index ec760635..bd85634f 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs @@ -42,7 +42,9 @@ use datafusion::{ }, error::{DataFusionError, Result}, physical_expr::{ - PhysicalExpr, conjunction, projection::ProjectionExprs, utils::collect_columns, + PhysicalExpr, conjunction, + projection::ProjectionExprs, + utils::{collect_columns, reassign_expr_columns}, }, physical_plan::{ filter_pushdown::{FilterPushdownPropagation, PushedDown}, @@ -370,6 +372,14 @@ impl AtlasDatasets { return Arc::new(CandidateFilter::KeepAll); } + // The predicate is indexed against the table schema, the pruning + // engine reads it against the projected one. A projection pushed down + // after the filter leaves the two out of step, so the columns are + // re-indexed by name here. + let Ok(predicate) = reassign_expr_columns(predicate, &self.logical_schema) else { + return Arc::new(CandidateFilter::KeepAll); + }; + let started = Instant::now(); let atlas = Arc::clone(atlas); let schema = Arc::clone(&self.logical_schema); @@ -405,7 +415,9 @@ impl AtlasDatasets { /// changed. async fn projected_names(&self, view: &DatasetView) -> Result>> { if self.projected_schema.fields().is_empty() { - return Ok(count_driver(view).await?.map(|driver| vec![driver])); + return Ok(count_driver(view, self.read_dimensions.as_deref()) + .await? + .map(|driver| vec![driver])); } let mut names: Vec = self @@ -437,7 +449,15 @@ impl AtlasDatasets { /// `None` for a dataset with no readable array, and the caller then builds what /// there is — an attribute-only dataset contributes the one row its scalars /// define. -async fn count_driver(view: &DatasetView) -> Result> { +/// +/// `read_dimensions` rules out the arrays the narrowing after the build would +/// drop. The widest array of a dataset is the one on the most dimensions, and +/// a query that asked for fewer would lose exactly that one and leave the read +/// with no grid at all. +async fn count_driver( + view: &DatasetView, + read_dimensions: Option<&[String]>, +) -> Result> { let readable: Vec = view .schema() .iter() @@ -453,6 +473,14 @@ async fn count_driver(view: &DatasetView) -> Result> { view.name() )) })?; + if let Some(wanted) = read_dimensions + && !layout + .dimension_names() + .iter() + .all(|dim| wanted.iter().any(|kept| kept == dim)) + { + continue; + } let cells = layout.element_count(); if widest.as_ref().is_none_or(|(_, held)| cells > *held) { widest = Some((array, cells)); diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs index 2718b27f..00c2a2b6 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs @@ -201,9 +201,23 @@ pub fn project_read_dimensions( log_label: Option<&str>, ) -> anyhow::Result { match resolve_read_dimensions(&dataset, read_dimensions, log_label) { - Some(dims) => dataset - .project(&DatasetProjection::new_with_dimension_projection(dims)) - .map_err(|e| anyhow::anyhow!("Failed to project the atlas dataset by dimension: {e}")), + Some(dims) => { + // A projection names the dimensions of a whole collection, and one + // dataset need not hold them all. Narrowing to those it does hold + // drops the same arrays — an array survives only when every one of + // its own dimensions is kept, and a dimension this dataset lacks + // cannot be one of them — and leaves an all-scalar projection with + // nothing to narrow rather than an unknown dimension to refuse. + let held: Vec = dims + .into_iter() + .filter(|dim| dataset.dataset().dimensions.contains_key(dim)) + .collect(); + dataset + .project(&DatasetProjection::new_with_dimension_projection(held)) + .map_err(|e| { + anyhow::anyhow!("Failed to project the atlas dataset by dimension: {e}") + }) + } None => Ok(dataset), } } @@ -801,4 +815,134 @@ mod tests { "a 2-D array does not fit a 1-D grid: {columns:?}" ); } + // ── a wide collection on two dimensions ───────────────────────────── + + /// The shapes every test in this section reads: three datasets whose + /// profile and level counts all differ. + const WIDE: &[(usize, usize)] = &[(3, 4), (5, 4), (2, 6)]; + + /// A dataset per name the writer added, in write order. + #[tokio::test] + async fn a_wide_collection_lists_every_dataset() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + let atlas = test_support::open(tmp.path()).await; + + assert_eq!(atlas.list_datasets(), vec!["set0", "set1", "set2"]); + } + + /// The footer types every array without a segment read, so this is what one + /// open already knows. + #[tokio::test] + async fn every_dataset_of_a_wide_collection_declares_the_same_arrays() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + let atlas = test_support::open(tmp.path()).await; + + for name in atlas.list_datasets() { + let view = atlas.dataset(&name).expect("the dataset is listed"); + assert_eq!(view.schema().len(), 8, "{name} declares eight arrays"); + } + } + + /// The datasets share one interned schema, so the merge derives once. + #[tokio::test] + async fn a_wide_collection_types_every_column_from_the_footer() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + let atlas = test_support::open(tmp.path()).await; + let schema = collection_schema(&atlas, None, "wide", &widening()) + .await + .unwrap(); + + // Eight arrays, four attributes each, and two dataset attributes. + assert_eq!(schema.fields().len(), 42); + + let column = |name: &str| { + schema + .field_with_name(name) + .unwrap_or_else(|_| panic!("{name} is a column of the collection")) + .data_type() + .clone() + }; + // Atlas has a native timestamp, so `time` arrives as one. + assert_eq!( + column("time"), + DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None) + ); + assert_eq!(column("latitude"), DataType::Float64); + assert_eq!(column("temperature"), DataType::Float32); + assert_eq!(column("platform"), DataType::Utf8); + // An attribute is a column under `{array}.{attr}`, and a dataset + // attribute under `.{attr}`. + assert_eq!(column("temperature.units"), DataType::Utf8); + assert_eq!(column("temperature.valid_max"), DataType::Float64); + assert_eq!(column(".title"), DataType::Utf8); + } + + /// `profile` and `level` are the two dimensions. A list that holds only the + /// first keeps the arrays on it and drops the arrays that need both. + #[tokio::test] + async fn narrowing_to_one_dimension_drops_the_two_dimensional_arrays() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + let atlas = test_support::open(tmp.path()).await; + let dims = ["profile".to_string()]; + let schema = collection_schema(&atlas, Some(&dims), "wide", &widening()) + .await + .unwrap(); + + // The four grid arrays are gone; their attributes are scalars and stay. + assert_eq!(schema.fields().len(), 38); + for kept in ["latitude", "longitude", "time", "platform"] { + assert!( + schema.field_with_name(kept).is_ok(), + "{kept} is per profile" + ); + } + for dropped in ["pressure", "temperature", "salinity", "quality"] { + assert!( + schema.field_with_name(dropped).is_err(), + "{dropped} needs level as well" + ); + } + assert!(schema.field_with_name("temperature.units").is_ok()); + } + + /// A projection that names only attributes builds a dataset of scalars. + /// A narrowing by a dimension no surviving array carries must leave it + /// alone, and not refuse a dimension the dataset no longer has. + #[tokio::test] + async fn narrowing_an_attribute_only_dataset_keeps_its_scalars() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + let view = view(tmp.path(), "set0").await; + let dataset = dataset_from_view( + view, + Some(&["temperature.units".to_string(), ".title".to_string()]), + ) + .await + .unwrap(); + + let narrowed = + project_read_dimensions(dataset, Some(vec!["profile".to_string()]), None).unwrap(); + assert_eq!(names(&narrowed), vec![".title", "temperature.units"]); + } + + /// The values themselves, out of the first dataset. + #[tokio::test] + async fn wide_collection_array_values_read_back() { + let tmp = tempfile::tempdir().unwrap(); + test_support::wide_profiles(tmp.path(), WIDE).await; + let view = view(tmp.path(), "set0").await; + let dataset = dataset_from_view( + view, + Some(&["latitude".to_string(), "temperature".to_string()]), + ) + .await + .unwrap(); + + assert_eq!(dataset.get_array("latitude").unwrap().shape(), &[3]); + assert_eq!(dataset.get_array("temperature").unwrap().shape(), &[3, 4]); + } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs index 19d81c4d..8631b6b0 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs @@ -351,3 +351,121 @@ pub async fn empty(dir: &Path) { .expect("create the collection"); writer.finish().await.expect("finish the collection"); } + +/// A wide collection on two dimensions, `profile` and `level`. +/// +/// Dataset `i` is named `set{i}` and holds `shapes[i]` as +/// `(profiles, levels)`. The pairs differ per dataset, so a row count is a sum +/// and not a product. +/// +/// Four arrays live on `profile` alone: `latitude`, `longitude`, `time` and +/// `platform`. Four live on `profile` and `level`: `pressure`, `temperature`, +/// `salinity` and `quality`. So a read narrowed to `profile` keeps the first +/// four and drops the rest. +/// +/// Every array carries four attributes, so the collection is wide. A predicate +/// column then sits far past the length of a small projection, which is what a +/// projection pushed down after a filter has to be re-indexed against. +/// +/// `temperature` holds `100 * i + level`, so each dataset owns a disjoint +/// range and a predicate over it has a known answer. +pub async fn wide_profiles(dir: &Path, shapes: &[(usize, usize)]) { + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + + for (index, &(profiles, levels)) in shapes.iter().enumerate() { + let mut set = writer + .add_dataset(&format!("set{index}")) + .await + .expect("add a dataset"); + let flat = vec!["profile".to_string()]; + let grid = vec!["profile".to_string(), "level".to_string()]; + + set.define_array::("latitude", flat.clone(), vec![profiles], None, None) + .await + .expect("define latitude"); + set.define_array::("longitude", flat.clone(), vec![profiles], None, None) + .await + .expect("define longitude"); + set.define_array::("time", flat.clone(), vec![profiles], None, None) + .await + .expect("define time"); + set.define_array::("platform", flat, vec![profiles], None, None) + .await + .expect("define platform"); + for name in ["pressure", "temperature", "salinity"] { + set.define_array::(name, grid.clone(), vec![profiles, levels], None, None) + .await + .unwrap_or_else(|e| panic!("define {name}: {e}")); + } + set.define_array::("quality", grid, vec![profiles, levels], None, None) + .await + .expect("define quality"); + + let latitudes: Vec = (0..profiles).map(|p| 10.0 + p as f64).collect(); + let longitudes: Vec = (0..profiles).map(|p| 100.0 + p as f64).collect(); + let times: Vec = (0..profiles) + .map(|p| TimestampNs(EPOCH_NANOS + p as i64 * DAY_NANOS)) + .collect(); + let platforms: Vec = (0..profiles).map(|_| format!("set{index}")).collect(); + set.write_array("latitude", vec![0], arr1(&latitudes).into_dyn().view()) + .await + .expect("write latitude"); + set.write_array("longitude", vec![0], arr1(&longitudes).into_dyn().view()) + .await + .expect("write longitude"); + set.write_array("time", vec![0], arr1(×).into_dyn().view()) + .await + .expect("write time"); + set.write_array("platform", vec![0], arr1(&platforms).into_dyn().view()) + .await + .expect("write platform"); + + let shape = IxDyn(&[profiles, levels]); + let base = 100.0 * index as f32; + let pressure = ArrayD::from_shape_fn(shape.clone(), |i| i[1] as f32); + let temperature = ArrayD::from_shape_fn(shape.clone(), |i| base + i[1] as f32); + let salinity = ArrayD::from_shape_fn(shape.clone(), |i| 30.0 + i[1] as f32); + let quality = ArrayD::from_shape_fn(shape, |_| "1".to_string()); + set.write_array("pressure", vec![0, 0], pressure.view()) + .await + .expect("write pressure"); + set.write_array("temperature", vec![0, 0], temperature.view()) + .await + .expect("write temperature"); + set.write_array("salinity", vec![0, 0], salinity.view()) + .await + .expect("write salinity"); + set.write_array("quality", vec![0, 0], quality.view()) + .await + .expect("write quality"); + + // The attributes make the collection wide. Four per array, so the + // column count is well past the arrays alone. + for name in [ + "latitude", + "longitude", + "time", + "platform", + "pressure", + "temperature", + "salinity", + "quality", + ] { + set.set_array_attribute(name, "long_name", Attr::String(format!("the {name}"))) + .expect("set long_name"); + set.set_array_attribute(name, "units", Attr::String("1".into())) + .expect("set units"); + set.set_array_attribute(name, "valid_min", Attr::Float64(-1000.0)) + .expect("set valid_min"); + set.set_array_attribute(name, "valid_max", Attr::Float64(1000.0)) + .expect("set valid_max"); + } + set.set_attribute("title", Attr::String("wide profiles".into())); + set.set_attribute("institution", Attr::String("test".into())); + set.finish().await.expect("finish a dataset"); + } + + writer.finish().await.expect("finish the collection"); +} From f5d78163795feb6ed3f4c9dd45c1b3ab8d01ff8c Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 8 Sep 2026 14:13:04 +0200 Subject: [PATCH 05/16] wip --- .../beacon-datafusion-ext/src/nd/encoding.rs | 24 +- beacon-db/beacon-datafusion-ext/src/nd/mod.rs | 2 +- .../beacon-datafusion-ext/src/scan_adapt.rs | 35 +- .../beacon-arrow-atlas/src/config.rs | 22 - .../src/datafusion/metrics.rs | 9 - .../beacon-arrow-atlas/src/datafusion/mod.rs | 532 +++++++++----- .../src/datafusion/pruning.rs | 40 -- .../src/datafusion/source.rs | 658 ++++++++---------- .../src/datafusion/statistics.rs | 296 -------- .../beacon-arrow-atlas/src/lib.rs | 12 +- .../beacon-arrow-atlas/src/reader.rs | 121 +--- .../beacon-arrow-atlas/src/test_support.rs | 79 +++ beacon-server/beacon-server-config/src/lib.rs | 6 - 13 files changed, 804 insertions(+), 1032 deletions(-) delete mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs diff --git a/beacon-db/beacon-datafusion-ext/src/nd/encoding.rs b/beacon-db/beacon-datafusion-ext/src/nd/encoding.rs index dfab0335..af0802e5 100644 --- a/beacon-db/beacon-datafusion-ext/src/nd/encoding.rs +++ b/beacon-db/beacon-datafusion-ext/src/nd/encoding.rs @@ -17,7 +17,6 @@ //! normal `RecordBatch`; [`NdSourceExec`](crate::nd::exec::NdSourceExec) decodes //! it back into an [`NdRecordBatch`] on the way out. -use std::collections::HashMap; use std::sync::Arc; use arrow::array::{ @@ -69,11 +68,22 @@ pub fn nd_encoded_type(value_type: &DataType) -> DataType { /// An nd column field named `name` carrying values of `value_type`, tagged with /// the `beacon.nd` extension type. pub fn nd_encoded_field(name: &str, value_type: &DataType) -> Field { - let metadata = HashMap::from([ - ("ARROW:extension:name".to_string(), ND_EXTENSION_NAME.to_string()), - ("ARROW:extension:metadata".to_string(), "{}".to_string()), - ]); - Field::new(name, nd_encoded_type(value_type), true).with_metadata(metadata) + nd_encoded_field_of(&Field::new(name, value_type.clone(), true)) +} + +/// `field`, with its values carried as a `beacon.nd` struct. +/// +/// The field's own metadata travels with it, and the extension keys go on top. +/// That matters because a scan's target schema is the *encoded* one: a column +/// the merge marked with +/// [`TYPE_CONFLICT_KEY`](crate::type_widening::TYPE_CONFLICT_KEY) has to stay +/// marked. Without the mark `scan_adapt` casts strictly, and a value the merged +/// type cannot hold fails the scan instead of reading as null. +pub fn nd_encoded_field_of(field: &Field) -> Field { + let mut metadata = field.metadata().clone(); + metadata.insert("ARROW:extension:name".to_string(), ND_EXTENSION_NAME.to_string()); + metadata.insert("ARROW:extension:metadata".to_string(), "{}".to_string()); + Field::new(field.name(), nd_encoded_type(field.data_type()), true).with_metadata(metadata) } /// The nd-encoded schema of a logical schema: every field becomes a `beacon.nd` @@ -83,7 +93,7 @@ pub fn encoded_schema(logical: &Schema) -> Schema { let fields: Vec = logical .fields() .iter() - .map(|f| nd_encoded_field(f.name(), f.data_type())) + .map(|f| nd_encoded_field_of(f)) .collect(); Schema::new_with_metadata(fields, logical.metadata().clone()) } diff --git a/beacon-db/beacon-datafusion-ext/src/nd/mod.rs b/beacon-db/beacon-datafusion-ext/src/nd/mod.rs index 7cee86e4..a0ea99f0 100644 --- a/beacon-db/beacon-datafusion-ext/src/nd/mod.rs +++ b/beacon-db/beacon-datafusion-ext/src/nd/mod.rs @@ -31,7 +31,7 @@ pub use broadcast::BroadcastMap; pub use dimensions::{Dimension, Dimensions}; pub use encoding::{ decode_nd_record_batch, encode_flat_batch_as_nd, encode_nd_record_batch, encoded_schema, - is_nd_encoded, logical_schema, nd_encoded_field, nd_encoded_type, + is_nd_encoded, logical_schema, nd_encoded_field, nd_encoded_field_of, nd_encoded_type, }; pub use optimizer::{NdFilterPushdown, NdProjectionPushdown, is_pushable_expr}; diff --git a/beacon-db/beacon-datafusion-ext/src/scan_adapt.rs b/beacon-db/beacon-datafusion-ext/src/scan_adapt.rs index 60dd62a5..d87f31d4 100644 --- a/beacon-db/beacon-datafusion-ext/src/scan_adapt.rs +++ b/beacon-db/beacon-datafusion-ext/src/scan_adapt.rs @@ -28,8 +28,17 @@ //! - A type no cast reaches reads as null for the whole file. A list beside a //! number is one such pair. //! -//! Every other cast stays strict, and a value it cannot hold is an error. The -//! merged schema carries the mark, so no scan reads the setting itself. +//! The merged schema carries the mark, so no scan reads the setting itself. +//! +//! # An nd column +//! +//! An nd column reads leniently too, whether or not the merge marked it. Its +//! cast lands on the `values` list inside the `beacon.nd` struct, and one +//! collection of a million datasets may store an array as text where another +//! stores numbers. One cell that does not parse would otherwise fail the whole +//! scan, and no single dataset is worth a collection. +//! +//! Every other cast stays strict, and a value it cannot hold is an error. //! //! [`TypeConflict::KeepFirst`]: crate::type_widening::TypeConflict::KeepFirst @@ -52,8 +61,24 @@ use datafusion::physical_expr_adapter::{ }; use futures::StreamExt; +use crate::nd::is_nd_encoded; use crate::type_widening::is_type_conflict; +/// Whether a cast onto `field` may read a value the type cannot hold as null. +/// +/// Two cases qualify. +/// +/// A column the merge could not join, marked by the widening rule. The sources +/// state two families, so no value of the other family is a value of this one. +/// +/// An nd column. Its cast lands on the `values` list inside the `beacon.nd` +/// struct, and a collection of a million datasets may store one array as text +/// where another stores numbers. One cell that does not parse must not fail the +/// whole scan, because no single dataset is worth the collection. +fn casts_leniently(field: &arrow::datatypes::Field) -> bool { + is_type_conflict(field) || is_nd_encoded(field) +} + /// Where one column of the target schema comes from. #[derive(Debug, Clone)] enum Source { @@ -114,7 +139,7 @@ impl BatchAdapter { // A column the merge could not join reads null where the cast // cannot answer, and null for the whole file where no cast // reaches its type. - Ok(at) if is_type_conflict(field) => { + Ok(at) if casts_leniently(field) => { let data_type = field.data_type().clone(); Ok( if can_cast_types(source.field(at).data_type(), &data_type) { @@ -286,7 +311,7 @@ impl LenientCastAdapter { fn conflicted(&self, column: &Column) -> Option<&arrow::datatypes::Field> { let at = self.logical_file_schema.index_of(column.name()).ok()?; let field = self.logical_file_schema.field(at); - is_type_conflict(field).then_some(field) + casts_leniently(field).then_some(field) } } @@ -323,7 +348,7 @@ impl PhysicalExprAdapter for LenientCastAdapter { let Some(cast) = expr.as_any().downcast_ref::() else { return Ok(Transformed::no(expr)); }; - if !is_type_conflict(cast.target_field()) { + if !casts_leniently(cast.target_field()) { return Ok(Transformed::no(expr)); } Ok(Transformed::yes(Arc::new(CastColumnExpr::new( diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs index 77ed9caf..e0a831a3 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs @@ -10,30 +10,9 @@ #[derive(Debug, Clone)] pub struct AtlasConfig { /// Whether a read consults the shared reader cache. - /// - /// A collection is immutable, so a cached handle stays valid until its - /// deletion mask changes. The cache saves the footer read and keeps the - /// decompressed blocks of a collection between queries. pub use_reader_cache: bool, - /// How many opened collections the shared reader cache holds. - /// - /// Each entry owns its own block cache — 256 MiB of decompressed blocks and - /// 64 MiB of raw slabs — so this is also a memory bound: the default of 32 - /// admits up to 10 GiB of cached blocks across every open collection. pub reader_cache_size: u64, - /// Whether a predicate scan drops the datasets that cannot match, before it - /// reads them. - /// - /// A pure optimization: pruning only ever removes datasets that hold no - /// matching row, and every path fails open. Off trades throughput for - /// skipping the index build. pub use_pruning: bool, - /// Whether the file analyzer measures a collection's column ranges. - /// - /// The ranges come from the footer, so they cost no array read. The switch - /// exists because a listing of many collections turns even a footer read - /// per collection into real I/O. - pub enable_statistics: bool, } impl Default for AtlasConfig { @@ -42,7 +21,6 @@ impl Default for AtlasConfig { use_reader_cache: true, reader_cache_size: 32, use_pruning: true, - enable_statistics: true, } } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs index 4365536d..c2cef023 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs @@ -22,16 +22,7 @@ pub struct AtlasScanMetrics { /// Wall time opening collections, or hitting the reader cache for them. pub open_time: Time, /// Wall time deciding which datasets a predicate can rule out. - /// - /// One partition builds a collection's index and the rest wait on it, so - /// this is the build for one of them and the wait for the others. pub prune_time: Time, - /// Wall time building lazy datasets: resolving the view, reading the - /// projected attribute values out of the footer, wiring the backends, and - /// planning the chunk queue. - /// - /// Array data is read later, as the queue is drained, and `ReadMetrics` - /// counts that. pub dataset_build_time: Time, /// Datasets this partition opened and read. pub datasets_scanned: Count, diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs index 635d8b06..9d2c272d 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs @@ -36,19 +36,16 @@ use object_store::{ObjectMeta, ObjectStore}; use crate::config::AtlasConfig; use crate::reader::collection_schema; -use crate::store::{ - ATLAS_MARKER, AtlasReaderCache, get_or_open_atlas, is_atlas_marker, top_level_atlas_markers, -}; +use crate::store::{ATLAS_MARKER, AtlasReaderCache, get_or_open_atlas, top_level_atlas_markers}; pub mod metrics; pub mod options; pub mod pruning; pub mod source; -pub mod statistics; pub mod table_function; pub use options::AtlasOptions; -pub use source::{AtlasEntry, AtlasSource}; +pub use source::AtlasSource; pub use table_function::ReadAtlasFunc; /// The name this format answers to: `STORED AS ATLAS`, `read_atlas`. @@ -90,26 +87,16 @@ impl AtlasFormatFactory { /// A format with this table's effective settings, wired to the shared cache /// when caching is on. - fn build( + pub(crate) fn build( &self, options: AtlasOptions, use_reader_cache: bool, use_pruning: bool, ) -> AtlasFormat { - AtlasFormat::new(options) - .with_cache(use_reader_cache.then(|| self.cache.clone())) - .with_pruning(use_pruning) - } - - /// Whether this table wants its columns measured at all. - /// - /// Only the file analyzer measures a collection, through - /// [`FileFormatFactoryExt::create_for_analysis`]. This is the switch that - /// turns even that off, per table or per runtime. - fn statistics_wanted(&self, format_options: &HashMap) -> Result { - match format_option(format_options, "enable_statistics") { - Some(value) => parse_bool_option("enable_statistics", value), - None => Ok(self.config.enable_statistics), + AtlasFormat { + options, + cache: use_reader_cache.then(|| self.cache.clone()), + use_pruning, } } } @@ -139,11 +126,6 @@ impl FileFormatFactory for AtlasFormatFactory { if let Some(value) = format_option(format_options, "use_pruning") { use_pruning = parse_bool_option("use_pruning", value)?; } - // Parsed here only so a bad value is an error at `CREATE EXTERNAL - // TABLE` rather than at the first analysis pass. A query measures - // nothing whatever it says: see `create_for_analysis`. - self.statistics_wanted(format_options)?; - Ok(Arc::new(self.build(options, use_reader_cache, use_pruning))) } @@ -184,11 +166,6 @@ impl FileFormatFactoryExt for AtlasFormatFactory { } /// One schema per collection, not per object. - /// - /// `infer_schema` reads the container and derives the schema from the - /// footer inside it, so the entry is keyed on the container and depends on - /// everything beside it — the deletion mask included, which changes which - /// datasets the schema covers. fn schema_units(&self, objects: &[ObjectMeta]) -> Vec { units_over_stores(objects, &top_level_atlas_markers(objects)) } @@ -207,29 +184,18 @@ impl FileFormatFactoryExt for AtlasFormatFactory { Some(SchemaOptions::new(ATLAS_FORMAT).finish()) } - /// The same format, with the column measurement switched on. + /// The plain format. Atlas measures no column. /// - /// `infer_stats` folds a collection's footer, which is cheap but not free - /// over a listing of thousands. Only the file analyzer asks for it, and a - /// scan prunes from what that recorded. See - /// [`FileFormatFactoryExt::create_for_analysis`]. + /// The analyzer asks a format to measure a collection, and this one reports + /// unknown for every column. See [`AtlasFormat::infer_stats`]. fn create_for_analysis( &self, - state: &dyn Session, - format_options: &HashMap, - url: &ListingTableUrl, - listing: &ListingFactory, + _state: &dyn Session, + _format_options: &HashMap, + _url: &ListingTableUrl, + _listing: &ListingFactory, ) -> Result> { - let wanted = self.statistics_wanted(format_options)?; - let format = self.create_with_native_root(state, format_options, url, listing)?; - let atlas = format - .as_any() - .downcast_ref::() - .ok_or_else(|| { - exec_datafusion_err!("the atlas factory did not produce an AtlasFormat") - })? - .clone(); - Ok(Arc::new(atlas.with_enable_statistics(wanted))) + Ok(Arc::new(AtlasFormat::default())) } } @@ -243,8 +209,6 @@ pub struct AtlasFormat { cache: Option, /// Whether a predicate scan drops the datasets it can rule out. use_pruning: bool, - /// Whether [`FileFormat::infer_stats`] measures a collection's columns. - enable_statistics: bool, } impl Default for AtlasFormat { @@ -261,30 +225,8 @@ impl AtlasFormat { cache: None, // A query prunes by default: it only ever saves reads. use_pruning: defaults.use_pruning, - // A query measures nothing. Only the analyzer asks, through - // `create_for_analysis`. - enable_statistics: false, } } - - /// Wire in a reader cache (`Some`), or bypass caching (`None`). - pub fn with_cache(mut self, cache: Option) -> Self { - self.cache = cache; - self - } - - /// Drop the datasets a predicate rules out, or read them all. - pub fn with_pruning(mut self, use_pruning: bool) -> Self { - self.use_pruning = use_pruning; - self - } - - /// Measure a collection's columns in [`FileFormat::infer_stats`], or report - /// them unknown. - pub fn with_enable_statistics(mut self, enable_statistics: bool) -> Self { - self.enable_statistics = enable_statistics; - self - } } /// Wrap a scan in the nd spine: `NdBroadcastExec` over `NdSourceExec` over the @@ -373,36 +315,29 @@ impl FileFormat for AtlasFormat { Ok(schema) } - /// The column ranges of one collection, folded out of its footer. + /// Unknown for every column. /// - /// Reporting unknown rather than erroring is deliberate throughout: absent - /// statistics are always a legal answer, and they only mean Beacon reads - /// what it might have skipped. A listing also hands this method every - /// object it matched, and only a container has a collection behind it. + /// Atlas measures nothing for the analyzer. A recorded range prunes whole + /// collections before a scan opens them, so a range that is too narrow + /// deletes matching rows from an answer. The scan prunes from the footer + /// itself instead, per dataset, where the numbers are exact and cost no + /// array read. See [`pruning`]. async fn infer_stats( &self, _state: &dyn Session, - store: &Arc, + _store: &Arc, table_schema: SchemaRef, - object: &ObjectMeta, + _object: &ObjectMeta, ) -> Result { - if !self.enable_statistics || !is_atlas_marker(object) { - return Ok(Statistics::new_unknown(&table_schema)); - } - - match get_or_open_atlas(self.cache.as_ref(), Arc::clone(store), object).await { - Ok(atlas) => Ok(statistics::collection_statistics(&atlas, &table_schema).await), - Err(e) => { - tracing::debug!(object = %object.location, "not measuring this collection: {e}"); - Ok(Statistics::new_unknown(&table_schema)) - } - } + Ok(Statistics::new_unknown(&table_schema)) } - /// Plan one entry per dataset, then wrap the scan in the nd spine. + /// Plan one entry per collection, then wrap the scan in the nd spine. /// - /// Opening each collection here is one metadata read, and the openers reuse - /// the same handle through the reader cache. `list_datasets` is in memory. + /// Nothing is opened here. The markers the listing found are deduped to the + /// outermost collections and dealt round-robin over the target partitions, + /// and each partition's opener lists, prunes and reads the collections it + /// holds. Parallelism is therefore bounded by the collection count. async fn create_physical_plan( &self, state: &dyn Session, @@ -410,11 +345,6 @@ impl FileFormat for AtlasFormat { ) -> Result> { beacon_nd_array::arrow::morsel::reject_partition_columns("Atlas", &conf)?; - let started = std::time::Instant::now(); - let object_store = state - .runtime_env() - .object_store(conf.object_store_url.clone())?; - let listed: Vec = conf .file_groups .iter() @@ -423,32 +353,24 @@ impl FileFormat for AtlasFormat { .collect(); let markers = top_level_atlas_markers(&listed); - let mut datasets = 0usize; - let mut file_groups: Vec = Vec::with_capacity(markers.len()); - for marker in &markers { - let atlas = get_or_open_atlas(self.cache.as_ref(), Arc::clone(&object_store), marker) - .await - .map_err(|e| exec_datafusion_err!("{e}"))?; - let names = atlas.list_datasets(); - datasets += names.len(); - - let entries: Vec = names - .into_iter() - .enumerate() - .map(|(position, dataset)| { - // `From` keeps the container's freshness, so - // the opener's cache key matches this plan-time open. - let mut entry = PartitionedFile::from(marker.clone()); - entry.extensions = Some(Arc::new(AtlasEntry { dataset, position })); - entry - }) - .collect(); - file_groups.push(FileGroup::new(entries)); + // One collection is one unit of work, and a container is never split, + // so the deal here is the whole distribution. + let partitions = state + .config() + .target_partitions() + .clamp(1, markers.len().max(1)); + let mut dealt: Vec> = vec![Vec::new(); partitions]; + for (index, marker) in markers.iter().enumerate() { + dealt[index % partitions].push(PartitionedFile::from(marker.clone())); } + let file_groups: Vec = dealt + .into_iter() + .filter(|group| !group.is_empty()) + .map(FileGroup::new) + .collect(); tracing::debug!( - elapsed_ms = started.elapsed().as_millis() as u64, collections = markers.len(), - datasets, + partitions = file_groups.len(), "atlas create_physical_plan", ); @@ -706,7 +628,7 @@ mod tests { #[tokio::test] async fn a_widened_column_is_cast_from_each_dataset() { - use arrow::array::Float64Array; + use arrow::array::{Array, Float64Array}; let tmp = tempfile::tempdir().unwrap(); test_support::widening(tmp.path()).await; @@ -807,9 +729,10 @@ mod tests { } } - /// The scan is planned across every partition, through the queue. + /// The unit of work is the collection, so one collection is one partition + /// however many datasets it holds. #[tokio::test] - async fn a_collection_is_planned_across_every_partition() { + async fn one_collection_is_one_partition() { use datafusion::physical_plan::ExecutionPlanProperties; let tmp = tempfile::tempdir().unwrap(); @@ -833,12 +756,58 @@ mod tests { } assert_eq!( scan.output_partitioning().partition_count(), - 4, - "the datasets divide over the partitions:\n{}", + 1, + "twelve datasets of one collection stay together: +{}", datafusion::physical_plan::displayable(plan.as_ref()).indent(false) ); } + /// Several collections deal round-robin over the partitions, and never + /// past the collection count. + #[tokio::test] + async fn collections_deal_across_the_partitions() { + use datafusion::physical_plan::ExecutionPlanProperties; + + let tmp = tempfile::tempdir().unwrap(); + for name in ["a", "b", "c"] { + test_support::ranged(&tmp.path().join(name), 2).await; + } + let root = tmp + .path() + .to_string_lossy() + .replace(std::path::MAIN_SEPARATOR, "/"); + + for (target, expected) in [(8, 3), (2, 2), (1, 1)] { + let ctx = ddl_context(target); + ctx.sql(&format!( + "CREATE EXTERNAL TABLE t STORED AS ATLAS LOCATION '{root}/**/data.atlas'" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let plan = ctx + .sql("SELECT temperature FROM t") + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let mut scan = Arc::clone(&plan); + while let Some(child) = scan.children().first() { + scan = Arc::clone(child); + } + assert_eq!( + scan.output_partitioning().partition_count(), + expected, + "three collections over {target} target partitions" + ); + } + } + // ── predicates ────────────────────────────────────────────────────── #[tokio::test] @@ -870,11 +839,23 @@ mod tests { // ── pruning, end to end ───────────────────────────────────────────── + /// A format with pruning on or off, built the way a table is. + /// + /// `use_pruning` reaches a format through the factory alone, as + /// `CREATE EXTERNAL TABLE ... OPTIONS ('use_pruning' '...')` does. + fn format_with_pruning(use_pruning: bool) -> Arc { + Arc::new( + AtlasFormatFactory::new(Default::default(), Default::default()).build( + AtlasOptions::default(), + false, + use_pruning, + ), + ) + } + /// Register the collection twice, once pruning and once not. async fn register_pruning(ctx: &SessionContext, dir: &Path, name: &str, use_pruning: bool) { - let format: Arc = - Arc::new(AtlasFormat::default().with_pruning(use_pruning)); - register_with(ctx, dir, name, format).await; + register_with(ctx, dir, name, format_with_pruning(use_pruning)).await; } async fn values(ctx: &SessionContext, sql: &str) -> Vec { @@ -1039,54 +1020,58 @@ mod tests { // ── measuring a collection ────────────────────────────────────────── - /// A query never measures a collection, whatever the option says. Only the - /// analyzer asks, through `create_for_analysis`. - #[test] - fn a_query_never_measures_a_collection() { - let factory = AtlasFormatFactory::new(Default::default(), Default::default()); + /// Atlas measures no column, for any caller. + /// + /// A recorded range prunes whole collections before a scan opens them, so a + /// range that is too narrow deletes matching rows from an answer. The scan + /// prunes from the footer instead, per dataset, where the numbers are exact + /// and cost no array read. + #[tokio::test] + async fn a_collection_reports_no_column_range() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 4).await; + let (store, marker) = test_support::store_and_marker(tmp.path()); let ctx = SessionContext::new(); - let on = HashMap::from([("enable_statistics".to_string(), "true".to_string())]); - for options in [HashMap::new(), on] { - let format = factory.create(&ctx.state(), &options).unwrap(); + let format = AtlasFormat::default(); + let schema = format + .infer_schema(&ctx.state(), &store, std::slice::from_ref(&marker)) + .await + .unwrap(); + let statistics = format + .infer_stats(&ctx.state(), &store, Arc::clone(&schema), &marker) + .await + .unwrap(); + + assert_eq!(statistics.column_statistics.len(), schema.fields().len()); + for (column, field) in statistics.column_statistics.iter().zip(schema.fields()) { assert!( - !format - .as_any() - .downcast_ref::() - .unwrap() - .enable_statistics, - "a format built for a query measures nothing" + column.min_value.is_exact().is_none(), + "{} reports a minimum", + field.name() + ); + assert!( + column.max_value.is_exact().is_none(), + "{} reports a maximum", + field.name() ); } + assert!(statistics.num_rows.is_exact().is_none()); } - /// The analyzer layers the per-table option over the runtime default. + /// The analyzer gets the same format a query does, because neither + /// measures anything. #[test] - fn analysis_layers_the_statistics_option_over_the_runtime() { - let measured = |config: AtlasConfig, options: HashMap| { - let ctx = SessionContext::new(); - let listing = Arc::new(ListingFactory::dynamic()); - let url = ListingTableUrl::parse("file:///tmp/").unwrap(); - AtlasFormatFactory::new(Default::default(), config) - .create_for_analysis(&ctx.state(), &options, &url, &listing) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .enable_statistics - }; - let off = HashMap::from([("enable_statistics".to_string(), "false".to_string())]); - let on = HashMap::from([("enable_statistics".to_string(), "yes".to_string())]); - - assert!(measured(AtlasConfig::default(), HashMap::new())); - assert!(!measured(AtlasConfig::default(), off)); + fn analysis_asks_for_no_measurement() { + let ctx = SessionContext::new(); + let listing = Arc::new(ListingFactory::dynamic()); + let url = ListingTableUrl::parse("file:///tmp/").unwrap(); + let factory = AtlasFormatFactory::new(Default::default(), Default::default()); - let disabled = AtlasConfig { - enable_statistics: false, - ..Default::default() - }; - assert!(!measured(disabled.clone(), HashMap::new())); - assert!(measured(disabled, on), "one table can turn them back on"); + let analysis = factory + .create_for_analysis(&ctx.state(), &HashMap::new(), &url, &listing) + .unwrap(); + assert!(analysis.as_any().downcast_ref::().is_some()); } // ── dimensions ────────────────────────────────────────────────────── @@ -1359,7 +1344,7 @@ mod tests { let pruned = context(4); let _a = wide_table(&pruned, Arc::new(AtlasFormat::default())).await; let whole = context(4); - let _b = wide_table(&whole, Arc::new(AtlasFormat::default().with_pruning(false))).await; + let _b = wide_table(&whole, format_with_pruning(false)).await; for sql in queries { assert_eq!( @@ -1738,4 +1723,211 @@ mod tests { assert!(error.contains("use_pruning"), "{error}"); assert!(error.contains("maybe"), "{error}"); } + /// A session whose merge keeps the first type of a conflicting column. + /// + /// The default refuses such a column outright, so a test of the cast has to + /// ask for `KeepFirst`. + fn keep_first_context(partitions: usize) -> SessionContext { + use beacon_datafusion_ext::type_widening::{ArrowTypeWidening, DefaultArrowTypeWidening}; + + SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(partitions) + .with_extension(Arc::new(ArrowTypeWidening::new(Arc::new( + DefaultArrowTypeWidening::keeping_first_type(), + )))), + ) + } + + /// A value the merged type cannot hold reads as null, and does not fail the + /// query. + /// + /// `KeepFirst` keeps `Float64` and marks the column. The mark has to + /// survive the nd encoding, because the scan's target schema is the encoded + /// one. Without it `scan_adapt` casts strictly and `'0.-90'` fails the whole + /// scan. + /// + /// `'3.5'` casts cleanly beside it, so this separates "one value is null" + /// from "the column gave up". + #[tokio::test] + async fn a_value_the_merged_type_cannot_hold_reads_as_null() { + use arrow::array::{Array, Float64Array}; + + let tmp = tempfile::tempdir().unwrap(); + test_support::conflicting_numbers(tmp.path()).await; + let ctx = keep_first_context(1); + register(&ctx, tmp.path(), "t").await; + + let batches = ctx + .sql("SELECT value FROM t") + .await + .unwrap() + .collect() + .await + .expect("an unparseable value is null, not an error"); + + let mut seen: Vec> = Vec::new(); + for batch in &batches { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("the merge kept the first type"); + for row in 0..column.len() { + seen.push((!column.is_null(row)).then(|| column.value(row))); + } + } + seen.sort_by(|a, b| a.partial_cmp(b).unwrap()); + assert_eq!(seen, vec![None, Some(1.5), Some(2.5), Some(3.5)]); + } + + /// The mark itself, on the schema the scan carries. + /// + /// The logical schema is marked, and the encoded schema has to stay marked. + /// This is the step that used to drop it. + #[test] + fn the_nd_encoding_keeps_the_type_conflict_mark() { + use beacon_datafusion_ext::nd::encoded_schema; + use beacon_datafusion_ext::type_widening::{ + TYPE_CONFLICT_FIRST_TYPE, TYPE_CONFLICT_KEY, is_type_conflict, + }; + + let marked = arrow::datatypes::Field::new("value", DataType::Float64, true).with_metadata( + std::collections::HashMap::from([( + TYPE_CONFLICT_KEY.to_string(), + TYPE_CONFLICT_FIRST_TYPE.to_string(), + )]), + ); + let plain = arrow::datatypes::Field::new("other", DataType::Float64, true); + let logical = Schema::new(vec![marked, plain]); + + let encoded = encoded_schema(&logical); + assert!( + is_type_conflict(encoded.field_with_name("value").unwrap()), + "the encoded field lost the mark" + ); + assert!( + !is_type_conflict(encoded.field_with_name("other").unwrap()), + "an unmarked column must not gain the mark" + ); + // The extension tag still identifies the column as nd-encoded. + assert!(beacon_datafusion_ext::nd::is_nd_encoded( + encoded.field_with_name("value").unwrap() + )); + } + /// What the *default* widening does with Float64 beside Utf8. + #[tokio::test] + async fn probe_default_widening() { + let tmp = tempfile::tempdir().unwrap(); + test_support::conflicting_numbers(tmp.path()).await; + + // Default session: no widening extension registered at all. + let ctx = context(1); + let dir = tmp + .path() + .to_string_lossy() + .replace(std::path::MAIN_SEPARATOR, "/"); + let url = ListingTableUrl::parse(format!("file://{dir}/")).unwrap(); + let format: Arc = Arc::new(AtlasFormat::default()); + let listing = ListingOptions::new(format).with_file_extension(ATLAS_MARKER); + match ListingTableConfig::new(url) + .with_listing_options(listing) + .infer_schema(&ctx.state()) + .await + { + Ok(config) => { + let schema = config.file_schema.clone().unwrap(); + let field = schema.field_with_name("value").unwrap(); + println!("PROBE default infer OK type={:?}", field.data_type()); + println!( + "PROBE default marked={}", + beacon_datafusion_ext::type_widening::is_type_conflict(field) + ); + ctx.register_table("t", Arc::new(ListingTable::try_new(config).unwrap())) + .unwrap(); + match ctx + .sql("SELECT value FROM t") + .await + .unwrap() + .collect() + .await + { + Ok(b) => println!( + "PROBE default query OK\n{}", + arrow::util::pretty::pretty_format_batches(&b).unwrap() + ), + Err(e) => println!("PROBE default query ERR {e}"), + } + } + Err(e) => println!("PROBE default infer ERR {e}"), + } + } + + /// The same value, with no mark on the column at all. + /// + /// An nd cast lands on the `values` list inside the `beacon.nd` struct, and + /// it reads leniently whether or not the merge marked the column. A table + /// whose schema says `Float64` therefore survives a dataset that stores the + /// array as text, however the schema came to say so. + #[tokio::test] + async fn an_unparseable_nd_value_reads_as_null_without_a_mark() { + use arrow::array::{Array, Float64Array}; + + let tmp = tempfile::tempdir().unwrap(); + test_support::value_typed(&tmp.path().join("one"), false).await; + test_support::value_typed(&tmp.path().join("two"), true).await; + + // Two paths, so the table schema comes from the first collection alone + // and carries no mark. That is the shape a scan cannot rely on one. + let ctx = keep_first_context(4); + let urls: Vec = ["one", "two"] + .iter() + .map(|name| { + let dir = tmp + .path() + .join(name) + .to_string_lossy() + .replace(std::path::MAIN_SEPARATOR, "/"); + ListingTableUrl::parse(format!("file://{dir}/")).unwrap() + }) + .collect(); + let format: Arc = Arc::new(AtlasFormat::default()); + let listing = ListingOptions::new(format).with_file_extension(ATLAS_MARKER); + let config = ListingTableConfig::new_with_multi_paths(urls) + .with_listing_options(listing) + .infer_schema(&ctx.state()) + .await + .expect("the collections type"); + let schema = config.file_schema.clone().unwrap(); + let field = schema.field_with_name("value").unwrap(); + assert_eq!(field.data_type(), &DataType::Float64); + assert!( + !beacon_datafusion_ext::type_widening::is_type_conflict(field), + "this shape is the one with no mark; the test is pointless with one" + ); + + ctx.register_table("t", Arc::new(ListingTable::try_new(config).unwrap())) + .unwrap(); + let batches = ctx + .sql("SELECT value FROM t") + .await + .unwrap() + .collect() + .await + .expect("an unparseable value is null, not an error"); + + let mut seen: Vec> = Vec::new(); + for batch in &batches { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("the table type"); + for row in 0..column.len() { + seen.push((!column.is_null(row)).then(|| column.value(row))); + } + } + seen.sort_by(|a, b| a.partial_cmp(b).unwrap()); + assert_eq!(seen, vec![None, Some(1.5), Some(2.5), Some(3.5)]); + } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs index 1909e73a..d1336804 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs @@ -94,46 +94,6 @@ impl CandidateFilter { } } -/// Each collection's [`CandidateFilter`], computed once per scan. -/// -/// Keyed by the container's path. The predicate and the schema are fixed for a -/// scan, so the container identifies the answer. Every partition's opener holds -/// a clone of this cache, and the clones share one store: the first opener to -/// reach a collection builds the index while the rest await the same future. -#[derive(Clone)] -pub struct PruneCache { - cache: moka::future::Cache>, -} - -impl PruneCache { - pub fn new() -> Self { - Self { - cache: moka::future::Cache::builder().max_capacity(256).build(), - } - } - - /// The memoized filter for `key`, computing it with `init` on first use. - /// Concurrent callers for one key share the one computation. - pub async fn get_or_compute(&self, key: String, init: F) -> Arc - where - F: std::future::Future>, - { - self.cache.get_with(key, init).await - } -} - -impl Default for PruneCache { - fn default() -> Self { - Self::new() - } -} - -impl std::fmt::Debug for PruneCache { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("PruneCache").finish_non_exhaustive() - } -} - // ─── The index ─────────────────────────────────────────────────────────────── /// One column's statistics, one row per dataset. diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs index bd85634f..38b8a3c6 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs @@ -1,23 +1,23 @@ //! The DataFusion [`FileSource`] and [`FileOpener`] for Atlas collections. //! -//! # One dataset is one unit of work +//! # One collection is one unit of work //! -//! A plan entry is a *dataset*, not a collection: [`AtlasFormat`] lists each -//! collection at plan time and emits one [`PartitionedFile`] per dataset, with -//! the dataset's name in [`PartitionedFile::extensions`]. Every entry carries -//! the collection's own marker, so the opener knows which container to open and -//! the reader cache keys on the same object the plan did. +//! A plan entry is a collection: its `data.atlas` container, as the listing +//! found it. [`AtlasFormat`] dedupes the markers and deals them round-robin +//! over the target partitions, and a partition reads each collection it holds +//! from end to end. A container is never split by byte range, because a byte +//! range of one means nothing. //! -//! Those entries go into one [`MorselSource`], and each partition holds a -//! standing entry pointing at it. A partition takes the next dataset when it is -//! free, and helps drain an open one when none is left. Balance follows -//! completion, so a collection of a million small datasets and a collection of -//! four large ones both divide over every core. +//! # What an open does //! -//! The second level of the queue is the dataset's own chunk grid, the one the -//! writer chose. A dataset stored as a single chunk is one unit of work; a -//! chunked one is as many as it has chunks, and several partitions drain it -//! together. Each unit becomes one nd record batch. +//! Opening a collection costs one footer read through the reader cache. The +//! opener then lists the datasets, prunes them in one vectorised pass over the +//! footer's statistics, and is left with the names it has to read. Those names +//! feed one stream: each dataset is built in turn, planned as a [`FileRead`], +//! and its batches are yielded before the next dataset is touched. +//! +//! So a pruned dataset costs nothing at all, and a kept one costs its build and +//! its read. Nothing is listed at plan time, and nothing is queued. //! //! [`AtlasFormat`]: super::AtlasFormat @@ -28,10 +28,7 @@ use std::time::Instant; use arrow::datatypes::SchemaRef; use atlas::{Atlas, DatasetView}; use beacon_nd_array::arrow::{ - file_read::FileRead, - metrics::ReadMetrics, - morsel::{MorselSource, OpenFile, morsel_scan}, - partition::FilePartitions, + file_read::FileRead, metrics::ReadMetrics, partition::FilePartitions, }; use datafusion::{ config::ConfigOptions, @@ -51,29 +48,15 @@ use datafusion::{ metrics::ExecutionPlanMetricsSet, }, }; -use futures::FutureExt; +use futures::{FutureExt, StreamExt, TryStreamExt}; use object_store::ObjectStore; use crate::compat; use crate::datafusion::metrics::AtlasScanMetrics; -use crate::datafusion::pruning::{CandidateFilter, PruneCache, candidate_filter, logical_schema}; +use crate::datafusion::pruning::{CandidateFilter, candidate_filter, logical_schema}; use crate::reader::{dataset_from_view, project_read_dimensions}; use crate::store::{AtlasReaderCache, get_or_open_atlas}; -/// Which dataset of a collection one plan entry stands for. -/// -/// Attached to [`PartitionedFile::extensions`] by -/// [`AtlasFormat::create_physical_plan`](super::AtlasFormat). The `position` is -/// the dataset's index in the `list_datasets()` call the plan made, which is -/// the row a collection-wide pruning index keys on. -#[derive(Debug, Clone)] -pub struct AtlasEntry { - /// The dataset's name, as the collection footer states it. - pub dataset: String, - /// Its row in the plan-time listing. - pub position: usize, -} - /// DataFusion [`FileSource`] for Atlas collections. #[derive(Debug, Clone)] pub struct AtlasSource { @@ -87,12 +70,6 @@ pub struct AtlasSource { cache: Option, /// Whether a predicate scan drops the datasets it can rule out. use_pruning: bool, - /// Each collection's pruning result, computed once for this scan and shared - /// by every partition's opener. - prune_cache: PruneCache, - /// The scan's dataset queue, when it is planned morsel-driven. See - /// [`morsel_scan`]. - morsel: Option>, } impl AtlasSource { @@ -106,8 +83,6 @@ impl AtlasSource { projection: None, cache: None, use_pruning: false, - prune_cache: PruneCache::new(), - morsel: None, } } @@ -131,12 +106,6 @@ impl AtlasSource { self.projection = projection; self } - - /// The datasets this scan's queue holds, when it is planned morsel-driven. - #[cfg(test)] - pub(crate) fn morsel_datasets(&self) -> Option { - self.morsel.as_ref().map(|source| source.files()) - } } impl FileSource for AtlasSource { @@ -147,10 +116,7 @@ impl FileSource for AtlasSource { partition: usize, ) -> Result> { let projected_schema = base_config.projected_schema()?; - let read_metrics = ReadMetrics::new(&self.execution_plan_metrics, partition); - let scan_metrics = AtlasScanMetrics::new(&self.execution_plan_metrics, partition); - - let datasets = Arc::new(AtlasDatasets { + Ok(Arc::new(AtlasOpener { object_store, cache: self.cache.clone(), // A predicate is written against the values, not the encoding the @@ -158,19 +124,11 @@ impl FileSource for AtlasSource { logical_schema: logical_schema(&projected_schema), projected_schema, use_pruning: self.use_pruning, - prune_cache: self.prune_cache.clone(), read_dimensions: self.read_dimensions.clone(), batch_size: self.batch_size, predicate: self.predicate.clone(), - read_metrics: read_metrics.clone(), - scan_metrics, - }); - - Ok(Arc::new(AtlasOpener { - datasets, - morsel: self.morsel.clone(), - partition, - read_metrics, + read_metrics: ReadMetrics::new(&self.execution_plan_metrics, partition), + scan_metrics: AtlasScanMetrics::new(&self.execution_plan_metrics, partition), })) } @@ -189,50 +147,10 @@ impl FileSource for AtlasSource { }) } - /// Put every dataset of the scan in one queue, and point each partition at - /// it. - /// - /// Nothing is assigned here. A dataset's cost is the cells the query keeps, - /// which no plan-time number states: two datasets of one collection differ - /// by orders of magnitude, and a predicate prunes them unevenly. So the - /// partitions divide the queue as they drain it. - /// - /// `repartition_file_min_size` is ignored, as it is for Zarr. It was the - /// size a file had to reach before sharing it was worth the seek, and a - /// queue makes no such bet. An atlas entry has no size of its own anyway: - /// every dataset of a collection reports the container's. - fn repartitioned( - &self, - target_partitions: usize, - _repartition_file_min_size: usize, - output_ordering: Option, - config: &FileScanConfig, - ) -> Result> { - if output_ordering.is_some() || target_partitions <= 1 { - // A partition holding an arbitrary share of the datasets cannot - // emit its rows in collection order. - return Ok(None); - } - - if let Some((morsel, file_groups)) = morsel_scan(&config.file_groups, target_partitions) { - tracing::debug!( - "AtlasSource morsel scan: {} datasets over {target_partitions} partitions", - morsel.files() - ); - let mut config = config.clone(); - config.file_groups = file_groups; - // The openers are built from the config's source, so the queue has - // to travel with it. - config.file_source = Arc::new(Self { - morsel: Some(morsel), - ..self.clone() - }); - return Ok(Some(config)); - } - - // The queue declined: one partition, or no datasets. Keeping the scan - // as planned is the answer to both. - Ok(None) + /// A container is one unit. A byte range of it names nothing a reader can + /// open, so the plan's groups stand as the format dealt them. + fn supports_repartitioning(&self) -> bool { + false } fn metrics(&self) -> &ExecutionPlanMetricsSet { @@ -293,48 +211,12 @@ impl FileSource for AtlasSource { // ─── The opener ────────────────────────────────────────────────────────────── -/// One partition's opener. -struct AtlasOpener { - /// How one dataset is opened, for the queue to call. - datasets: Arc, - /// The scan's queue, when it is planned morsel-driven. `Some` means the - /// entry `FileStream` hands this opener stands for the whole scan. - morsel: Option>, - partition: usize, - read_metrics: ReadMetrics, -} - -impl FileOpener for AtlasOpener { - fn open(&self, file: PartitionedFile) -> Result { - // A morsel-driven scan hands every partition the same standing entry. - // It is not a dataset: the datasets are in the queue, and this - // partition reads whatever it hands out until the scan is done. - if let Some(morsel) = &self.morsel { - let stream = morsel.stream( - self.partition, - Arc::clone(&self.datasets), - Some(self.read_metrics.clone()), - ); - return Ok(futures::future::ready(Ok(stream)).boxed()); - } - - // One partition, or no datasets: `FileStream` walks the real entries - // and this opener reads each one whole. - let datasets = Arc::clone(&self.datasets); - let metrics = self.read_metrics.clone(); - Ok(async move { - let read = datasets.open(&file).await?; - Ok(read.stream(Some(metrics))) - } - .boxed()) - } -} - -/// How one Atlas dataset becomes a planned [`FileRead`]. +/// One partition's opener: a collection in, its batches out. /// -/// This is everything a [`MorselSource`] needs of the format. The queue holds -/// the datasets; this says what opening one means. -struct AtlasDatasets { +/// Every field is a handle or a clone, so the opener itself is cloned into the +/// stream it returns and outlives the call that made it. +#[derive(Clone)] +struct AtlasOpener { object_store: Arc, cache: Option, /// The scan's output schema, nd-encoded. Its field *names* are the columns @@ -344,7 +226,6 @@ struct AtlasDatasets { /// and the pruning engine are written against. logical_schema: SchemaRef, use_pruning: bool, - prune_cache: PruneCache, read_dimensions: Option>, batch_size: usize, predicate: Option>, @@ -352,24 +233,85 @@ struct AtlasDatasets { scan_metrics: AtlasScanMetrics, } -impl AtlasDatasets { - /// Which datasets of one collection this scan's predicate can still match. +impl FileOpener for AtlasOpener { + fn open(&self, file: PartitionedFile) -> Result { + let opener = self.clone(); + Ok(async move { + let collection = file.object_meta.location.to_string(); + + let open_start = Instant::now(); + let atlas = get_or_open_atlas( + opener.cache.as_ref(), + Arc::clone(&opener.object_store), + &file.object_meta, + ) + .await + .map_err(|e| DataFusionError::Execution(format!("{e}")))?; + opener.scan_metrics.open_time.add_elapsed(open_start); + + let listed = atlas.list_datasets(); + let total = listed.len(); + let names = opener.survivors(&atlas, listed).await; + tracing::debug!( + collection = %collection, + datasets = total, + kept = names.len(), + "atlas scan: collection opened and pruned" + ); + + // One stream over the kept datasets. Each is built when the stream + // reaches it, and its batches drain before the next is touched. + let metrics = opener.read_metrics.clone(); + let batches = futures::stream::iter(names) + .then(move |name| { + let opener = opener.clone(); + let atlas = Arc::clone(&atlas); + let collection = collection.clone(); + async move { opener.dataset_read(&atlas, &collection, &name).await } + }) + .map_ok(move |read| read.stream(Some(metrics.clone()))) + .try_flatten(); + + Ok(batches.boxed()) + } + .boxed()) + } +} + +impl AtlasOpener { + /// The datasets of `atlas` this scan reads, in listing order. /// - /// Built once per collection per scan. Every partition's opener shares one - /// memo, so the first to arrive builds the index while the rest await it; - /// each then reads its own dataset's bit out of the result. + /// One pruning pass over the whole collection decides it, and a dataset + /// ruled out never comes back: no handle, no build, no read. + async fn survivors(&self, atlas: &Arc, listed: Vec) -> Vec { + let filter = self.candidates(atlas).await; + + let mut kept = Vec::with_capacity(listed.len()); + let mut pruned = 0usize; + for (position, name) in listed.into_iter().enumerate() { + if filter.keeps(position, &name) { + kept.push(name); + } else { + pruned += 1; + } + } + self.scan_metrics.datasets_pruned.add(pruned); + kept + } + + /// Which datasets of `atlas` this scan's predicate can still match. /// /// Without a predicate, or with pruning off, nothing is ruled out and no /// index is built. - async fn candidates(&self, atlas: &Arc, marker: &str) -> Arc { + async fn candidates(&self, atlas: &Arc) -> CandidateFilter { let Some(predicate) = self.predicate.clone().filter(|_| self.use_pruning) else { - return Arc::new(CandidateFilter::KeepAll); + return CandidateFilter::KeepAll; }; // DataFusion offers a scan its filters even when it has none, and an // empty conjunction is the literal `true`. Such a predicate names no // column, so it can rule nothing out and is not worth a pass. if collect_columns(&predicate).is_empty() { - return Arc::new(CandidateFilter::KeepAll); + return CandidateFilter::KeepAll; } // The predicate is indexed against the table schema, the pruning @@ -377,30 +319,64 @@ impl AtlasDatasets { // after the filter leaves the two out of step, so the columns are // re-indexed by name here. let Ok(predicate) = reassign_expr_columns(predicate, &self.logical_schema) else { - return Arc::new(CandidateFilter::KeepAll); + return CandidateFilter::KeepAll; }; let started = Instant::now(); - let atlas = Arc::clone(atlas); - let schema = Arc::clone(&self.logical_schema); - let metrics = self.scan_metrics.clone(); - let filter = self - .prune_cache - .get_or_compute(marker.to_string(), async move { - let filter = Arc::new(candidate_filter(&atlas, &predicate, &schema).await); - // Only an index that exists is a build. Pruning that did not - // apply read nothing and judged nothing. - if filter.is_index() { - metrics.index_builds.add(1); - metrics.index_rows.add(filter.rows()); - } - filter - }) - .await; + let filter = candidate_filter(atlas, &predicate, &self.logical_schema).await; + // Only an index that exists is a build. Pruning that did not apply + // read nothing and judged nothing. + if filter.is_index() { + self.scan_metrics.index_builds.add(1); + self.scan_metrics.index_rows.add(filter.rows()); + } self.scan_metrics.prune_time.add_elapsed(started); filter } + /// One dataset, built and planned as a [`FileRead`]. + async fn dataset_read( + &self, + atlas: &Arc, + collection: &str, + name: &str, + ) -> Result> { + let build_start = Instant::now(); + let view = Arc::new(atlas.dataset(name).map_err(|e| { + DataFusionError::Execution(format!( + "Failed to open atlas dataset '{name}' of '{collection}': {e}" + )) + })?); + + let projected = self.projected_names(&view).await?; + let dataset = dataset_from_view(view, projected.as_deref()) + .await + .map_err(|e| DataFusionError::Execution(format!("{e}")))?; + // Explicit dimensions, or a broadcast-compatible default. No log label: + // this runs per dataset, and schema inference already logged the choice. + let dataset = project_read_dimensions(dataset, self.read_dimensions.clone(), None) + .map_err(|e| DataFusionError::Execution(format!("{e}")))?; + + let read = FileRead::plan( + dataset, + Arc::clone(&self.projected_schema), + self.batch_size, + self.predicate.clone(), + // A dataset lives inside a container, not at a path, so no + // `PARTITIONED BY` value can be read off it. The format refuses + // such a table outright. + FilePartitions::none(), + Some(&self.read_metrics), + ) + .await?; + + self.scan_metrics + .dataset_build_time + .add_elapsed(build_start); + self.scan_metrics.datasets_scanned.add(1); + Ok(read) + } + /// The columns to build for one dataset, or `None` to build every one. /// /// The projection reaches the build, so an unprojected array gets no @@ -489,85 +465,11 @@ async fn count_driver( Ok(widest.map(|(array, _)| array)) } -#[async_trait::async_trait] -impl OpenFile for AtlasDatasets { - async fn open(&self, file: &PartitionedFile) -> Result> { - let entry = file - .extensions - .as_ref() - .and_then(|extension| (extension.as_ref() as &dyn Any).downcast_ref::()) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "the atlas scan entry at '{}' names no dataset", - file.object_meta.location - )) - })?; - - let open_start = Instant::now(); - let atlas = get_or_open_atlas( - self.cache.as_ref(), - Arc::clone(&self.object_store), - &file.object_meta, - ) - .await - .map_err(|e| DataFusionError::Execution(format!("{e}")))?; - self.scan_metrics.open_time.add_elapsed(open_start); - - // One index per collection decides this, and the first opener to reach - // the collection builds it. A dataset ruled out costs one pop and no - // read at all. - if !self - .candidates(&atlas, file.object_meta.location.as_ref()) - .await - .keeps(entry.position, &entry.dataset) - { - self.scan_metrics.datasets_pruned.add(1); - return Ok(FileRead::skipped()); - } - - let build_start = Instant::now(); - let view = Arc::new(atlas.dataset(&entry.dataset).map_err(|e| { - DataFusionError::Execution(format!( - "Failed to open atlas dataset '{}' of '{}': {e}", - entry.dataset, file.object_meta.location - )) - })?); - - let projected = self.projected_names(&view).await?; - let dataset = dataset_from_view(view, projected.as_deref()) - .await - .map_err(|e| DataFusionError::Execution(format!("{e}")))?; - // Explicit dimensions, or a broadcast-compatible default. No log label: - // this runs per dataset, and schema inference already logged the choice. - let dataset = project_read_dimensions(dataset, self.read_dimensions.clone(), None) - .map_err(|e| DataFusionError::Execution(format!("{e}")))?; - - let read = FileRead::plan( - dataset, - Arc::clone(&self.projected_schema), - self.batch_size, - self.predicate.clone(), - // A dataset lives inside a container, not at a path, so no - // `PARTITIONED BY` value can be read off it. The format refuses - // such a table outright. - FilePartitions::none(), - Some(&self.read_metrics), - ) - .await?; - - self.scan_metrics - .dataset_build_time - .add_elapsed(build_start); - self.scan_metrics.datasets_scanned.add(1); - Ok(read) - } -} - #[cfg(test)] mod tests { use super::*; - use datafusion::datasource::physical_plan::FileScanConfigBuilder; - use datafusion::execution::object_store::ObjectStoreUrl; + use crate::test_support; + use datafusion::physical_plan::metrics::MetricsSet; fn source() -> AtlasSource { AtlasSource::new( @@ -576,122 +478,175 @@ mod tests { ) } - fn entry(dataset: &str, position: usize) -> PartitionedFile { - let mut file = PartitionedFile::new("obs/data.atlas", 4096); - file.extensions = Some(Arc::new(AtlasEntry { - dataset: dataset.to_string(), - position, - })); - file + /// A container is one unit of work, so DataFusion may not split it into + /// byte ranges the way it would a Parquet file. + #[test] + fn a_container_is_never_split() { + assert!(!source().supports_repartitioning()); } - /// Every dataset of the scan goes into one queue, and each partition gets a - /// standing entry pointing at it. - #[test] - fn the_datasets_go_into_one_queue() { - const PARTITIONS: usize = 4; + // ── opening a collection ──────────────────────────────────────────── - let source = source(); - let mut builder = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), - Arc::new(source.clone()) as Arc, - ); - for (position, name) in ["a", "b", "c", "d", "e"].iter().enumerate() { - builder = builder.with_file(entry(name, position)); - } - let config = builder.build(); + /// An opener over the collection in `dir`, projecting every column, with + /// the metrics it reports. + async fn opener( + dir: &std::path::Path, + predicate: Option>, + use_pruning: bool, + ) -> (AtlasOpener, ExecutionPlanMetricsSet) { + use crate::reader::collection_schema; + use beacon_datafusion_ext::type_widening::ArrowTypeWidening; + + let atlas = test_support::open(dir).await; + let logical = collection_schema(&atlas, None, "c", &ArrowTypeWidening::default_extension()) + .await + .unwrap(); + let projected_schema: SchemaRef = + Arc::new(beacon_datafusion_ext::nd::encoded_schema(&logical)); + + let metrics = ExecutionPlanMetricsSet::new(); + let (store, _) = test_support::store_and_marker(dir); + let opener = AtlasOpener { + object_store: store, + cache: None, + logical_schema: logical_schema(&projected_schema), + projected_schema, + use_pruning, + read_dimensions: None, + batch_size: usize::MAX, + predicate, + read_metrics: ReadMetrics::new(&metrics, 0), + scan_metrics: AtlasScanMetrics::new(&metrics, 0), + }; + (opener, metrics) + } + + fn count_of(metrics: &ExecutionPlanMetricsSet, name: &str) -> usize { + let set: MetricsSet = metrics.clone_inner(); + set.sum_by_name(name).map_or(0, |value| value.as_usize()) + } + + /// One open, then every dataset of the collection, in listing order. + #[tokio::test] + async fn an_open_streams_every_dataset_of_the_collection() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 6).await; + let (opener, metrics) = opener(tmp.path(), None, true).await; + let (_, marker) = test_support::store_and_marker(tmp.path()); - let planned = source - .repartitioned(PARTITIONS, 10 * 1024 * 1024, None, &config) + let stream = opener + .open(PartitionedFile::from(marker)) .unwrap() - .expect("the datasets are planned across the partitions"); + .await + .expect("the collection opens"); + let batches: Vec<_> = stream.try_collect().await.expect("every dataset reads"); - assert_eq!(planned.file_groups.len(), PARTITIONS); - for group in &planned.file_groups { - assert_eq!(group.len(), 1, "one standing entry per partition"); - } - let planned = planned - .file_source() - .as_any() - .downcast_ref::() - .expect("the config carries an AtlasSource"); + // Six datasets, one chunk each, one nd batch per chunk. + assert_eq!(batches.len(), 6); + assert_eq!(count_of(&metrics, "atlas_datasets_scanned"), 6); + assert_eq!(count_of(&metrics, "atlas_datasets_pruned"), 0); assert_eq!( - planned.morsel_datasets(), - Some(5), - "and the queue holds every dataset" + count_of(&metrics, "atlas_index_builds"), + 0, + "no predicate, no index" ); } - /// An ordered scan cannot share: a partition holding an arbitrary share of - /// the datasets cannot emit its rows in collection order. - #[test] - fn an_ordered_scan_is_left_alone() { - use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; - use datafusion::physical_plan::expressions::Column; - - let source = source(); - let config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), - Arc::new(source.clone()) as Arc, - ) - .with_file(entry("a", 0)) - .build(); + /// A predicate prunes once for the collection, before any dataset is + /// built. A ruled-out dataset then costs no build and no read. + #[tokio::test] + async fn an_open_prunes_the_collection_before_it_reads() { + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; + use datafusion::scalar::ScalarValue; - let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( - Column::new("time", 0), - ))]); - assert!( - source - .repartitioned(4, 0, ordering, &config) - .unwrap() - .is_none() - ); + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 6).await; + // Datasets hold 10i..=10i+3, so only d5 (50..=53) reaches past 45. + let predicate: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("temperature", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Float32(Some(45.0)))), + )); + let (opener, metrics) = opener(tmp.path(), Some(predicate), true).await; + let (_, marker) = test_support::store_and_marker(tmp.path()); + + let stream = opener + .open(PartitionedFile::from(marker)) + .unwrap() + .await + .expect("the collection opens"); + let batches: Vec<_> = stream.try_collect().await.unwrap(); + + assert_eq!(batches.len(), 1, "one dataset survives"); + assert_eq!(count_of(&metrics, "atlas_datasets_scanned"), 1); + assert_eq!(count_of(&metrics, "atlas_datasets_pruned"), 5); + assert_eq!(count_of(&metrics, "atlas_index_builds"), 1); + assert_eq!(count_of(&metrics, "atlas_index_rows"), 6); } - #[test] - fn one_partition_divides_nothing() { - let source = source(); - let config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), - Arc::new(source.clone()) as Arc, - ) - .with_file(entry("a", 0)) - .build(); + /// With pruning off, every dataset is built and read. + /// + /// The predicate still reaches [`FileRead::plan`], where the chunk mask over + /// a 1-D coordinate skips chunks no row of which can match. That layer is + /// always on and is not what the switch controls, so this pins the dataset + /// metrics and not the batch count. + #[tokio::test] + async fn pruning_off_reads_every_dataset() { + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; + use datafusion::scalar::ScalarValue; + + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 6).await; + let predicate: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("temperature", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Float32(Some(45.0)))), + )); + let (opener, metrics) = opener(tmp.path(), Some(predicate), false).await; + let (_, marker) = test_support::store_and_marker(tmp.path()); + + let stream = opener + .open(PartitionedFile::from(marker)) + .unwrap() + .await + .unwrap(); + let _batches: Vec<_> = stream.try_collect().await.unwrap(); - assert!(source.repartitioned(1, 0, None, &config).unwrap().is_none()); + assert_eq!( + count_of(&metrics, "atlas_datasets_scanned"), + 6, + "every dataset is built" + ); + assert_eq!(count_of(&metrics, "atlas_datasets_pruned"), 0); + assert_eq!( + count_of(&metrics, "atlas_index_builds"), + 0, + "no index is built" + ); } - /// An entry that names no dataset is a bug in the planner, not bad input, - /// and the error says which collection it came from. + /// A path that is not a container fails at the open, and the error names + /// what a collection is called. #[tokio::test] - async fn an_entry_without_a_dataset_is_an_internal_error() { - let datasets = AtlasDatasets { - object_store: Arc::new(object_store::memory::InMemory::new()), - cache: None, - projected_schema: Arc::new(arrow::datatypes::Schema::empty()), - logical_schema: Arc::new(arrow::datatypes::Schema::empty()), - use_pruning: false, - prune_cache: PruneCache::new(), - read_dimensions: None, - batch_size: usize::MAX, - predicate: None, - read_metrics: ReadMetrics::new(&ExecutionPlanMetricsSet::new(), 0), - scan_metrics: AtlasScanMetrics::new(&ExecutionPlanMetricsSet::new(), 0), - }; + async fn a_path_that_is_not_a_container_fails_to_open() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let (opener, _) = opener(tmp.path(), None, false).await; - let error = datasets - .open(&PartitionedFile::new("obs/data.atlas", 1)) + let error = opener + .open(PartitionedFile::new("obs/index.json", 1)) + .unwrap() .await - .expect_err("an entry must name its dataset") + .err() + .expect("only the container names a collection") .to_string(); - assert!(error.contains("names no dataset"), "{error}"); - assert!(error.contains("obs/data.atlas"), "{error}"); + assert!(error.contains("data.atlas"), "{error}"); } // ── which columns a dataset is built with ─────────────────────────── - use crate::test_support; - /// One dataset's view. It owns what it needs, so the collection handle /// behind it may go. async fn view(dir: &std::path::Path, dataset: &str) -> DatasetView { @@ -701,22 +656,21 @@ mod tests { .expect("the dataset") } - fn datasets_wanting( + fn opener_wanting( projected: Vec<&str>, predicate: Option>, - ) -> AtlasDatasets { + ) -> AtlasOpener { let fields: Vec = projected .into_iter() .map(|name| arrow::datatypes::Field::new(name, arrow::datatypes::DataType::Null, true)) .collect(); let projected_schema = Arc::new(arrow::datatypes::Schema::new(fields)); - AtlasDatasets { + AtlasOpener { object_store: Arc::new(object_store::memory::InMemory::new()), cache: None, logical_schema: Arc::clone(&projected_schema), projected_schema, use_pruning: false, - prune_cache: PruneCache::new(), read_dimensions: None, batch_size: usize::MAX, predicate, @@ -730,8 +684,8 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); test_support::two_datasets(tmp.path()).await; - let datasets = datasets_wanting(vec!["temperature"], None); - let names = datasets + let opener = opener_wanting(vec!["temperature"], None); + let names = opener .projected_names(&view(tmp.path(), "winter").await) .await .expect("the layouts resolve") @@ -756,8 +710,8 @@ mod tests { Operator::Gt, Arc::new(Literal::new(ScalarValue::Int32(Some(20)))), )); - let datasets = datasets_wanting(vec!["temperature"], Some(predicate)); - let names = datasets + let opener = opener_wanting(vec!["temperature"], Some(predicate)); + let names = opener .projected_names(&view(tmp.path(), "winter").await) .await .expect("the layouts resolve") @@ -777,8 +731,8 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); test_support::chunked_grid(tmp.path()).await; - let datasets = datasets_wanting(vec![], None); - let names = datasets + let opener = opener_wanting(vec![], None); + let names = opener .projected_names(&view(tmp.path(), "grid").await) .await .expect("the layouts resolve") diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs deleted file mode 100644 index d6699c52..00000000 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/statistics.rs +++ /dev/null @@ -1,296 +0,0 @@ -//! The column ranges of a whole collection, for the file analyzer. -//! -//! A collection reports one range per column, folded over its live datasets. -//! It costs no array read: the writer computed each dataset's minimum and -//! maximum while it staged the data, and stored them in that variable's -//! segment. One open per column answers for the whole collection. -//! -//! # Why a wrong answer here is worse than no answer -//! -//! A recorded range prunes whole collections before a scan opens them, so a -//! range that is too narrow silently deletes matching rows from an answer. -//! Every path below reports unknown unless it can *prove* the bound. Unknown is -//! always legal: it only means Beacon reads what it might have skipped. - -use arrow::datatypes::{DataType, Schema}; -use atlas::{Atlas, StatValue}; -use datafusion::common::{ColumnStatistics, Statistics, stats::Precision}; -use datafusion::scalar::ScalarValue; - -/// The statistics of one collection, in `table_schema` order. -pub async fn collection_statistics(atlas: &Atlas, table_schema: &Schema) -> Statistics { - let arrays = atlas.list_arrays(); - let live = atlas.dataset_count(); - - let mut statistics = Statistics::default(); - for field in table_schema.fields() { - let range = if arrays.iter().any(|array| array == field.name()) { - column_range(atlas, field.name(), field.data_type(), live).await - } else { - // An attribute is exact rather than a range, and bounding a column - // by it says nothing a scan can use. The query-time pruning index - // reads attributes; the analyzer has no use for them. - None - }; - - statistics = statistics.add_column_statistics(match range { - Some((min, max)) => ColumnStatistics::new_unknown() - .with_min_value(Precision::Exact(min)) - .with_max_value(Precision::Exact(max)), - None => ColumnStatistics::new_unknown(), - }); - } - statistics -} - -/// The range of one array column over every live dataset, or `None`. -/// -/// # Every live dataset must report -/// -/// A dataset that declares an array and never writes it has no statistics -/// entry, and its cells read back as the array's fill — or, when it declares -/// none, as zeros that nothing nulls. Folding only the datasets that *do* -/// report would then produce a range those zeros sit outside of, and pruning -/// would drop the collection for a query that matches them. -/// -/// A segment holds an entry only for a dataset that wrote the array, so the -/// count is the proof: a bound is claimed only when every live dataset reported -/// one. A uniform collection — which is what `atlas create` writes, and what -/// this format exists for — satisfies that; a heterogeneous one goes unknown -/// and is read in full. -async fn column_range( - atlas: &Atlas, - column: &str, - target: &DataType, - live: usize, -) -> Option<(ScalarValue, ScalarValue)> { - let per_dataset = atlas.array_stats_by_dataset(column).await.ok()?; - if per_dataset.is_empty() || per_dataset.len() != live { - return None; - } - - let mut low: Option = None; - let mut high: Option = None; - for stats in per_dataset { - // One dataset without a bound leaves the column unbounded: its values - // may lie anywhere. - let min = bound(stats.min.as_ref(), target)?; - let max = bound(stats.max.as_ref(), target)?; - low = Some(match low { - Some(held) => smaller(held, min)?, - None => min, - }); - high = Some(match high { - Some(held) => larger(held, max)?, - None => max, - }); - } - Some((low?, high?)) -} - -/// One statistic as a scalar of the table's type, or `None` when it proves -/// nothing: absent, `NaN` — which sorts last and bounds nothing — or a value -/// that will not cast. -fn bound(value: Option<&StatValue>, target: &DataType) -> Option { - let canonical = match value? { - StatValue::Int(v) => ScalarValue::Int64(Some(*v)), - StatValue::UInt(v) => ScalarValue::UInt64(Some(*v)), - StatValue::Float(v) if v.is_nan() => return None, - StatValue::Float(v) => ScalarValue::Float64(Some(*v)), - StatValue::TimestampNs(v) => ScalarValue::TimestampNanosecond(Some(*v), None), - StatValue::Bytes(bytes) => match std::str::from_utf8(bytes) { - Ok(text) => ScalarValue::Utf8(Some(text.to_string())), - Err(_) => ScalarValue::Binary(Some(bytes.clone())), - }, - }; - let cast = canonical.cast_to(target).ok()?; - // A cast that lands on null has lost the value, and a null bounds nothing. - if cast.is_null() { None } else { Some(cast) } -} - -/// The lower of two bounds, or `None` when they do not compare. -fn smaller(held: ScalarValue, next: ScalarValue) -> Option { - match held.partial_cmp(&next)? { - std::cmp::Ordering::Greater => Some(next), - _ => Some(held), - } -} - -/// The higher of two bounds, or `None` when they do not compare. -fn larger(held: ScalarValue, next: ScalarValue) -> Option { - match held.partial_cmp(&next)? { - std::cmp::Ordering::Less => Some(next), - _ => Some(held), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::test_support; - use arrow::datatypes::Field; - - fn schema(fields: Vec) -> Schema { - Schema::new(fields) - } - - fn range( - statistics: &Statistics, - index: usize, - ) -> (Precision, Precision) { - let column = &statistics.column_statistics[index]; - (column.min_value.clone(), column.max_value.clone()) - } - - /// A uniform collection — every dataset holding every array — reports the - /// union of its datasets' ranges. - #[tokio::test] - async fn a_uniform_collection_reports_the_union_of_its_datasets() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 10).await; - let atlas = test_support::open(tmp.path()).await; - - let statistics = collection_statistics( - &atlas, - &schema(vec![Field::new("temperature", DataType::Float32, true)]), - ) - .await; - let (min, max) = range(&statistics, 0); - // d0 starts at 0 and d9 ends at 93. - assert_eq!(min, Precision::Exact(ScalarValue::Float32(Some(0.0)))); - assert_eq!(max, Precision::Exact(ScalarValue::Float32(Some(93.0)))); - } - - /// A column only some datasets declare goes unknown, because a dataset that - /// declared it and never wrote it is indistinguishable from one that never - /// declared it, and the first reads back as values this fold cannot see. - #[tokio::test] - async fn a_column_not_every_dataset_reports_goes_unknown() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - let atlas = test_support::open(tmp.path()).await; - - let statistics = collection_statistics( - &atlas, - &schema(vec![ - // winter alone declares `cycle`. - Field::new("cycle", DataType::Int32, true), - // both declare `temperature`. - Field::new("temperature", DataType::Float32, true), - ]), - ) - .await; - assert_eq!(range(&statistics, 0).0, Precision::Absent); - assert_eq!( - range(&statistics, 1).0, - Precision::Exact(ScalarValue::Float32(Some(1.0))), - "a column every dataset reports still bounds" - ); - assert_eq!( - range(&statistics, 1).1, - Precision::Exact(ScalarValue::Float32(Some(22.0))), - ); - } - - /// A column two datasets type differently is folded on the table's type. - #[tokio::test] - async fn a_mixed_dtype_column_folds_on_the_table_type() { - let tmp = tempfile::tempdir().unwrap(); - test_support::widening(tmp.path()).await; - let atlas = test_support::open(tmp.path()).await; - - let statistics = collection_statistics( - &atlas, - &schema(vec![Field::new("value", DataType::Float64, true)]), - ) - .await; - let (min, max) = range(&statistics, 0); - // a holds [1, 2] as Int16 and b holds [3.5, 4.5] as Float32. - assert_eq!(min, Precision::Exact(ScalarValue::Float64(Some(1.0)))); - assert_eq!(max, Precision::Exact(ScalarValue::Float64(Some(4.5)))); - } - - /// A column the collection does not hold, and a column whose values are - /// text, both report unknown rather than a guess. - #[tokio::test] - async fn an_unknown_column_reports_unknown() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 4).await; - let atlas = test_support::open(tmp.path()).await; - - let statistics = collection_statistics( - &atlas, - &schema(vec![ - Field::new("ghost", DataType::Float32, true), - Field::new(".platform", DataType::Utf8, true), - ]), - ) - .await; - assert_eq!(range(&statistics, 0).0, Precision::Absent); - assert_eq!( - range(&statistics, 1).0, - Precision::Absent, - "an attribute is not measured here" - ); - } - - #[tokio::test] - async fn an_empty_collection_bounds_nothing() { - let tmp = tempfile::tempdir().unwrap(); - test_support::empty(tmp.path()).await; - let atlas = test_support::open(tmp.path()).await; - - let statistics = collection_statistics( - &atlas, - &schema(vec![Field::new("temperature", DataType::Float32, true)]), - ) - .await; - assert_eq!(range(&statistics, 0).0, Precision::Absent); - } - - /// A deleted dataset counts toward nothing: neither the fold nor the count - /// that guards it. - #[tokio::test] - async fn a_deleted_dataset_leaves_the_range() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 10).await; - let atlas = test_support::open(tmp.path()).await; - atlas.delete_dataset("d9").await.unwrap(); - - // Reopen, so the handle reads the mask that was just written. - let atlas = test_support::open(tmp.path()).await; - let statistics = collection_statistics( - &atlas, - &schema(vec![Field::new("temperature", DataType::Float32, true)]), - ) - .await; - assert_eq!( - range(&statistics, 0).1, - Precision::Exact(ScalarValue::Float32(Some(83.0))), - "d9 held the values up to 93 and is gone" - ); - } - - // ── the pieces ────────────────────────────────────────────────────── - - #[test] - fn a_nan_bound_proves_nothing() { - assert!(bound(Some(&StatValue::Float(f64::NAN)), &DataType::Float64).is_none()); - } - - #[test] - fn a_bound_that_will_not_cast_proves_nothing() { - assert!( - bound( - Some(&StatValue::Bytes(b"argo".to_vec())), - &DataType::Float64 - ) - .is_none() - ); - } - - #[test] - fn an_absent_bound_proves_nothing() { - assert!(bound(None, &DataType::Float64).is_none()); - } -} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs index 3aaf667e..281fdf5e 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs @@ -35,13 +35,13 @@ //! derives the Arrow schema of a whole collection. [`compat`] holds the type //! and column-name mapping the two share. //! -//! # One dataset is one unit of work +//! # One collection is one unit of work //! -//! The scan plans one entry per dataset and puts them all in one morsel queue, -//! as the netCDF and Zarr scans do for their files. A partition takes the next -//! dataset when it is free, and each dataset is cut on the chunk grid the -//! writer actually chose, so a dataset stored as one chunk yields one unit and -//! a chunked one yields many. See [`datafusion::source`]. +//! The scan plans one entry per collection and deals them over the partitions. +//! A partition opens each collection it holds once, prunes every dataset in one +//! pass over the footer, and streams the survivors one after another. A pruned +//! dataset therefore costs nothing, and parallelism is bounded by the +//! collection count. See [`datafusion::source`]. //! //! # Columns //! diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs index 00c2a2b6..69de28a5 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs @@ -6,7 +6,6 @@ //! footer, which the open already held, and its columns fetch their bytes when //! the scan asks for them. -use std::collections::HashSet; use std::sync::Arc; use arrow::datatypes::{Schema, SchemaRef}; @@ -202,12 +201,6 @@ pub fn project_read_dimensions( ) -> anyhow::Result { match resolve_read_dimensions(&dataset, read_dimensions, log_label) { Some(dims) => { - // A projection names the dimensions of a whole collection, and one - // dataset need not hold them all. Narrowing to those it does hold - // drops the same arrays — an array survives only when every one of - // its own dimensions is kept, and a dimension this dataset lacks - // cannot be one of them — and leaves an all-scalar projection with - // nothing to narrow rather than an unknown dimension to refuse. let held: Vec = dims .into_iter() .filter(|dim| dataset.dataset().dimensions.contains_key(dim)) @@ -241,31 +234,14 @@ pub async fn open_dataset( /// The Arrow schema of a whole collection: every live dataset, merged. /// /// Atlas reconciles nothing. Two datasets may declare one array name with two -/// dtypes, so the collection's schema is the widening merge of its datasets' +/// dtypes, so the collection's schema is the widening merge of its datasets /// schemas, under the rule the session carries. -/// -/// # One schema per shape, not per dataset -/// -/// Datasets that declare the same arrays share one interned schema in the -/// footer, and `atlas create` writes a fleet of files that way. Each dataset is -/// therefore reduced to a key over its interned schema and its attribute -/// namespace, and a schema is derived once per distinct key. A thousand -/// datasets of one shape cost one derivation. -/// -/// # Cost -/// -/// Linear in the dataset count, and every step is in memory: the footer keys -/// its datasets by name, so resolving one is a single hash lookup, and a key is -/// built from the footer and from segments that one open serves collection -/// wide. Only the distinct keys are derived into a schema. The result is cached -/// above this crate, so a table pays even that once rather than once per query. pub async fn collection_schema( atlas: &Arc, read_dimensions: Option<&[String]>, label: &str, widening: &ArrowTypeWidening, ) -> anyhow::Result { - let mut seen: HashSet = HashSet::new(); let mut schemas: Vec = Vec::new(); for name in atlas.list_datasets() { @@ -273,10 +249,6 @@ pub async fn collection_schema( .dataset(&name) .map_err(|e| anyhow::anyhow!("Failed to open atlas dataset '{name}': {e}"))?; - if !seen.insert(shape_key(&view).await?) { - continue; - } - let dataset = dataset_from_view(Arc::new(view), None).await?; // The same narrowing the scan applies, so the schema states what a // query can actually return. @@ -303,75 +275,6 @@ pub async fn collection_schema( .map_err(|e| anyhow::anyhow!("Failed to merge the schemas of the atlas datasets: {e}")) } -/// What makes two datasets produce the same columns and types. -/// -/// Three things decide a dataset's Arrow schema, and the key holds all three: -/// -/// - The arrays it declares, with their element types. Datasets that declare -/// the same ones share one interned schema in the footer, and `atlas create` -/// writes a fleet of files that way. -/// - Its attribute keys and their types, at both scopes. Those are named in the -/// interned schema too, so two datasets that differ only in an attribute's -/// *value* share a key. That is exactly the fleet case. -/// - Each array's dimension names, which the interned schema does **not** hold. -/// They pick the default grid, and a different grid keeps different columns, -/// so two datasets that agree on everything else can still differ here. -/// -/// A shape is deliberately left out: an array of a different length is the same -/// column. -/// -/// # Cost -/// -/// The names and the types come from the footer. A dimension name comes from -/// its variable's segment, which one open serves for every dataset of the -/// collection, so the lookup is in memory after the first. -async fn shape_key(view: &DatasetView) -> anyhow::Result { - let mut key = String::new(); - - for (array, dtype) in declared_arrays(view) { - key.push('|'); - key.push_str(&array); - key.push(':'); - key.push_str(&compat::dtype_tag(&dtype)); - - // An array Beacon cannot read is no column, so its grid decides - // nothing and its segment stays shut. - if compat::array_dtype_to_nd(&dtype).is_some() { - let layout = view.array_layout(&array).await.map_err(|e| { - anyhow::anyhow!( - "Failed to read the layout of atlas array '{array}' of dataset '{}': {e}", - view.name() - ) - })?; - key.push('@'); - key.push_str(&layout.dimension_names().join(",")); - } - } - - // Attribute keys and types are in the interned schema, so this reads - // nothing. The values are not, and they do not belong in the key. - fn push(key: &mut String, array: &str, attr: &str, dtype: &DType) { - key.push('|'); - key.push_str(array); - key.push('.'); - key.push_str(attr); - key.push(':'); - key.push_str(&compat::dtype_tag(dtype)); - } - let schema = view.schema(); - for meta in schema.iter() { - let array = meta.name(); - for (attr, dtype) in meta.attribute_pairs() { - push(&mut key, array, attr, dtype); - } - } - for (attr, dtype) in schema.attribute_pairs() { - push(&mut key, "", attr, dtype); - } - - Ok(key) -} - #[cfg(test)] mod tests { use super::*; @@ -763,19 +666,14 @@ mod tests { /// gives ten datasets one shape, and they differ only in an attribute /// value. #[tokio::test] - async fn a_fleet_of_one_shape_derives_one_schema() { + async fn a_fleet_of_one_shape_merges_to_one_set_of_columns() { let tmp = tempfile::tempdir().unwrap(); test_support::ranged(tmp.path(), 10).await; let atlas = test_support::open(tmp.path()).await; assert_eq!(atlas.interned_schemas(), 1, "the fixture shares its schema"); - let mut keys = HashSet::new(); - for name in atlas.list_datasets() { - keys.insert(shape_key(&atlas.dataset(&name).unwrap()).await.unwrap()); - } - assert_eq!(keys.len(), 1, "and every dataset reduces to one key"); - + // Ten datasets, one schema. The merge folds the repeats away. let schema = collection_schema(&atlas, None, "c", &widening()) .await .unwrap(); @@ -783,19 +681,6 @@ mod tests { assert_eq!(columns, vec![".platform", "temperature"]); } - /// Two datasets that share arrays but not attribute *keys* produce two - /// column sets, so the key has to separate them. - #[tokio::test] - async fn a_different_attribute_namespace_is_a_different_shape() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - let atlas = test_support::open(tmp.path()).await; - - let winter = shape_key(&atlas.dataset("winter").unwrap()).await.unwrap(); - let summer = shape_key(&atlas.dataset("summer").unwrap()).await.unwrap(); - assert_ne!(winter, summer); - } - // ── dimensions ────────────────────────────────────────────────────── #[tokio::test] diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs index 8631b6b0..66122372 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs @@ -469,3 +469,82 @@ pub async fn wide_profiles(dir: &Path, shapes: &[(usize, usize)]) { writer.finish().await.expect("finish the collection"); } + +/// A number column that one dataset stores as text a cast cannot read. +/// +/// - `a`: `value: Float64[2] = [1.5, 2.5]`. +/// - `b`: `value: String[2] = ["0.-90", "3.5"]`. +/// +/// `Float64` and `String` share no type, so `KeepFirst` keeps `Float64` and +/// marks the column. `'0.-90'` is not a number, so its cast has to read as +/// null. `'3.5'` casts cleanly, which separates "the cast ran" from "the cast +/// gave up on the column". +pub async fn conflicting_numbers(dir: &Path) { + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + + { + let mut a = writer.add_dataset("a").await.expect("add a"); + a.define_array::("value", vec!["obs".into()], vec![2], None, None) + .await + .expect("define value"); + a.write_array("value", vec![0], arr1(&[1.5f64, 2.5]).into_dyn().view()) + .await + .expect("write value"); + a.finish().await.expect("finish a"); + } + + { + let mut b = writer.add_dataset("b").await.expect("add b"); + b.define_array::("value", vec!["obs".into()], vec![2], None, None) + .await + .expect("define value"); + b.write_array( + "value", + vec![0], + arr1(&["0.-90".to_string(), "3.5".to_string()]) + .into_dyn() + .view(), + ) + .await + .expect("write value"); + b.finish().await.expect("finish b"); + } + + writer.finish().await.expect("finish the collection"); +} + +/// One dataset, `d`, holding `value` as a number or as text. +/// +/// Two collections built this way put the conflict across the *collection* +/// merge rather than inside one collection's own merge. +pub async fn value_typed(dir: &Path, as_text: bool) { + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + let mut d = writer.add_dataset("d").await.expect("add d"); + if as_text { + d.define_array::("value", vec!["obs".into()], vec![2], None, None) + .await + .expect("define value"); + d.write_array( + "value", + vec![0], + arr1(&["0.-30".to_string(), "3.5".to_string()]) + .into_dyn() + .view(), + ) + .await + .expect("write value"); + } else { + d.define_array::("value", vec!["obs".into()], vec![2], None, None) + .await + .expect("define value"); + d.write_array("value", vec![0], arr1(&[1.5f64, 2.5]).into_dyn().view()) + .await + .expect("write value"); + } + d.finish().await.expect("finish d"); + writer.finish().await.expect("finish the collection"); +} diff --git a/beacon-server/beacon-server-config/src/lib.rs b/beacon-server/beacon-server-config/src/lib.rs index 8e828974..55a06a23 100644 --- a/beacon-server/beacon-server-config/src/lib.rs +++ b/beacon-server/beacon-server-config/src/lib.rs @@ -515,11 +515,6 @@ struct RawConfig { #[envconfig(from = "BEACON_ATLAS_USE_PRUNING", default = "true")] atlas_use_pruning: bool, - /// Whether `ANALYZE FILES` measures the column ranges of an Atlas - /// collection. They come from its footer, so they cost no array read. - #[envconfig(from = "BEACON_ATLAS_ENABLE_STATISTICS", default = "true")] - atlas_enable_statistics: bool, - /// The batch size for NetCDF reads, in number of rows. This is used for both local and MPIO reads. #[envconfig(from = "BEACON_BATCH_SIZE", default = "64000")] beacon_batch_size: usize, @@ -716,7 +711,6 @@ impl From for Config { use_reader_cache: raw.atlas_use_reader_cache, reader_cache_size: raw.atlas_reader_cache_size, use_pruning: raw.atlas_use_pruning, - enable_statistics: raw.atlas_enable_statistics, }, bbf: BbfConfig { split_streams_slice: raw.bbf_split_streams_slice, From a459478896c4bce9487bc1e267759753bcb342fc Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 8 Sep 2026 16:44:13 +0200 Subject: [PATCH 06/16] wip --- Cargo.lock | 4 +- Cargo.toml | 2 +- beacon-db/beacon-core/src/runtime_builder.rs | 15 +- .../beacon-arrow-atlas/src/compat.rs | 276 ++- .../beacon-arrow-atlas/src/config.rs | 26 - .../src/datafusion/metrics.rs | 12 - .../beacon-arrow-atlas/src/datafusion/mod.rs | 1633 +---------------- .../src/datafusion/opener.rs | 48 + .../src/datafusion/source.rs | 563 +----- .../beacon-arrow-atlas/src/lib.rs | 11 +- .../beacon-arrow-atlas/src/reader.rs | 833 --------- beacon-server/beacon-server-config/src/lib.rs | 29 +- 12 files changed, 381 insertions(+), 3071 deletions(-) delete mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs create mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs delete mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs diff --git a/Cargo.lock b/Cargo.lock index 4753c51a..632bd5e5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1193,9 +1193,9 @@ dependencies = [ [[package]] name = "atlas-rust" -version = "0.17.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4dfadaa65a76f424f053c764cd439ab61c22ffade15ade54fe35432df2561e5b" +checksum = "7af7d700a1036f8630d57847273e9a5f438b6ee0c5361ada94991f46f7feb6ef" dependencies = [ "array-format", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index c4dcaaeb..e4c90b4a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,7 +83,7 @@ object_store = { version = "0.13.1", features = ["aws", "gcp", "azure", "http"] # The atlas container embeds array-format files verbatim, so a change to the # crate's bytes is a change to the format Beacon reads. Pin it exactly. -atlas-rust = "=0.17.0" +atlas-rust = "=0.17.1" oxcdf = { git = "https://github.com/robinskil/oxcdf.git", version = "0.4.0", features = ["async", "object-store", "ndarray"] } diff --git a/beacon-db/beacon-core/src/runtime_builder.rs b/beacon-db/beacon-core/src/runtime_builder.rs index 56e6326e..7a33a004 100644 --- a/beacon-db/beacon-core/src/runtime_builder.rs +++ b/beacon-db/beacon-core/src/runtime_builder.rs @@ -1,12 +1,11 @@ use std::{ - collections::HashMap, path::PathBuf, sync::{Arc, OnceLock}, }; use crate::crawler::{new_crawler_manager_handle, CrawlerConfig, CrawlerManager}; use crate::schema_persistence::{init_tables, PersistentSchemaProvider}; -use beacon_arrow_atlas::{AtlasConfig, AtlasFormatFactory, AtlasOptions}; +use beacon_arrow_atlas::{AtlasFormatFactory, AtlasOptions}; use beacon_arrow_bbf::datafusion::BBFFormatFactory; use beacon_arrow_csv::datafusion::CsvFormatFactory; use beacon_arrow_geoparquet::datafusion::GeoParquetFormatFactory; @@ -99,7 +98,6 @@ pub struct RuntimeBuilder { pub netcdf: NetcdfConfig, pub hdf5: Hdf5Config, pub zarr: ZarrConfig, - pub atlas: AtlasConfig, pub auth_provider: Option>, pub secrets_encryption_key: Option<[u8; 32]>, @@ -238,12 +236,6 @@ impl RuntimeBuilder { self } - /// Replaces the whole Atlas reader configuration. - pub fn with_atlas_config(mut self, atlas: AtlasConfig) -> Self { - self.atlas = atlas; - self - } - pub fn with_auth_provider(mut self, provider: Arc) -> Self { self.auth_provider = Some(provider); self @@ -799,10 +791,7 @@ fn register_file_formats( Arc::new(ArrowFormatFactory), Arc::new(TiffFormatFactory::new(Default::default())), Arc::new(ZarrFormatFactory::new(builder.zarr.clone())), - Arc::new(AtlasFormatFactory::new( - AtlasOptions::default(), - builder.atlas.clone(), - )), + Arc::new(AtlasFormatFactory::new(AtlasOptions::default())), Arc::new(BBFFormatFactory::new(Default::default())), Arc::new(GeoParquetFormatFactory::default()), Arc::new(NetCDFFormatFactory::new( diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs index d8377bbb..3dede666 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs @@ -5,10 +5,13 @@ //! [`NdArrayDataType`] through `beacon-nd-array`'s own conversion, so a schema //! derived here and a batch produced by a scan can never disagree. +use std::collections::BTreeMap; use std::sync::Arc; -use arrow::datatypes::DataType; -use atlas::{ArrayLayout, Attr, DType, DatasetView, FillValue}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::error::ArrowError; +use atlas::{ArrayLayout, Attr, CollectionSchema, DType, DatasetView, FillValue}; +use beacon_datafusion_ext::type_widening::{ArrowTypeWidening, LabeledSchema}; use beacon_nd_array::{ NdArray, NdArrayD, datatypes::NdArrayDataType, datatypes::TimestampNanosecond, }; @@ -100,6 +103,101 @@ pub(crate) fn dtype_tag(dtype: &DType) -> String { format!("{dtype:?}") } +// ─── Collection schema ─────────────────────────────────────────────────────── + +/// The Arrow schema of one collection, from its footer alone. +/// +/// One nullable field per array, under the array's own name. A dataset +/// attribute becomes `.{attr}`, and an array attribute `{array}.{attr}`. A +/// name that two datasets type differently takes the type `widening` gives +/// the set, with the conflict mark it applies. A dtype Beacon cannot read is +/// dropped with a `debug` log. +/// +/// Every dataset in the container counts, deleted ones too. A column only a +/// deleted dataset declares reads as null. `read_dimensions` does not narrow +/// this schema: the footer holds no dimension name. +/// +/// The fields are sorted by name. Atlas permits an array named `.season` or +/// `temperature.units`, so one column name can come from two maps. Their types +/// then merge as one column. +pub fn collection_arrow_schema( + schema: &CollectionSchema<'_>, + widening: &ArrowTypeWidening, +) -> Result { + let mut columns: BTreeMap> = BTreeMap::new(); + for (array, dtypes) in &schema.arrays { + let types = readable_types(array, dtypes, array_dtype_to_arrow); + columns.entry(array.to_string()).or_default().extend(types); + } + for (key, dtypes) in &schema.attributes { + let column = global_attr_column(key); + let types = readable_types(&column, dtypes, attr_dtype_to_arrow); + columns.entry(column).or_default().extend(types); + } + for (array, attributes) in &schema.array_attributes { + for (key, dtypes) in attributes { + let column = array_attr_column(array, key); + let types = readable_types(&column, dtypes, attr_dtype_to_arrow); + columns.entry(column).or_default().extend(types); + } + } + + let mut fields = Vec::with_capacity(columns.len()); + for (name, types) in &columns { + if types.is_empty() { + continue; + } + fields.push(merge_types(widening, name, types)?); + } + Ok(Schema::new(fields)) +} + +/// The Arrow types of `dtypes` that Beacon can read as column `column`. +/// +/// A dtype `to_arrow` refuses is logged at `debug` and dropped. A collection +/// can hold a million datasets, so a `warn` per skip would be a flood. +fn readable_types( + column: &str, + dtypes: &[&DType], + to_arrow: fn(&DType) -> Option, +) -> Vec { + dtypes + .iter() + .filter_map(|dtype| { + let data_type = to_arrow(dtype); + if data_type.is_none() { + tracing::debug!(column, ?dtype, "no column for this atlas dtype, skipped"); + } + data_type + }) + .collect() +} + +/// The field column `name` takes when its sources state `types`. +/// +/// One nullable single-field schema per type, merged under the session rule. +/// That is [`ArrowTypeWidening::merge_schemas`] for one column: the same +/// widening, the same conflict setting, and the same conflict mark, which the +/// scan reads to cast a source the type cannot hold as null. +fn merge_types( + widening: &ArrowTypeWidening, + name: &str, + types: &[DataType], +) -> Result { + if let [only] = types { + return Ok(Field::new(name, only.clone(), true)); + } + let schemas: Vec = types + .iter() + .map(|data_type| { + let field = Field::new(name, data_type.clone(), true); + LabeledSchema::unlabeled(Arc::new(Schema::new(vec![field]))) + }) + .collect(); + let merged = widening.merge_schemas(&schemas)?; + Ok(merged.field(0).clone()) +} + // ─── Lazy arrays ───────────────────────────────────────────────────────────── /// Wrap one atlas array as a lazy [`NdArrayD`] over `view`. @@ -319,4 +417,178 @@ mod tests { .to_string(); assert!(error.contains("list"), "{error}"); } + + // ── the collection schema ─────────────────────────────────────────── + + use crate::test_support; + use arrow::datatypes::TimeUnit; + use beacon_datafusion_ext::type_widening::{DefaultArrowTypeWidening, is_type_conflict}; + + fn widening() -> Arc { + ArrowTypeWidening::default_extension() + } + + fn names(schema: &Schema) -> Vec<&str> { + schema.fields().iter().map(|f| f.name().as_str()).collect() + } + + /// Every array and every attribute of every dataset is a column, in name + /// order. The footer's maps have no order of their own. + #[tokio::test] + async fn a_collections_schema_is_the_union_of_its_datasets_in_name_order() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let schema = + collection_arrow_schema(&atlas.footer().collection_schema(), &widening()).unwrap(); + + assert_eq!( + names(&schema), + vec![ + ".season", + ".year", + "cycle", + "temperature", + "temperature.units", + "time" + ] + ); + } + + #[tokio::test] + async fn every_column_keeps_the_type_the_footer_gave_it() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let schema = + collection_arrow_schema(&atlas.footer().collection_schema(), &widening()).unwrap(); + let field = |name: &str| schema.field_with_name(name).unwrap(); + + assert_eq!(field("temperature").data_type(), &DataType::Float32); + assert_eq!(field("cycle").data_type(), &DataType::Int32); + assert_eq!( + field("time").data_type(), + &DataType::Timestamp(TimeUnit::Nanosecond, None) + ); + assert_eq!(field(".year").data_type(), &DataType::Int64); + assert_eq!(field(".season").data_type(), &DataType::Utf8); + assert_eq!(field("temperature.units").data_type(), &DataType::Utf8); + assert!( + schema.fields().iter().all(|f| f.is_nullable()), + "a dataset may lack any column, so every column is nullable" + ); + } + + /// Two datasets that give one array two numeric types merge to the type + /// that holds both, by the rule of the session rather than one of atlas's + /// own. `Int16` beside `Float32` gives `Float64`. See issue #377. + #[tokio::test] + async fn a_shared_array_widens_to_a_type_that_holds_both() { + let tmp = tempfile::tempdir().unwrap(); + test_support::widening(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let schema = + collection_arrow_schema(&atlas.footer().collection_schema(), &widening()).unwrap(); + + assert_eq!( + schema.field_with_name("value").unwrap().data_type(), + &DataType::Float64, + "Int16 and Float32 widen to Float64" + ); + assert_eq!( + schema.field_with_name("flag").unwrap().data_type(), + &DataType::Int32, + "a column only one dataset declares keeps its own type" + ); + } + + /// A column two datasets type in two families is refused, and the error + /// names the column and both types. The footer's type set names no + /// dataset, so the error cannot. + #[tokio::test] + async fn types_that_do_not_widen_are_refused_by_name() { + let tmp = tempfile::tempdir().unwrap(); + test_support::incompatible(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let error = collection_arrow_schema(&atlas.footer().collection_schema(), &widening()) + .expect_err("Utf8 and Int64 are two families") + .to_string(); + + assert!(error.contains("value"), "the column: {error}"); + assert!( + error.contains("Utf8") && error.contains("Int64"), + "both types: {error}" + ); + } + + /// A deployment that reads such a collection anyway sets `keep_first`. The + /// column then takes the type the footer states first, which is the type + /// of the dataset written first, and carries the mark the scan reads. + #[tokio::test] + async fn keep_first_settles_a_conflict_with_the_first_type_and_marks_it() { + let tmp = tempfile::tempdir().unwrap(); + test_support::incompatible(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + let keep_first = + ArrowTypeWidening::new(Arc::new(DefaultArrowTypeWidening::keeping_first_type())); + + let schema = + collection_arrow_schema(&atlas.footer().collection_schema(), &keep_first).unwrap(); + let value = schema.field_with_name("value").unwrap(); + + assert_eq!(value.data_type(), &DataType::Utf8, "dataset `a` came first"); + assert!(is_type_conflict(value), "the scan must cast `b` to null"); + assert_eq!( + schema.field_with_name("only_a").unwrap().data_type(), + &DataType::Int32 + ); + } + + #[tokio::test] + async fn a_list_attribute_is_dropped_and_the_rest_survives() { + let tmp = tempfile::tempdir().unwrap(); + test_support::skips(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let schema = + collection_arrow_schema(&atlas.footer().collection_schema(), &widening()).unwrap(); + + assert_eq!(names(&schema), vec![".title", "value", "value.units"]); + } + + #[tokio::test] + async fn an_empty_collection_has_no_column() { + let tmp = tempfile::tempdir().unwrap(); + test_support::empty(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + + let schema = + collection_arrow_schema(&atlas.footer().collection_schema(), &widening()).unwrap(); + + assert!(schema.fields().is_empty(), "{:?}", names(&schema)); + } + + /// The footer reports every dataset the container holds. A deleted one + /// still shapes the schema: its columns stay, and read as null. + #[tokio::test] + async fn a_deleted_dataset_still_shapes_the_schema() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + atlas.delete_dataset("winter").await.unwrap(); + let atlas = test_support::open(tmp.path()).await; + + let schema = + collection_arrow_schema(&atlas.footer().collection_schema(), &widening()).unwrap(); + + assert!( + schema.field_with_name("cycle").is_ok(), + "only `winter` declares `cycle`: {:?}", + names(&schema) + ); + } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs deleted file mode 100644 index e0a831a3..00000000 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/config.rs +++ /dev/null @@ -1,26 +0,0 @@ -//! [`AtlasConfig`]: the runtime settings of the Atlas format. - -/// Runtime configuration for the Atlas format. -/// -/// Plain data with sensible defaults; the caller populates it. There is no -/// environment parsing here, so the crate stays reusable and the host decides -/// where the values come from. Each field is the *default* for a runtime, and -/// each can be overridden per table via -/// `CREATE EXTERNAL TABLE ... OPTIONS (...)`. -#[derive(Debug, Clone)] -pub struct AtlasConfig { - /// Whether a read consults the shared reader cache. - pub use_reader_cache: bool, - pub reader_cache_size: u64, - pub use_pruning: bool, -} - -impl Default for AtlasConfig { - fn default() -> Self { - Self { - use_reader_cache: true, - reader_cache_size: 32, - use_pruning: true, - } - } -} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs index c2cef023..22323d5c 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs @@ -1,15 +1,3 @@ -//! What one Atlas scan partition did, reported through DataFusion's metrics. -//! -//! These complement -//! [`ReadMetrics`](beacon_nd_array::arrow::metrics::ReadMetrics), which counts -//! the chunks and rows the shared queue handed out. What it cannot see is the -//! cost of reaching a dataset at all: opening the collection, deciding whether -//! the dataset is worth reading, and building its lazy columns. -//! -//! Every name is `atlas_`-prefixed. DataFusion sums metrics that share a name, -//! and `output_rows` and `output_batches` are already registered for this -//! partition by the scan itself. - use datafusion::physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder, Time}; /// Per-partition timings and counts for one Atlas scan partition. diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs index 9d2c272d..adc14497 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs @@ -34,11 +34,11 @@ use datafusion::{ }; use object_store::{ObjectMeta, ObjectStore}; -use crate::config::AtlasConfig; -use crate::reader::collection_schema; +use crate::compat; use crate::store::{ATLAS_MARKER, AtlasReaderCache, get_or_open_atlas, top_level_atlas_markers}; pub mod metrics; +pub mod opener; pub mod options; pub mod pruning; pub mod source; @@ -51,53 +51,22 @@ pub use table_function::ReadAtlasFunc; /// The name this format answers to: `STORED AS ATLAS`, `read_atlas`. pub const ATLAS_FORMAT: &str = "atlas"; -/// Parse a boolean supplied through `CREATE EXTERNAL TABLE ... OPTIONS`. -fn parse_bool_option(key: &str, value: &str) -> Result { - match value.trim().to_ascii_lowercase().as_str() { - "true" | "1" | "yes" | "on" => Ok(true), - "false" | "0" | "no" | "off" => Ok(false), - other => Err(exec_datafusion_err!( - "invalid boolean for atlas option '{key}': '{other}'" - )), - } -} - -// ─── Factory ───────────────────────────────────────────────────────────────── - /// Builds an [`AtlasFormat`] per table, over one runtime's settings and one /// shared reader cache. #[derive(Debug, Clone)] pub struct AtlasFormatFactory { pub options: AtlasOptions, - pub config: AtlasConfig, - /// The runtime's reader cache, sized from `config` and shared by every - /// format, source and opener this factory builds. - cache: AtlasReaderCache, } impl AtlasFormatFactory { - pub fn new(options: AtlasOptions, config: AtlasConfig) -> Self { - let cache = AtlasReaderCache::new(config.reader_cache_size); - Self { - options, - config, - cache, - } + pub fn new(options: AtlasOptions) -> Self { + Self { options } } /// A format with this table's effective settings, wired to the shared cache /// when caching is on. - pub(crate) fn build( - &self, - options: AtlasOptions, - use_reader_cache: bool, - use_pruning: bool, - ) -> AtlasFormat { - AtlasFormat { - options, - cache: use_reader_cache.then(|| self.cache.clone()), - use_pruning, - } + pub(crate) fn build(&self, options: AtlasOptions) -> AtlasFormat { + AtlasFormat::new(options) } } @@ -108,8 +77,6 @@ impl FileFormatFactory for AtlasFormatFactory { format_options: &HashMap, ) -> Result> { let mut options = self.options.clone(); - let mut use_reader_cache = self.config.use_reader_cache; - let mut use_pruning = self.config.use_pruning; if let Some(value) = format_option(format_options, "read_dimensions") { options.read_dimensions = Some( @@ -120,21 +87,11 @@ impl FileFormatFactory for AtlasFormatFactory { .collect(), ); } - if let Some(value) = format_option(format_options, "use_reader_cache") { - use_reader_cache = parse_bool_option("use_reader_cache", value)?; - } - if let Some(value) = format_option(format_options, "use_pruning") { - use_pruning = parse_bool_option("use_pruning", value)?; - } - Ok(Arc::new(self.build(options, use_reader_cache, use_pruning))) + Ok(Arc::new(self.build(options))) } fn default(&self) -> Arc { - Arc::new(self.build( - self.options.clone(), - self.config.use_reader_cache, - self.config.use_pruning, - )) + Arc::new(self.build(self.options.clone())) } fn as_any(&self) -> &dyn Any { @@ -199,16 +156,11 @@ impl FileFormatFactoryExt for AtlasFormatFactory { } } -// ─── Format ────────────────────────────────────────────────────────────────── - /// Reads one table's worth of Atlas collections. #[derive(Debug, Clone)] pub struct AtlasFormat { pub options: AtlasOptions, - /// The reader cache to consult, or `None` to bypass caching. - cache: Option, - /// Whether a predicate scan drops the datasets it can rule out. - use_pruning: bool, + cache: AtlasReaderCache, } impl Default for AtlasFormat { @@ -219,12 +171,9 @@ impl Default for AtlasFormat { impl AtlasFormat { pub fn new(options: AtlasOptions) -> Self { - let defaults = AtlasConfig::default(); Self { options, - cache: None, - // A query prunes by default: it only ever saves reads. - use_pruning: defaults.use_pruning, + cache: AtlasReaderCache::new(512), } } } @@ -272,6 +221,11 @@ impl FileFormat for AtlasFormat { /// Each collection costs one open — a footer read — and the datasets behind /// it cost no I/O at all. Never enumerate a collection's datasets any other /// way here: at a million datasets that would turn planning into a scan. + /// + /// The footer counts every dataset, deleted ones too, so a column only a + /// deleted dataset declares is in the schema and reads as null. + /// `read_dimensions` does not narrow the schema: the footer holds no + /// dimension name. The scan fills an array it drops with nulls. async fn infer_schema( &self, state: &dyn Session, @@ -287,20 +241,25 @@ impl FileFormat for AtlasFormat { // One rule for both merges: the datasets inside a collection, and the // collections of this table. let widening = session_widening(state); - let read_dimensions = self.options.read_dimensions.clone(); let mut schemas = Vec::with_capacity(markers.len()); for marker in &markers { - let atlas = get_or_open_atlas(self.cache.as_ref(), Arc::clone(store), marker) - .await - .map_err(|e| exec_datafusion_err!("{e}"))?; - let label = marker.location.as_ref(); - let schema = collection_schema(&atlas, read_dimensions.as_deref(), label, &widening) + let atlas = get_or_open_atlas(Some(&self.cache), Arc::clone(store), marker) .await .map_err(|e| exec_datafusion_err!("{e}"))?; - // The container names the schema, so a refused column names both - // collections. - schemas.push(LabeledSchema::new(schema, label)); + + let schema = + compat::collection_arrow_schema(&atlas.footer().collection_schema(), &widening) + .map_err(|e| { + exec_datafusion_err!( + "Failed to read the schema of atlas collection '{}': {e}", + marker.location + ) + })?; + schemas.push(LabeledSchema::new( + Arc::new(schema), + marker.location.as_ref(), + )); } let schema = widening.merge_schemas(&schemas).map_err(|e| { @@ -384,10 +343,12 @@ impl FileFormat for AtlasFormat { // rebuilding it below would otherwise drop it. let projection = conf.file_source().projection().cloned(); - let source = AtlasSource::new(self.options.read_dimensions.clone(), table_schema) - .with_cache(self.cache.clone()) - .with_pruning(self.use_pruning) - .with_projection(projection); + let source = AtlasSource::new( + self.options.read_dimensions.clone(), + table_schema, + self.cache.clone(), + ) + .with_projection(projection); let conf = FileScanConfigBuilder::from(conf) .with_file_groups(file_groups) .with_source(Arc::new(source)) @@ -397,11 +358,11 @@ impl FileFormat for AtlasFormat { } fn file_source(&self, table_schema: TableSchema) -> Arc { - Arc::new( - AtlasSource::new(self.options.read_dimensions.clone(), table_schema) - .with_cache(self.cache.clone()) - .with_pruning(self.use_pruning), - ) + Arc::new(AtlasSource::new( + self.options.read_dimensions.clone(), + table_schema, + self.cache.clone(), + )) } async fn create_writer_physical_plan( @@ -417,1517 +378,3 @@ impl FileFormat for AtlasFormat { )) } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::test_support; - use arrow::datatypes::DataType; - use datafusion::datasource::listing::{ - ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, - }; - use datafusion::physical_plan::ExecutionPlan; - use datafusion::prelude::{SessionConfig, SessionContext}; - use std::path::Path; - - /// Register the collection in `dir` as `name`, through a listing table. - async fn register(ctx: &SessionContext, dir: &Path, name: &str) { - let format: Arc = Arc::new(AtlasFormat::default()); - register_with(ctx, dir, name, format).await; - } - - async fn register_with( - ctx: &SessionContext, - dir: &Path, - name: &str, - format: Arc, - ) { - let directory = dir.to_string_lossy().replace('\\', "/"); - let url = ListingTableUrl::parse(format!("file://{directory}/")).unwrap(); - let listing = ListingOptions::new(format).with_file_extension(ATLAS_MARKER); - let config = ListingTableConfig::new(url) - .with_listing_options(listing) - .infer_schema(&ctx.state()) - .await - .expect("the collection types"); - ctx.register_table(name, Arc::new(ListingTable::try_new(config).unwrap())) - .unwrap(); - } - - fn context(partitions: usize) -> SessionContext { - SessionContext::new_with_config(SessionConfig::new().with_target_partitions(partitions)) - } - - async fn rows(ctx: &SessionContext, sql: &str) -> usize { - ctx.sql(sql) - .await - .unwrap() - .collect() - .await - .unwrap() - .iter() - .map(|batch| batch.num_rows()) - .sum() - } - - async fn count(ctx: &SessionContext, sql: &str) -> i64 { - use arrow::array::Int64Array; - let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); - batches[0] - .column(0) - .as_any() - .downcast_ref::() - .expect("a count is an i64") - .value(0) - } - - // ── discovery ─────────────────────────────────────────────────────── - - #[test] - fn the_factory_answers_to_atlas() { - let factory = AtlasFormatFactory::new(Default::default(), Default::default()); - assert_eq!(factory.get_ext(), "atlas"); - assert_eq!(factory.file_format_name(), "atlas"); - assert_eq!(AtlasFormat::default().get_ext(), "data.atlas"); - } - - #[test] - fn one_dataset_entry_per_collection() { - fn object(path: &str) -> ObjectMeta { - ObjectMeta { - location: object_store::path::Path::from(path), - last_modified: Default::default(), - size: 0, - e_tag: None, - version: None, - } - } - - let factory = AtlasFormatFactory::new(Default::default(), Default::default()); - let discovered = factory - .discover_datasets(&[ - object("a/data.atlas"), - object("a/deleted.mask"), - object("b/data.atlas"), - object("b/notes.txt"), - ]) - .unwrap(); - - let paths: Vec<&str> = discovered.iter().map(|d| d.file_path.as_str()).collect(); - assert_eq!(paths, vec!["a/data.atlas", "b/data.atlas"]); - assert!(discovered.iter().all(|d| d.format == "atlas")); - } - - // ── reading, end to end ───────────────────────────────────────────── - - #[tokio::test] - async fn every_dataset_of_a_collection_is_read() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - let ctx = context(1); - register(&ctx, tmp.path(), "obs").await; - - // winter contributes 4 rows and summer 3. - assert_eq!(rows(&ctx, "SELECT temperature FROM obs").await, 7); - } - - #[tokio::test] - async fn count_star_counts_every_dataset() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - let ctx = context(1); - register(&ctx, tmp.path(), "obs").await; - - assert_eq!(count(&ctx, "SELECT COUNT(*) FROM obs").await, 7); - } - - /// The plan is the nd spine over the scan, in that nesting order. - #[tokio::test] - async fn the_plan_is_the_nd_spine_over_the_scan() { - use datafusion::physical_plan::displayable; - - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - let ctx = context(1); - register(&ctx, tmp.path(), "obs").await; - - let plan = ctx - .sql("SELECT temperature FROM obs") - .await - .unwrap() - .create_physical_plan() - .await - .unwrap(); - let rendered = displayable(plan.as_ref()).indent(true).to_string(); - - let broadcast = rendered.find("NdBroadcastExec"); - let source = rendered.find("NdSourceExec"); - let scan = rendered.find("DataSourceExec"); - assert!( - broadcast.is_some() && source.is_some() && scan.is_some(), - "the spine must be present:\n{rendered}" - ); - assert!( - broadcast < source && source < scan, - "expected NdBroadcastExec over NdSourceExec over DataSourceExec:\n{rendered}" - ); - } - - #[tokio::test] - async fn a_projection_reaches_the_result() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - let ctx = context(1); - register(&ctx, tmp.path(), "obs").await; - - let df = ctx.sql("SELECT temperature FROM obs").await.unwrap(); - let columns: Vec = df - .schema() - .fields() - .iter() - .map(|field| field.name().clone()) - .collect(); - assert_eq!(columns, vec!["temperature".to_string()]); - } - - /// An attribute rides along as a constant column on every row its dataset - /// contributes. - #[tokio::test] - async fn an_attribute_is_constant_across_its_datasets_rows() { - use arrow::array::{Array, StringArray}; - - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - let ctx = context(1); - register(&ctx, tmp.path(), "obs").await; - - let batches = ctx - .sql(r#"SELECT ".season" AS season FROM obs WHERE temperature < 10"#) - .await - .unwrap() - .collect() - .await - .unwrap(); - - let mut seen = Vec::new(); - for batch in &batches { - let column = batch - .column(0) - .as_any() - .downcast_ref::() - .expect("a text column"); - for row in 0..column.len() { - seen.push(column.value(row).to_string()); - } - } - // Only winter's four rows are below 10 degrees. - assert_eq!(seen, vec!["winter".to_string(); 4]); - } - - // ── datasets that disagree ────────────────────────────────────────── - - #[tokio::test] - async fn a_widened_column_is_cast_from_each_dataset() { - use arrow::array::{Array, Float64Array}; - - let tmp = tempfile::tempdir().unwrap(); - test_support::widening(tmp.path()).await; - let ctx = context(1); - register(&ctx, tmp.path(), "w").await; - - let df = ctx.sql("SELECT value FROM w ORDER BY value").await.unwrap(); - assert_eq!( - df.schema() - .field_with_unqualified_name("value") - .unwrap() - .data_type(), - &DataType::Float64 - ); - - let batches = df.collect().await.unwrap(); - let mut values = Vec::new(); - for batch in &batches { - let column = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - values.extend(column.iter().flatten()); - } - // a.value = [1, 2] as Int16, b.value = [3.5, 4.5] as Float32. - assert_eq!(values, vec![1.0, 2.0, 3.5, 4.5]); - } - - /// A dataset that lacks a projected column contributes its rows with that - /// column null, rather than dropping them. - #[tokio::test] - async fn a_column_one_dataset_lacks_is_null_filled() { - let tmp = tempfile::tempdir().unwrap(); - test_support::widening(tmp.path()).await; - let ctx = context(1); - register(&ctx, tmp.path(), "w").await; - - let batches = ctx - .sql("SELECT value, flag FROM w") - .await - .unwrap() - .collect() - .await - .unwrap(); - let (rows, nulls) = batches.iter().fold((0, 0), |(rows, nulls), batch| { - ( - rows + batch.num_rows(), - nulls + batch.column(1).null_count(), - ) - }); - assert_eq!(rows, 4, "both datasets contribute their rows"); - assert_eq!(nulls, 2, "dataset b declares no flag"); - } - - // ── the deletion mask ─────────────────────────────────────────────── - - /// A deleted dataset is gone from the result, and its rows with it. - #[tokio::test] - async fn a_deleted_dataset_is_not_read() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - test_support::open(tmp.path()) - .await - .delete_dataset("winter") - .await - .expect("delete winter"); - - let ctx = context(1); - register(&ctx, tmp.path(), "obs").await; - // Summer's three rows alone. - assert_eq!(rows(&ctx, "SELECT temperature FROM obs").await, 3); - } - - // ── dividing the scan ─────────────────────────────────────────────── - - /// Every row is read exactly once, however many partitions share the - /// collection. A dataset popped twice is a row returned twice, and one - /// popped by nobody is a row lost; neither raises an error. - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn a_partitioned_scan_reads_every_row_once() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 12).await; - - for partitions in [1_usize, 2, 4, 8] { - let ctx = context(partitions); - register(&ctx, tmp.path(), "ranged").await; - assert_eq!( - rows(&ctx, "SELECT temperature FROM ranged").await, - 48, - "partitions={partitions}: 12 datasets of 4 rows" - ); - assert_eq!( - count(&ctx, "SELECT COUNT(*) FROM ranged").await, - 48, - "partitions={partitions}: and the count agrees" - ); - } - } - - /// The unit of work is the collection, so one collection is one partition - /// however many datasets it holds. - #[tokio::test] - async fn one_collection_is_one_partition() { - use datafusion::physical_plan::ExecutionPlanProperties; - - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 12).await; - let ctx = context(4); - register(&ctx, tmp.path(), "ranged").await; - - let plan = ctx - .sql("SELECT temperature FROM ranged") - .await - .unwrap() - .create_physical_plan() - .await - .unwrap(); - - // The count comes off the scan, not the plan root: DataFusion adds a - // round robin above a single-partition scan either way. - let mut scan = Arc::clone(&plan); - while let Some(child) = scan.children().first() { - scan = Arc::clone(child); - } - assert_eq!( - scan.output_partitioning().partition_count(), - 1, - "twelve datasets of one collection stay together: -{}", - datafusion::physical_plan::displayable(plan.as_ref()).indent(false) - ); - } - - /// Several collections deal round-robin over the partitions, and never - /// past the collection count. - #[tokio::test] - async fn collections_deal_across_the_partitions() { - use datafusion::physical_plan::ExecutionPlanProperties; - - let tmp = tempfile::tempdir().unwrap(); - for name in ["a", "b", "c"] { - test_support::ranged(&tmp.path().join(name), 2).await; - } - let root = tmp - .path() - .to_string_lossy() - .replace(std::path::MAIN_SEPARATOR, "/"); - - for (target, expected) in [(8, 3), (2, 2), (1, 1)] { - let ctx = ddl_context(target); - ctx.sql(&format!( - "CREATE EXTERNAL TABLE t STORED AS ATLAS LOCATION '{root}/**/data.atlas'" - )) - .await - .unwrap() - .collect() - .await - .unwrap(); - - let plan = ctx - .sql("SELECT temperature FROM t") - .await - .unwrap() - .create_physical_plan() - .await - .unwrap(); - let mut scan = Arc::clone(&plan); - while let Some(child) = scan.children().first() { - scan = Arc::clone(child); - } - assert_eq!( - scan.output_partitioning().partition_count(), - expected, - "three collections over {target} target partitions" - ); - } - } - - // ── predicates ────────────────────────────────────────────────────── - - #[tokio::test] - async fn a_predicate_keeps_only_the_rows_that_match() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 10).await; - let ctx = context(2); - register(&ctx, tmp.path(), "ranged").await; - - // d5..d9 hold [50..53] … [90..93]: 20 rows above 45. - assert_eq!( - rows( - &ctx, - "SELECT temperature FROM ranged WHERE temperature > 45" - ) - .await, - 20 - ); - assert_eq!( - rows( - &ctx, - "SELECT temperature FROM ranged WHERE temperature > 100000" - ) - .await, - 0, - "a predicate nothing meets returns nothing" - ); - } - - // ── pruning, end to end ───────────────────────────────────────────── - - /// A format with pruning on or off, built the way a table is. - /// - /// `use_pruning` reaches a format through the factory alone, as - /// `CREATE EXTERNAL TABLE ... OPTIONS ('use_pruning' '...')` does. - fn format_with_pruning(use_pruning: bool) -> Arc { - Arc::new( - AtlasFormatFactory::new(Default::default(), Default::default()).build( - AtlasOptions::default(), - false, - use_pruning, - ), - ) - } - - /// Register the collection twice, once pruning and once not. - async fn register_pruning(ctx: &SessionContext, dir: &Path, name: &str, use_pruning: bool) { - register_with(ctx, dir, name, format_with_pruning(use_pruning)).await; - } - - async fn values(ctx: &SessionContext, sql: &str) -> Vec { - use arrow::array::Float32Array; - let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); - let mut out = Vec::new(); - for batch in &batches { - let column = batch - .column(0) - .as_any() - .downcast_ref::() - .expect("a float column"); - out.extend(column.iter().flatten()); - } - out - } - - /// The switch changes what is read, never what is returned. - /// - /// This is the property pruning has to hold above all others: it drops - /// datasets that cannot contain a matching row, and nothing else. - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn pruning_does_not_change_the_answer() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 10).await; - - for predicate in [ - "temperature > 45", - "temperature < 25", - "temperature > 1000", - "temperature >= 0", - "temperature > 45 AND temperature < 75", - ] { - let sql = - format!("SELECT temperature FROM ranged WHERE {predicate} ORDER BY temperature"); - - let on = context(4); - register_pruning(&on, tmp.path(), "ranged", true).await; - let off = context(4); - register_pruning(&off, tmp.path(), "ranged", false).await; - - assert_eq!( - values(&on, &sql).await, - values(&off, &sql).await, - "pruning changed the answer for `{predicate}`" - ); - } - } - - /// Find the scan's metrics by the names only this format registers. - fn atlas_metrics( - plan: &Arc, - ) -> Option { - if let Some(metrics) = plan.metrics() - && metrics.sum_by_name("atlas_datasets_scanned").is_some() - { - return Some(metrics); - } - plan.children().into_iter().find_map(atlas_metrics) - } - - /// The scan reports what it read and what it skipped, and the two add up to - /// the collection. - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn the_metrics_report_what_was_read_and_what_was_skipped() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 10).await; - let ctx = context(4); - register_pruning(&ctx, tmp.path(), "ranged", true).await; - - let plan = ctx - .sql("SELECT temperature FROM ranged WHERE temperature > 45") - .await - .unwrap() - .create_physical_plan() - .await - .unwrap(); - datafusion::physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()) - .await - .unwrap(); - - let metrics = atlas_metrics(&plan).expect("the atlas scan reports metrics"); - let sum = |name: &str| metrics.sum_by_name(name).map(|value| value.as_usize()); - - // d5..d9 hold values above 45; d0..d4 cannot. - assert_eq!(sum("atlas_datasets_scanned"), Some(5)); - assert_eq!(sum("atlas_datasets_pruned"), Some(5)); - assert_eq!( - sum("atlas_index_rows"), - Some(10), - "the index covered them all" - ); - assert!(metrics.sum_by_name("atlas_prune_time").is_some()); - } - - /// One index per collection, however many partitions share it. - /// - /// Every partition's opener holds the same memo, so the first to reach the - /// collection builds the index and the rest await it. Without that, a - /// twenty-four-partition scan would build it twenty-four times. - #[tokio::test(flavor = "multi_thread", worker_threads = 8)] - async fn one_index_is_built_per_collection() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 24).await; - let ctx = context(8); - register_pruning(&ctx, tmp.path(), "ranged", true).await; - - let plan = ctx - .sql("SELECT temperature FROM ranged WHERE temperature > 100") - .await - .unwrap() - .create_physical_plan() - .await - .unwrap(); - datafusion::physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()) - .await - .unwrap(); - - let metrics = atlas_metrics(&plan).expect("the atlas scan reports metrics"); - assert_eq!( - metrics - .sum_by_name("atlas_index_builds") - .map(|value| value.as_usize()), - Some(1), - "eight partitions must share one index" - ); - } - - /// A scan with no predicate builds no index at all. - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn a_scan_without_a_predicate_builds_no_index() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 6).await; - let ctx = context(4); - register_pruning(&ctx, tmp.path(), "ranged", true).await; - - let plan = ctx - .sql("SELECT temperature FROM ranged") - .await - .unwrap() - .create_physical_plan() - .await - .unwrap(); - datafusion::physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()) - .await - .unwrap(); - - let metrics = atlas_metrics(&plan).expect("the atlas scan reports metrics"); - assert_eq!( - metrics - .sum_by_name("atlas_index_builds") - .map(|value| value.as_usize()), - Some(0) - ); - assert_eq!( - metrics - .sum_by_name("atlas_datasets_pruned") - .map(|value| value.as_usize()), - Some(0) - ); - } - - // ── measuring a collection ────────────────────────────────────────── - - /// Atlas measures no column, for any caller. - /// - /// A recorded range prunes whole collections before a scan opens them, so a - /// range that is too narrow deletes matching rows from an answer. The scan - /// prunes from the footer instead, per dataset, where the numbers are exact - /// and cost no array read. - #[tokio::test] - async fn a_collection_reports_no_column_range() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 4).await; - let (store, marker) = test_support::store_and_marker(tmp.path()); - let ctx = SessionContext::new(); - - let format = AtlasFormat::default(); - let schema = format - .infer_schema(&ctx.state(), &store, std::slice::from_ref(&marker)) - .await - .unwrap(); - let statistics = format - .infer_stats(&ctx.state(), &store, Arc::clone(&schema), &marker) - .await - .unwrap(); - - assert_eq!(statistics.column_statistics.len(), schema.fields().len()); - for (column, field) in statistics.column_statistics.iter().zip(schema.fields()) { - assert!( - column.min_value.is_exact().is_none(), - "{} reports a minimum", - field.name() - ); - assert!( - column.max_value.is_exact().is_none(), - "{} reports a maximum", - field.name() - ); - } - assert!(statistics.num_rows.is_exact().is_none()); - } - - /// The analyzer gets the same format a query does, because neither - /// measures anything. - #[test] - fn analysis_asks_for_no_measurement() { - let ctx = SessionContext::new(); - let listing = Arc::new(ListingFactory::dynamic()); - let url = ListingTableUrl::parse("file:///tmp/").unwrap(); - let factory = AtlasFormatFactory::new(Default::default(), Default::default()); - - let analysis = factory - .create_for_analysis(&ctx.state(), &HashMap::new(), &url, &listing) - .unwrap(); - assert!(analysis.as_any().downcast_ref::().is_some()); - } - - // ── dimensions ────────────────────────────────────────────────────── - - #[tokio::test] - async fn read_dimensions_narrow_the_table() { - let tmp = tempfile::tempdir().unwrap(); - test_support::chunked_grid(tmp.path()).await; - let ctx = context(1); - let format: Arc = Arc::new(AtlasFormat::new(AtlasOptions { - read_dimensions: Some(vec!["lat".to_string()]), - })); - register_with(&ctx, tmp.path(), "grid", format).await; - - let columns: Vec = ctx - .table_provider("grid") - .await - .unwrap() - .schema() - .fields() - .iter() - .map(|field| field.name().clone()) - .collect(); - assert!( - !columns.contains(&"temperature".to_string()), - "a 2-D array does not fit a 1-D grid: {columns:?}" - ); - } - - // ── the table this crate is registered as ─────────────────────────── - - /// The same collection through `FastObjectTable`, which is what - /// `read_atlas` builds. - /// - /// A collection is one file, and the reader takes it as the marker it is; - /// every other test here goes through `ListingTable`, which would not - /// notice if that stopped being true. - #[tokio::test] - async fn a_collection_reads_through_the_fast_object_table() { - use beacon_datafusion_ext::fast_object::FastObjectTable; - use beacon_datafusion_ext::type_widening::ArrowTypeWidening; - use datafusion::execution::SessionStateBuilder; - - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let state = SessionStateBuilder::new() - .with_config( - SessionConfig::new() - .with_target_partitions(4) - .with_extension(ArrowTypeWidening::default_extension()), - ) - .with_default_features() - .build(); - let ctx = SessionContext::new_with_state(state); - - let directory = tmp.path().to_string_lossy().replace('\\', "/"); - let url = ListingTableUrl::parse(format!("file://{directory}/")).unwrap(); - let table = - FastObjectTable::try_new(&ctx.state(), Arc::new(AtlasFormat::default()), vec![url]) - .await - .expect("a collection registers as a table"); - ctx.register_table("obs", Arc::new(table)).unwrap(); - - assert_eq!(rows(&ctx, "SELECT temperature FROM obs").await, 7); - } - - // ── refusals ──────────────────────────────────────────────────────── - - /// A dataset lives inside a container, not at a path, so no `PARTITIONED - /// BY` value can be read off it. Saying so beats returning the column - /// silently empty. - #[tokio::test] - async fn a_partitioned_table_is_refused_by_name() { - use datafusion::datasource::physical_plan::FileScanConfigBuilder; - use datafusion::execution::object_store::ObjectStoreUrl; - - let table_schema = TableSchema::new( - Arc::new(arrow::datatypes::Schema::empty()), - vec![Arc::new(arrow::datatypes::Field::new( - "year", - DataType::Utf8, - false, - ))], - ); - let source = AtlasSource::new(None, table_schema); - let conf = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), - Arc::new(source) as Arc, - ) - .build(); - - let ctx = SessionContext::new(); - let error = AtlasFormat::default() - .create_physical_plan(&ctx.state(), conf) - .await - .expect_err("a partitioned atlas table is refused") - .to_string(); - assert!(error.contains("Atlas"), "{error}"); - assert!(error.contains("year"), "{error}"); - } - - #[test] - fn an_unparseable_option_is_an_error() { - let error = parse_bool_option("use_reader_cache", "maybe") - .unwrap_err() - .to_string(); - assert!(error.contains("use_reader_cache"), "{error}"); - assert!(error.contains("maybe"), "{error}"); - } - - /// A read that names dimensions stays out of the schema cache, because the - /// key does not carry the dimension set. See the `TODO(#367)` above. - #[test] - fn a_dimension_projected_read_is_not_schema_cached() { - let factory = AtlasFormatFactory::new(Default::default(), Default::default()); - assert!( - factory - .schema_options_fingerprint(&AtlasFormat::default()) - .is_some() - ); - assert!( - factory - .schema_options_fingerprint(&AtlasFormat::new(AtlasOptions { - read_dimensions: Some(vec!["time".to_string()]), - })) - .is_none() - ); - } - // ── a wide collection on two dimensions, through SQL ───────────────── - - /// The shapes every test in this section reads: three datasets whose - /// profile and level counts all differ. - const WIDE: &[(usize, usize)] = &[(3, 4), (5, 4), (2, 6)]; - - /// Rows a full-grid read returns: `profiles * levels`, summed per dataset. - const WIDE_GRID_ROWS: usize = 44; - - /// Rows a read narrowed to `profile` returns: the profiles themselves. - const WIDE_PROFILE_ROWS: usize = 10; - - /// Columns the merged collection carries: eight arrays, four attributes - /// each, and two dataset attributes. - const WIDE_COLUMNS: usize = 42; - - /// Columns that survive a narrowing to `profile`. The four grid arrays go; - /// their attributes stay. - const WIDE_PROFILE_COLUMNS: usize = 38; - - /// A wide collection in a temporary directory, and a table over it. - /// - /// The caller holds the returned directory. It deletes the collection when - /// it drops. - async fn wide_table(ctx: &SessionContext, format: Arc) -> tempfile::TempDir { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - register_with(ctx, tmp.path(), "wide", format).await; - tmp - } - - /// One scalar, as its rendered text. Enough to pin a value without a - /// downcast per Arrow type. - async fn scalar(ctx: &SessionContext, sql: &str) -> String { - let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); - arrow::util::pretty::pretty_format_batches(&batches) - .unwrap() - .to_string() - .lines() - .nth(3) - .expect("a one-row result") - .trim() - .trim_matches('|') - .trim() - .to_string() - } - - /// The row count is the grid every dataset contributes, and not the - /// profile count. - #[tokio::test] - async fn a_wide_collection_reads_every_row() { - let ctx = context(4); - let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; - - assert_eq!( - count(&ctx, "SELECT COUNT(*) FROM wide").await as usize, - WIDE_GRID_ROWS - ); - } - - #[tokio::test] - async fn a_wide_table_carries_every_column() { - let ctx = context(1); - let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; - - let df = ctx.sql("SELECT * FROM wide").await.unwrap(); - assert_eq!(df.schema().fields().len(), WIDE_COLUMNS); - } - - /// A per-profile array beside a per-level one, in one result. - #[tokio::test] - async fn a_wide_row_carries_both_of_its_grids() { - let ctx = context(1); - let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; - - let batches = ctx - .sql( - "SELECT platform, latitude, pressure, temperature, salinity \ - FROM wide WHERE platform = 'set1' AND pressure = 2.0 \ - ORDER BY latitude LIMIT 1", - ) - .await - .unwrap() - .collect() - .await - .unwrap(); - - let rendered = arrow::util::pretty::pretty_format_batches(&batches) - .unwrap() - .to_string(); - // set1 holds temperature = 100 + level, and level 2 is pressure 2. - assert!(rendered.contains("102.0"), "{rendered}"); - assert!(rendered.contains("32.0"), "{rendered}"); - } - - /// An attribute rides along as a constant column. - #[tokio::test] - async fn an_attribute_of_a_wide_collection_reads_as_a_column() { - let ctx = context(1); - let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; - - assert_eq!( - scalar(&ctx, "SELECT DISTINCT \".title\" FROM wide").await, - "wide profiles" - ); - assert_eq!( - scalar(&ctx, "SELECT DISTINCT \"temperature.long_name\" FROM wide").await, - "the temperature" - ); - } - - /// A predicate over a real column, with a known answer. Each dataset owns a - /// disjoint `temperature` range, so this is the pruning arithmetic too. - #[tokio::test] - async fn a_predicate_over_a_wide_collection_selects_the_rows_that_match() { - let ctx = context(4); - let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; - - // Only set2 reaches past 150: its 2 * 6 cells hold 200 to 205. - assert_eq!( - count(&ctx, "SELECT COUNT(*) FROM wide WHERE temperature > 150.0").await, - 12 - ); - assert_eq!( - count(&ctx, "SELECT COUNT(DISTINCT platform) FROM wide").await, - 3 - ); - } - - /// Pruning reads the footer's per-dataset ranges and drops what cannot - /// match. It must not change an answer, whichever way the switch is set. - #[tokio::test] - async fn pruning_does_not_change_the_answers_over_a_wide_collection() { - let queries = [ - "SELECT COUNT(*) FROM wide WHERE temperature > 150.0", - "SELECT COUNT(*) FROM wide WHERE latitude > 12.0", - "SELECT COUNT(*) FROM wide WHERE platform = 'set1'", - "SELECT COUNT(*) FROM wide WHERE salinity < 32.0", - ]; - - let pruned = context(4); - let _a = wide_table(&pruned, Arc::new(AtlasFormat::default())).await; - let whole = context(4); - let _b = wide_table(&whole, format_with_pruning(false)).await; - - for sql in queries { - assert_eq!( - count(&pruned, sql).await, - count(&whole, sql).await, - "pruning changed the answer to: {sql}" - ); - } - } - - /// A projected scan over a pruned collection. - /// - /// The predicate is indexed against the table schema, and the pruning - /// engine reads it against the projected one. `temperature` sits at column - /// 32 of 42, so a five-column projection puts that index out of range. This - /// is the case where the two must be brought into step. - #[tokio::test] - async fn a_projected_scan_prunes_without_losing_its_columns() { - let ctx = context(4); - let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; - - let rows = rows( - &ctx, - "SELECT platform, latitude, longitude, pressure, temperature \ - FROM wide WHERE temperature IS NOT NULL \ - ORDER BY platform, latitude, pressure LIMIT 5", - ) - .await; - assert_eq!(rows, 5); - } - - /// A narrowing to `profile` reads the profiles and not their levels. - /// - /// `COUNT(*)` projects nothing, so the scan picks the widest array of the - /// dataset to count. That array lives on both dimensions, and this - /// narrowing drops it. So the driver has to respect the dimensions too. - #[tokio::test] - async fn narrowing_a_wide_table_to_one_dimension_reads_one_row_per_profile() { - let ctx = context(4); - let _tmp = wide_table( - &ctx, - Arc::new(AtlasFormat::new(AtlasOptions { - read_dimensions: Some(vec!["profile".to_string()]), - })), - ) - .await; - - assert_eq!( - count(&ctx, "SELECT COUNT(*) FROM wide").await as usize, - WIDE_PROFILE_ROWS - ); - // The same number through a column, rather than the count driver. - assert_eq!( - count(&ctx, "SELECT COUNT(latitude) FROM wide").await as usize, - WIDE_PROFILE_ROWS - ); - assert_eq!( - ctx.sql("SELECT * FROM wide") - .await - .unwrap() - .schema() - .fields() - .len(), - WIDE_PROFILE_COLUMNS - ); - } - - /// A narrowed read whose projection names only attributes builds no - /// dimensioned array. The narrowing has nothing to drop and must still - /// succeed. - #[tokio::test] - async fn a_narrowed_read_of_attributes_alone_succeeds() { - let ctx = context(1); - let _tmp = wide_table( - &ctx, - Arc::new(AtlasFormat::new(AtlasOptions { - read_dimensions: Some(vec!["profile".to_string()]), - })), - ) - .await; - - assert_eq!( - scalar(&ctx, "SELECT DISTINCT \"temperature.units\" FROM wide").await, - "1" - ); - assert_eq!( - scalar(&ctx, "SELECT DISTINCT \".institution\" FROM wide").await, - "test" - ); - } - - /// Every dataset is one unit of work, and a partitioned scan divides them - /// without reading one twice. - #[tokio::test] - async fn a_partitioned_scan_of_a_wide_collection_reads_every_row_once() { - for partitions in [1, 2, 8] { - let ctx = context(partitions); - let _tmp = wide_table(&ctx, Arc::new(AtlasFormat::default())).await; - assert_eq!( - count(&ctx, "SELECT COUNT(*) FROM wide").await as usize, - WIDE_GRID_ROWS, - "over {partitions} partitions" - ); - } - } - - // ── STORED AS ATLAS ────────────────────────────────────────────────── - - /// A session that answers `STORED AS ATLAS`, as the runtime builds one. - /// - /// Two lookups sit behind that clause, under two spellings of one name. - /// DataFusion resolves the `STORED AS` word in `table_factories`, upper - /// cased, and Beacon registers [`ListingTableFactoryExt`] there. That - /// factory then resolves the *file format* by the same word lower cased, - /// which is where [`ATLAS_FORMAT`] answers. - fn ddl_context(partitions: usize) -> SessionContext { - use beacon_datafusion_ext::listing_table_factory_ext::ListingTableFactoryExt; - use datafusion::execution::session_state::SessionStateBuilder; - - let mut config = SessionConfig::new() - .with_target_partitions(partitions) - .with_extension(Arc::new(ListingFactory::dynamic())) - .with_extension(Arc::new(ListingTableFactoryExt)); - // What the runtime sets. DataFusion's default matches a glob against - // the file name alone, and a collection is always one directory down, - // so `**/data.atlas` would list nothing and the table would be empty. - config - .options_mut() - .execution - .listing_table_ignore_subdirectory = false; - let state = SessionStateBuilder::new() - .with_config(config) - .with_default_features() - .with_table_factory( - ATLAS_FORMAT.to_uppercase(), - Arc::new(ListingTableFactoryExt), - ) - .build(); - let ctx = SessionContext::new_with_state(state); - ctx.state_ref() - .write() - .register_file_format( - Arc::new(AtlasFormatFactory::new( - Default::default(), - Default::default(), - )), - true, - ) - .expect("the atlas format registers under its own name"); - ctx - } - - /// A directory as SQL takes it: forward slashes, whatever the platform. - fn location(dir: &Path) -> String { - dir.to_string_lossy() - .replace(std::path::MAIN_SEPARATOR, "/") - } - - /// Run one DDL statement, and name the statement if it is refused. - async fn run_ddl(ctx: &SessionContext, sql: &str) { - ctx.sql(sql) - .await - .unwrap_or_else(|e| panic!("refused `{sql}`: {e}")) - .collect() - .await - .unwrap(); - } - - /// The clause a user writes, over the directory that holds the container. - #[tokio::test] - async fn stored_as_atlas_reads_a_collection_directory() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - let ctx = ddl_context(4); - - run_ddl( - &ctx, - &format!( - "CREATE EXTERNAL TABLE t STORED AS ATLAS LOCATION '{}/'", - location(tmp.path()) - ), - ) - .await; - - assert_eq!( - count(&ctx, "SELECT COUNT(*) FROM t").await as usize, - WIDE_GRID_ROWS - ); - assert_eq!( - ctx.sql("SELECT * FROM t") - .await - .unwrap() - .schema() - .fields() - .len(), - WIDE_COLUMNS - ); - } - - /// The container object names the collection just as well as its directory - /// does. That is what a `LOCATION` copied out of a listing looks like. - #[tokio::test] - async fn stored_as_atlas_reads_the_container_object() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - let ctx = ddl_context(4); - - run_ddl( - &ctx, - &format!( - "CREATE EXTERNAL TABLE t STORED AS ATLAS LOCATION '{}/{ATLAS_MARKER}'", - location(tmp.path()) - ), - ) - .await; - - assert_eq!( - count(&ctx, "SELECT COUNT(*) FROM t").await as usize, - WIDE_GRID_ROWS - ); - } - - /// SQL folds the `STORED AS` word, so either spelling reaches the format. - #[tokio::test] - async fn stored_as_atlas_is_case_insensitive() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - - for spelling in ["ATLAS", "atlas", "Atlas"] { - let ctx = ddl_context(1); - run_ddl( - &ctx, - &format!( - "CREATE EXTERNAL TABLE t STORED AS {spelling} LOCATION '{}/'", - location(tmp.path()) - ), - ) - .await; - - assert_eq!( - count(&ctx, "SELECT COUNT(*) FROM t").await as usize, - WIDE_GRID_ROWS, - "STORED AS {spelling}" - ); - } - } - - /// `OPTIONS` reaches the format, so a table can name its dimensions - /// without the `read_atlas` function. - #[tokio::test] - async fn stored_as_atlas_takes_its_options() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - let ctx = ddl_context(4); - - run_ddl( - &ctx, - &format!( - "CREATE EXTERNAL TABLE t STORED AS ATLAS \ - OPTIONS ('read_dimensions' 'profile') LOCATION '{}/'", - location(tmp.path()) - ), - ) - .await; - - assert_eq!( - count(&ctx, "SELECT COUNT(*) FROM t").await as usize, - WIDE_PROFILE_ROWS - ); - } - - /// An `OPTIONS` value is a string, so a list of dimensions is one comma - /// separated string. `read_atlas` takes a real SQL list and joins it into - /// the same option, so both spellings reach one place. - /// - /// Both dimensions of this collection reads it whole, which is what the - /// default already does. The point is the parse, not the answer. - #[tokio::test] - async fn stored_as_atlas_takes_a_comma_separated_dimension_list() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - - // A space after a comma, and a trailing comma, are both tolerated. - for list in ["profile,level", "profile, level", "level,profile,"] { - let ctx = ddl_context(4); - run_ddl( - &ctx, - &format!( - "CREATE EXTERNAL TABLE t STORED AS ATLAS \ - OPTIONS ('read_dimensions' '{list}') LOCATION '{}/'", - location(tmp.path()) - ), - ) - .await; - - assert_eq!( - count(&ctx, "SELECT COUNT(*) FROM t").await as usize, - WIDE_GRID_ROWS, - "'{list}' did not name the whole grid" - ); - assert_eq!( - ctx.sql("SELECT * FROM t") - .await - .unwrap() - .schema() - .fields() - .len(), - WIDE_COLUMNS, - "'{list}'" - ); - } - } - - /// A glob puts several collections in one table, and the rows are their - /// union. - /// - /// This is the form the docs give for a data lake: - /// `LOCATION 'collections/**/data.atlas'`. It rests on - /// `listing_table_ignore_subdirectory` being off, because a collection's - /// container always sits one directory below the glob's prefix. - #[tokio::test] - async fn stored_as_atlas_globs_several_collections_into_one_table() { - // Two collections under one root, so the union is a number neither one - // could produce alone. - let tmp = tempfile::tempdir().unwrap(); - for name in ["one", "two"] { - test_support::wide_profiles(&tmp.path().join(name), WIDE).await; - } - let root = location(tmp.path()); - - for glob in ["**/data.atlas", "*/data.atlas"] { - let ctx = ddl_context(4); - run_ddl( - &ctx, - &format!("CREATE EXTERNAL TABLE t STORED AS ATLAS LOCATION '{root}/{glob}'"), - ) - .await; - - assert_eq!( - count(&ctx, "SELECT COUNT(*) FROM t").await as usize, - 2 * WIDE_GRID_ROWS, - "'{glob}' did not read both collections" - ); - // One schema over both, merged under the session's widening rule. - assert_eq!( - ctx.sql("SELECT * FROM t") - .await - .unwrap() - .schema() - .fields() - .len(), - WIDE_COLUMNS, - "'{glob}'" - ); - } - } - - /// A bad option is an error at `CREATE EXTERNAL TABLE`, not at the first - /// query against the table. - #[tokio::test] - async fn stored_as_atlas_refuses_a_bad_option_at_ddl() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - let ctx = ddl_context(1); - - let error = ctx - .sql(&format!( - "CREATE EXTERNAL TABLE t STORED AS ATLAS \ - OPTIONS ('use_pruning' 'maybe') LOCATION '{}/'", - location(tmp.path()) - )) - .await - .expect_err("'maybe' is no boolean") - .to_string(); - - assert!(error.contains("use_pruning"), "{error}"); - assert!(error.contains("maybe"), "{error}"); - } - /// A session whose merge keeps the first type of a conflicting column. - /// - /// The default refuses such a column outright, so a test of the cast has to - /// ask for `KeepFirst`. - fn keep_first_context(partitions: usize) -> SessionContext { - use beacon_datafusion_ext::type_widening::{ArrowTypeWidening, DefaultArrowTypeWidening}; - - SessionContext::new_with_config( - SessionConfig::new() - .with_target_partitions(partitions) - .with_extension(Arc::new(ArrowTypeWidening::new(Arc::new( - DefaultArrowTypeWidening::keeping_first_type(), - )))), - ) - } - - /// A value the merged type cannot hold reads as null, and does not fail the - /// query. - /// - /// `KeepFirst` keeps `Float64` and marks the column. The mark has to - /// survive the nd encoding, because the scan's target schema is the encoded - /// one. Without it `scan_adapt` casts strictly and `'0.-90'` fails the whole - /// scan. - /// - /// `'3.5'` casts cleanly beside it, so this separates "one value is null" - /// from "the column gave up". - #[tokio::test] - async fn a_value_the_merged_type_cannot_hold_reads_as_null() { - use arrow::array::{Array, Float64Array}; - - let tmp = tempfile::tempdir().unwrap(); - test_support::conflicting_numbers(tmp.path()).await; - let ctx = keep_first_context(1); - register(&ctx, tmp.path(), "t").await; - - let batches = ctx - .sql("SELECT value FROM t") - .await - .unwrap() - .collect() - .await - .expect("an unparseable value is null, not an error"); - - let mut seen: Vec> = Vec::new(); - for batch in &batches { - let column = batch - .column(0) - .as_any() - .downcast_ref::() - .expect("the merge kept the first type"); - for row in 0..column.len() { - seen.push((!column.is_null(row)).then(|| column.value(row))); - } - } - seen.sort_by(|a, b| a.partial_cmp(b).unwrap()); - assert_eq!(seen, vec![None, Some(1.5), Some(2.5), Some(3.5)]); - } - - /// The mark itself, on the schema the scan carries. - /// - /// The logical schema is marked, and the encoded schema has to stay marked. - /// This is the step that used to drop it. - #[test] - fn the_nd_encoding_keeps_the_type_conflict_mark() { - use beacon_datafusion_ext::nd::encoded_schema; - use beacon_datafusion_ext::type_widening::{ - TYPE_CONFLICT_FIRST_TYPE, TYPE_CONFLICT_KEY, is_type_conflict, - }; - - let marked = arrow::datatypes::Field::new("value", DataType::Float64, true).with_metadata( - std::collections::HashMap::from([( - TYPE_CONFLICT_KEY.to_string(), - TYPE_CONFLICT_FIRST_TYPE.to_string(), - )]), - ); - let plain = arrow::datatypes::Field::new("other", DataType::Float64, true); - let logical = Schema::new(vec![marked, plain]); - - let encoded = encoded_schema(&logical); - assert!( - is_type_conflict(encoded.field_with_name("value").unwrap()), - "the encoded field lost the mark" - ); - assert!( - !is_type_conflict(encoded.field_with_name("other").unwrap()), - "an unmarked column must not gain the mark" - ); - // The extension tag still identifies the column as nd-encoded. - assert!(beacon_datafusion_ext::nd::is_nd_encoded( - encoded.field_with_name("value").unwrap() - )); - } - /// What the *default* widening does with Float64 beside Utf8. - #[tokio::test] - async fn probe_default_widening() { - let tmp = tempfile::tempdir().unwrap(); - test_support::conflicting_numbers(tmp.path()).await; - - // Default session: no widening extension registered at all. - let ctx = context(1); - let dir = tmp - .path() - .to_string_lossy() - .replace(std::path::MAIN_SEPARATOR, "/"); - let url = ListingTableUrl::parse(format!("file://{dir}/")).unwrap(); - let format: Arc = Arc::new(AtlasFormat::default()); - let listing = ListingOptions::new(format).with_file_extension(ATLAS_MARKER); - match ListingTableConfig::new(url) - .with_listing_options(listing) - .infer_schema(&ctx.state()) - .await - { - Ok(config) => { - let schema = config.file_schema.clone().unwrap(); - let field = schema.field_with_name("value").unwrap(); - println!("PROBE default infer OK type={:?}", field.data_type()); - println!( - "PROBE default marked={}", - beacon_datafusion_ext::type_widening::is_type_conflict(field) - ); - ctx.register_table("t", Arc::new(ListingTable::try_new(config).unwrap())) - .unwrap(); - match ctx - .sql("SELECT value FROM t") - .await - .unwrap() - .collect() - .await - { - Ok(b) => println!( - "PROBE default query OK\n{}", - arrow::util::pretty::pretty_format_batches(&b).unwrap() - ), - Err(e) => println!("PROBE default query ERR {e}"), - } - } - Err(e) => println!("PROBE default infer ERR {e}"), - } - } - - /// The same value, with no mark on the column at all. - /// - /// An nd cast lands on the `values` list inside the `beacon.nd` struct, and - /// it reads leniently whether or not the merge marked the column. A table - /// whose schema says `Float64` therefore survives a dataset that stores the - /// array as text, however the schema came to say so. - #[tokio::test] - async fn an_unparseable_nd_value_reads_as_null_without_a_mark() { - use arrow::array::{Array, Float64Array}; - - let tmp = tempfile::tempdir().unwrap(); - test_support::value_typed(&tmp.path().join("one"), false).await; - test_support::value_typed(&tmp.path().join("two"), true).await; - - // Two paths, so the table schema comes from the first collection alone - // and carries no mark. That is the shape a scan cannot rely on one. - let ctx = keep_first_context(4); - let urls: Vec = ["one", "two"] - .iter() - .map(|name| { - let dir = tmp - .path() - .join(name) - .to_string_lossy() - .replace(std::path::MAIN_SEPARATOR, "/"); - ListingTableUrl::parse(format!("file://{dir}/")).unwrap() - }) - .collect(); - let format: Arc = Arc::new(AtlasFormat::default()); - let listing = ListingOptions::new(format).with_file_extension(ATLAS_MARKER); - let config = ListingTableConfig::new_with_multi_paths(urls) - .with_listing_options(listing) - .infer_schema(&ctx.state()) - .await - .expect("the collections type"); - let schema = config.file_schema.clone().unwrap(); - let field = schema.field_with_name("value").unwrap(); - assert_eq!(field.data_type(), &DataType::Float64); - assert!( - !beacon_datafusion_ext::type_widening::is_type_conflict(field), - "this shape is the one with no mark; the test is pointless with one" - ); - - ctx.register_table("t", Arc::new(ListingTable::try_new(config).unwrap())) - .unwrap(); - let batches = ctx - .sql("SELECT value FROM t") - .await - .unwrap() - .collect() - .await - .expect("an unparseable value is null, not an error"); - - let mut seen: Vec> = Vec::new(); - for batch in &batches { - let column = batch - .column(0) - .as_any() - .downcast_ref::() - .expect("the table type"); - for row in 0..column.len() { - seen.push((!column.is_null(row)).then(|| column.value(row))); - } - } - seen.sort_by(|a, b| a.partial_cmp(b).unwrap()); - assert_eq!(seen, vec![None, Some(1.5), Some(2.5), Some(3.5)]); - } -} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs new file mode 100644 index 00000000..9337af88 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs @@ -0,0 +1,48 @@ +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use beacon_nd_array::arrow::metrics::ReadMetrics; +use datafusion::{ + datasource::{ + listing::PartitionedFile, + physical_plan::{FileOpenFuture, FileOpener}, + }, + physical_plan::PhysicalExpr, +}; +use futures::FutureExt; +use object_store::ObjectStore; + +use crate::{datafusion::metrics::AtlasScanMetrics, store::AtlasReaderCache}; + +/// One partition's opener: a collection in, its batches out. +/// +/// Every field is a handle or a clone, so the opener itself is cloned into the +/// stream it returns and outlives the call that made it. +#[derive(Clone)] +pub struct AtlasOpener { + pub object_store: Arc, + pub cache: AtlasReaderCache, + /// The scan's output schema, nd-encoded. Its field *names* are the columns + /// to keep, and the encoding leaves names alone. + pub projected_schema: SchemaRef, + /// The same schema with the encoding unwrapped, which is what a predicate + /// and the pruning engine are written against. + pub logical_schema: SchemaRef, + pub read_dimensions: Option>, + pub batch_size: usize, + pub predicate: Option>, + pub read_metrics: ReadMetrics, + pub scan_metrics: AtlasScanMetrics, +} + +impl FileOpener for AtlasOpener { + fn open(&self, file: PartitionedFile) -> datafusion::error::Result { + let fut = async move { + return Err(datafusion::error::DataFusionError::NotImplemented( + "AtlasOpener::open is not implemented yet".to_string(), + )); + }; + + Ok(fut.boxed()) + } +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs index 38b8a3c6..f73b3010 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs @@ -51,11 +51,10 @@ use datafusion::{ use futures::{FutureExt, StreamExt, TryStreamExt}; use object_store::ObjectStore; -use crate::compat; use crate::datafusion::metrics::AtlasScanMetrics; use crate::datafusion::pruning::{CandidateFilter, candidate_filter, logical_schema}; -use crate::reader::{dataset_from_view, project_read_dimensions}; use crate::store::{AtlasReaderCache, get_or_open_atlas}; +use crate::{compat, datafusion::opener::AtlasOpener}; /// DataFusion [`FileSource`] for Atlas collections. #[derive(Debug, Clone)] @@ -67,13 +66,15 @@ pub struct AtlasSource { read_dimensions: Option>, projection: Option, /// The reader cache to consult, or `None` to open every collection afresh. - cache: Option, - /// Whether a predicate scan drops the datasets it can rule out. - use_pruning: bool, + cache: AtlasReaderCache, } impl AtlasSource { - pub fn new(read_dimensions: Option>, table_schema: TableSchema) -> Self { + pub fn new( + read_dimensions: Option>, + table_schema: TableSchema, + cache: AtlasReaderCache, + ) -> Self { Self { table_schema, execution_plan_metrics: ExecutionPlanMetricsSet::new(), @@ -81,23 +82,16 @@ impl AtlasSource { predicate: None, read_dimensions, projection: None, - cache: None, - use_pruning: false, + cache, } } - /// Consult `cache` for opened collections, or open them afresh with `None`. - pub fn with_cache(mut self, cache: Option) -> Self { + /// Consult `cache` for opened collections, or open them afresh. + pub fn with_cache(mut self, cache: AtlasReaderCache) -> Self { self.cache = cache; self } - /// Drop the datasets a predicate rules out, or read them all. - pub fn with_pruning(mut self, use_pruning: bool) -> Self { - self.use_pruning = use_pruning; - self - } - /// Carry a projection the scan pushed down. /// /// The format rebuilds the source in `create_physical_plan`, and without @@ -123,7 +117,6 @@ impl FileSource for AtlasSource { // scan carries them in. logical_schema: logical_schema(&projected_schema), projected_schema, - use_pruning: self.use_pruning, read_dimensions: self.read_dimensions.clone(), batch_size: self.batch_size, predicate: self.predicate.clone(), @@ -208,539 +201,3 @@ impl FileSource for AtlasSource { .with_updated_node(Arc::new(source))) } } - -// ─── The opener ────────────────────────────────────────────────────────────── - -/// One partition's opener: a collection in, its batches out. -/// -/// Every field is a handle or a clone, so the opener itself is cloned into the -/// stream it returns and outlives the call that made it. -#[derive(Clone)] -struct AtlasOpener { - object_store: Arc, - cache: Option, - /// The scan's output schema, nd-encoded. Its field *names* are the columns - /// to keep, and the encoding leaves names alone. - projected_schema: SchemaRef, - /// The same schema with the encoding unwrapped, which is what a predicate - /// and the pruning engine are written against. - logical_schema: SchemaRef, - use_pruning: bool, - read_dimensions: Option>, - batch_size: usize, - predicate: Option>, - read_metrics: ReadMetrics, - scan_metrics: AtlasScanMetrics, -} - -impl FileOpener for AtlasOpener { - fn open(&self, file: PartitionedFile) -> Result { - let opener = self.clone(); - Ok(async move { - let collection = file.object_meta.location.to_string(); - - let open_start = Instant::now(); - let atlas = get_or_open_atlas( - opener.cache.as_ref(), - Arc::clone(&opener.object_store), - &file.object_meta, - ) - .await - .map_err(|e| DataFusionError::Execution(format!("{e}")))?; - opener.scan_metrics.open_time.add_elapsed(open_start); - - let listed = atlas.list_datasets(); - let total = listed.len(); - let names = opener.survivors(&atlas, listed).await; - tracing::debug!( - collection = %collection, - datasets = total, - kept = names.len(), - "atlas scan: collection opened and pruned" - ); - - // One stream over the kept datasets. Each is built when the stream - // reaches it, and its batches drain before the next is touched. - let metrics = opener.read_metrics.clone(); - let batches = futures::stream::iter(names) - .then(move |name| { - let opener = opener.clone(); - let atlas = Arc::clone(&atlas); - let collection = collection.clone(); - async move { opener.dataset_read(&atlas, &collection, &name).await } - }) - .map_ok(move |read| read.stream(Some(metrics.clone()))) - .try_flatten(); - - Ok(batches.boxed()) - } - .boxed()) - } -} - -impl AtlasOpener { - /// The datasets of `atlas` this scan reads, in listing order. - /// - /// One pruning pass over the whole collection decides it, and a dataset - /// ruled out never comes back: no handle, no build, no read. - async fn survivors(&self, atlas: &Arc, listed: Vec) -> Vec { - let filter = self.candidates(atlas).await; - - let mut kept = Vec::with_capacity(listed.len()); - let mut pruned = 0usize; - for (position, name) in listed.into_iter().enumerate() { - if filter.keeps(position, &name) { - kept.push(name); - } else { - pruned += 1; - } - } - self.scan_metrics.datasets_pruned.add(pruned); - kept - } - - /// Which datasets of `atlas` this scan's predicate can still match. - /// - /// Without a predicate, or with pruning off, nothing is ruled out and no - /// index is built. - async fn candidates(&self, atlas: &Arc) -> CandidateFilter { - let Some(predicate) = self.predicate.clone().filter(|_| self.use_pruning) else { - return CandidateFilter::KeepAll; - }; - // DataFusion offers a scan its filters even when it has none, and an - // empty conjunction is the literal `true`. Such a predicate names no - // column, so it can rule nothing out and is not worth a pass. - if collect_columns(&predicate).is_empty() { - return CandidateFilter::KeepAll; - } - - // The predicate is indexed against the table schema, the pruning - // engine reads it against the projected one. A projection pushed down - // after the filter leaves the two out of step, so the columns are - // re-indexed by name here. - let Ok(predicate) = reassign_expr_columns(predicate, &self.logical_schema) else { - return CandidateFilter::KeepAll; - }; - - let started = Instant::now(); - let filter = candidate_filter(atlas, &predicate, &self.logical_schema).await; - // Only an index that exists is a build. Pruning that did not apply - // read nothing and judged nothing. - if filter.is_index() { - self.scan_metrics.index_builds.add(1); - self.scan_metrics.index_rows.add(filter.rows()); - } - self.scan_metrics.prune_time.add_elapsed(started); - filter - } - - /// One dataset, built and planned as a [`FileRead`]. - async fn dataset_read( - &self, - atlas: &Arc, - collection: &str, - name: &str, - ) -> Result> { - let build_start = Instant::now(); - let view = Arc::new(atlas.dataset(name).map_err(|e| { - DataFusionError::Execution(format!( - "Failed to open atlas dataset '{name}' of '{collection}': {e}" - )) - })?); - - let projected = self.projected_names(&view).await?; - let dataset = dataset_from_view(view, projected.as_deref()) - .await - .map_err(|e| DataFusionError::Execution(format!("{e}")))?; - // Explicit dimensions, or a broadcast-compatible default. No log label: - // this runs per dataset, and schema inference already logged the choice. - let dataset = project_read_dimensions(dataset, self.read_dimensions.clone(), None) - .map_err(|e| DataFusionError::Execution(format!("{e}")))?; - - let read = FileRead::plan( - dataset, - Arc::clone(&self.projected_schema), - self.batch_size, - self.predicate.clone(), - // A dataset lives inside a container, not at a path, so no - // `PARTITIONED BY` value can be read off it. The format refuses - // such a table outright. - FilePartitions::none(), - Some(&self.read_metrics), - ) - .await?; - - self.scan_metrics - .dataset_build_time - .add_elapsed(build_start); - self.scan_metrics.datasets_scanned.add(1); - Ok(read) - } - - /// The columns to build for one dataset, or `None` to build every one. - /// - /// The projection reaches the build, so an unprojected array gets no - /// backend and an unprojected attribute is never read out of the footer. - /// - /// Two cases need care. A `COUNT(*)` projects nothing, and building nothing - /// would leave the read with no grid to count; it takes the widest array of - /// the dataset instead, which is what states the row count. And a predicate - /// column is added to the set: the filter above the scan forces such a - /// column into the projection today, but the chunk pruning inside the read - /// matches columns by name and would silently stop pruning if that ever - /// changed. - async fn projected_names(&self, view: &DatasetView) -> Result>> { - if self.projected_schema.fields().is_empty() { - return Ok(count_driver(view, self.read_dimensions.as_deref()) - .await? - .map(|driver| vec![driver])); - } - - let mut names: Vec = self - .projected_schema - .fields() - .iter() - .map(|field| field.name().clone()) - .collect(); - - if let Some(predicate) = &self.predicate { - for column in collect_columns(predicate) { - if !names.iter().any(|name| name == column.name()) { - names.push(column.name().to_string()); - } - } - } - Ok(Some(names)) - } -} - -/// The array a `COUNT(*)` reads to establish a dataset's row count: the widest -/// one Beacon can read. -/// -/// The footer names the candidates, and their sizes come from the segments -/// those arrays live in — one open each for the whole collection, however many -/// datasets a `COUNT(*)` walks. No array data is read, and the driver gets the -/// only backend the dataset builds. -/// -/// `None` for a dataset with no readable array, and the caller then builds what -/// there is — an attribute-only dataset contributes the one row its scalars -/// define. -/// -/// `read_dimensions` rules out the arrays the narrowing after the build would -/// drop. The widest array of a dataset is the one on the most dimensions, and -/// a query that asked for fewer would lose exactly that one and leave the read -/// with no grid at all. -async fn count_driver( - view: &DatasetView, - read_dimensions: Option<&[String]>, -) -> Result> { - let readable: Vec = view - .schema() - .iter() - .filter(|meta| compat::array_dtype_to_nd(meta.dtype()).is_some()) - .map(|meta| meta.name().to_string()) - .collect(); - - let mut widest: Option<(String, usize)> = None; - for array in readable { - let layout = view.array_layout(&array).await.map_err(|e| { - DataFusionError::Execution(format!( - "Failed to read the layout of atlas array '{array}' of dataset '{}': {e}", - view.name() - )) - })?; - if let Some(wanted) = read_dimensions - && !layout - .dimension_names() - .iter() - .all(|dim| wanted.iter().any(|kept| kept == dim)) - { - continue; - } - let cells = layout.element_count(); - if widest.as_ref().is_none_or(|(_, held)| cells > *held) { - widest = Some((array, cells)); - } - } - Ok(widest.map(|(array, _)| array)) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::test_support; - use datafusion::physical_plan::metrics::MetricsSet; - - fn source() -> AtlasSource { - AtlasSource::new( - None, - TableSchema::from_file_schema(Arc::new(arrow::datatypes::Schema::empty())), - ) - } - - /// A container is one unit of work, so DataFusion may not split it into - /// byte ranges the way it would a Parquet file. - #[test] - fn a_container_is_never_split() { - assert!(!source().supports_repartitioning()); - } - - // ── opening a collection ──────────────────────────────────────────── - - /// An opener over the collection in `dir`, projecting every column, with - /// the metrics it reports. - async fn opener( - dir: &std::path::Path, - predicate: Option>, - use_pruning: bool, - ) -> (AtlasOpener, ExecutionPlanMetricsSet) { - use crate::reader::collection_schema; - use beacon_datafusion_ext::type_widening::ArrowTypeWidening; - - let atlas = test_support::open(dir).await; - let logical = collection_schema(&atlas, None, "c", &ArrowTypeWidening::default_extension()) - .await - .unwrap(); - let projected_schema: SchemaRef = - Arc::new(beacon_datafusion_ext::nd::encoded_schema(&logical)); - - let metrics = ExecutionPlanMetricsSet::new(); - let (store, _) = test_support::store_and_marker(dir); - let opener = AtlasOpener { - object_store: store, - cache: None, - logical_schema: logical_schema(&projected_schema), - projected_schema, - use_pruning, - read_dimensions: None, - batch_size: usize::MAX, - predicate, - read_metrics: ReadMetrics::new(&metrics, 0), - scan_metrics: AtlasScanMetrics::new(&metrics, 0), - }; - (opener, metrics) - } - - fn count_of(metrics: &ExecutionPlanMetricsSet, name: &str) -> usize { - let set: MetricsSet = metrics.clone_inner(); - set.sum_by_name(name).map_or(0, |value| value.as_usize()) - } - - /// One open, then every dataset of the collection, in listing order. - #[tokio::test] - async fn an_open_streams_every_dataset_of_the_collection() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 6).await; - let (opener, metrics) = opener(tmp.path(), None, true).await; - let (_, marker) = test_support::store_and_marker(tmp.path()); - - let stream = opener - .open(PartitionedFile::from(marker)) - .unwrap() - .await - .expect("the collection opens"); - let batches: Vec<_> = stream.try_collect().await.expect("every dataset reads"); - - // Six datasets, one chunk each, one nd batch per chunk. - assert_eq!(batches.len(), 6); - assert_eq!(count_of(&metrics, "atlas_datasets_scanned"), 6); - assert_eq!(count_of(&metrics, "atlas_datasets_pruned"), 0); - assert_eq!( - count_of(&metrics, "atlas_index_builds"), - 0, - "no predicate, no index" - ); - } - - /// A predicate prunes once for the collection, before any dataset is - /// built. A ruled-out dataset then costs no build and no read. - #[tokio::test] - async fn an_open_prunes_the_collection_before_it_reads() { - use datafusion::logical_expr::Operator; - use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; - use datafusion::scalar::ScalarValue; - - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 6).await; - // Datasets hold 10i..=10i+3, so only d5 (50..=53) reaches past 45. - let predicate: Arc = Arc::new(BinaryExpr::new( - Arc::new(Column::new("temperature", 0)), - Operator::Gt, - Arc::new(Literal::new(ScalarValue::Float32(Some(45.0)))), - )); - let (opener, metrics) = opener(tmp.path(), Some(predicate), true).await; - let (_, marker) = test_support::store_and_marker(tmp.path()); - - let stream = opener - .open(PartitionedFile::from(marker)) - .unwrap() - .await - .expect("the collection opens"); - let batches: Vec<_> = stream.try_collect().await.unwrap(); - - assert_eq!(batches.len(), 1, "one dataset survives"); - assert_eq!(count_of(&metrics, "atlas_datasets_scanned"), 1); - assert_eq!(count_of(&metrics, "atlas_datasets_pruned"), 5); - assert_eq!(count_of(&metrics, "atlas_index_builds"), 1); - assert_eq!(count_of(&metrics, "atlas_index_rows"), 6); - } - - /// With pruning off, every dataset is built and read. - /// - /// The predicate still reaches [`FileRead::plan`], where the chunk mask over - /// a 1-D coordinate skips chunks no row of which can match. That layer is - /// always on and is not what the switch controls, so this pins the dataset - /// metrics and not the batch count. - #[tokio::test] - async fn pruning_off_reads_every_dataset() { - use datafusion::logical_expr::Operator; - use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; - use datafusion::scalar::ScalarValue; - - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 6).await; - let predicate: Arc = Arc::new(BinaryExpr::new( - Arc::new(Column::new("temperature", 0)), - Operator::Gt, - Arc::new(Literal::new(ScalarValue::Float32(Some(45.0)))), - )); - let (opener, metrics) = opener(tmp.path(), Some(predicate), false).await; - let (_, marker) = test_support::store_and_marker(tmp.path()); - - let stream = opener - .open(PartitionedFile::from(marker)) - .unwrap() - .await - .unwrap(); - let _batches: Vec<_> = stream.try_collect().await.unwrap(); - - assert_eq!( - count_of(&metrics, "atlas_datasets_scanned"), - 6, - "every dataset is built" - ); - assert_eq!(count_of(&metrics, "atlas_datasets_pruned"), 0); - assert_eq!( - count_of(&metrics, "atlas_index_builds"), - 0, - "no index is built" - ); - } - - /// A path that is not a container fails at the open, and the error names - /// what a collection is called. - #[tokio::test] - async fn a_path_that_is_not_a_container_fails_to_open() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - let (opener, _) = opener(tmp.path(), None, false).await; - - let error = opener - .open(PartitionedFile::new("obs/index.json", 1)) - .unwrap() - .await - .err() - .expect("only the container names a collection") - .to_string(); - assert!(error.contains("data.atlas"), "{error}"); - } - - // ── which columns a dataset is built with ─────────────────────────── - - /// One dataset's view. It owns what it needs, so the collection handle - /// behind it may go. - async fn view(dir: &std::path::Path, dataset: &str) -> DatasetView { - test_support::open(dir) - .await - .dataset(dataset) - .expect("the dataset") - } - - fn opener_wanting( - projected: Vec<&str>, - predicate: Option>, - ) -> AtlasOpener { - let fields: Vec = projected - .into_iter() - .map(|name| arrow::datatypes::Field::new(name, arrow::datatypes::DataType::Null, true)) - .collect(); - let projected_schema = Arc::new(arrow::datatypes::Schema::new(fields)); - AtlasOpener { - object_store: Arc::new(object_store::memory::InMemory::new()), - cache: None, - logical_schema: Arc::clone(&projected_schema), - projected_schema, - use_pruning: false, - read_dimensions: None, - batch_size: usize::MAX, - predicate, - read_metrics: ReadMetrics::new(&ExecutionPlanMetricsSet::new(), 0), - scan_metrics: AtlasScanMetrics::new(&ExecutionPlanMetricsSet::new(), 0), - } - } - - #[tokio::test] - async fn a_scan_builds_the_columns_it_projects() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let opener = opener_wanting(vec!["temperature"], None); - let names = opener - .projected_names(&view(tmp.path(), "winter").await) - .await - .expect("the layouts resolve") - .expect("a projection"); - assert_eq!(names, vec!["temperature".to_string()]); - } - - /// A predicate column joins the set even when the projection leaves it out. - /// The chunk pruning inside the read matches by name, so a missing column - /// would silently stop it pruning. - #[tokio::test] - async fn a_predicate_column_is_built_too() { - use datafusion::logical_expr::Operator; - use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; - use datafusion::scalar::ScalarValue; - - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let predicate: Arc = Arc::new(BinaryExpr::new( - Arc::new(Column::new("cycle", 0)), - Operator::Gt, - Arc::new(Literal::new(ScalarValue::Int32(Some(20)))), - )); - let opener = opener_wanting(vec!["temperature"], Some(predicate)); - let names = opener - .projected_names(&view(tmp.path(), "winter").await) - .await - .expect("the layouts resolve") - .expect("a projection"); - assert_eq!( - names, - vec!["temperature".to_string(), "cycle".to_string()], - "the predicate's column is kept alongside the projection" - ); - } - - /// `COUNT(*)` projects nothing and reads one array: the row count is a - /// property of the grid, and building every column to find it would read - /// every attribute of the dataset for nothing. - #[tokio::test] - async fn a_count_reads_the_widest_array_alone() { - let tmp = tempfile::tempdir().unwrap(); - test_support::chunked_grid(tmp.path()).await; - - let opener = opener_wanting(vec![], None); - let names = opener - .projected_names(&view(tmp.path(), "grid").await) - .await - .expect("the layouts resolve") - .expect("a driver"); - assert_eq!(names.len(), 1); - assert!( - names[0] == "temperature" || names[0] == "sparse", - "either 4x6 array states the row count: {names:?}" - ); - } -} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs index 281fdf5e..e8e6611b 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs @@ -29,11 +29,9 @@ //! # What this crate does with it //! //! [`store`] finds a collection's marker and opens it, through a reader cache. -//! [`reader`] turns one dataset into a Beacon -//! [`AnyDataset`](beacon_nd_array::dataset::AnyDataset) whose columns are lazy -//! [`NdArrayD`](beacon_nd_array::NdArrayD) values backed by [`backend`], and -//! derives the Arrow schema of a whole collection. [`compat`] holds the type -//! and column-name mapping the two share. +//! [`compat`] holds the column-name and type mapping, and derives the Arrow +//! schema of a whole collection from its footer. [`backend`] holds the lazy +//! [`NdArrayD`](beacon_nd_array::NdArrayD) values a scan reads through. //! //! # One collection is one unit of work //! @@ -77,12 +75,9 @@ pub use atlas; pub mod backend; pub mod compat; -pub mod config; pub mod datafusion; -pub mod reader; pub mod store; -pub use config::AtlasConfig; pub use datafusion::{AtlasFormat, AtlasFormatFactory, AtlasOptions, ReadAtlasFunc}; #[cfg(test)] diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs deleted file mode 100644 index 69de28a5..00000000 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs +++ /dev/null @@ -1,833 +0,0 @@ -//! Turning an Atlas collection into what Beacon's engine reads: one dataset as -//! an [`AnyDataset`] of lazy columns, and a whole collection as one Arrow -//! schema. -//! -//! Nothing here reads array data. A dataset is built from the collection -//! footer, which the open already held, and its columns fetch their bytes when -//! the scan asks for them. - -use std::sync::Arc; - -use arrow::datatypes::{Schema, SchemaRef}; -use atlas::{Atlas, Attr, DType, DatasetView}; -use beacon_datafusion_ext::type_widening::{ArrowTypeWidening, LabeledSchema}; -use beacon_nd_array::{ - NdArrayD, - arrow::schema::any_dataset_to_arrow_schema, - dataset::{AnyDataset, Dataset, resolve_read_dimensions}, - projection::DatasetProjection, -}; -use indexmap::IndexMap; -use object_store::{ObjectMeta, ObjectStore}; - -use crate::compat; - -/// Build one dataset as an [`AnyDataset`] of lazy columns. -/// -/// `projected_names` is the column set the query wants: -/// -/// - `None` — every array and attribute. -/// - `Some(names)` — only the columns named. A name the dataset does not hold -/// is ignored, so a projection may name columns from any dataset of the -/// collection. -/// -/// The projection reaches the *build*, not just the result: an array outside it -/// gets no backend, and an attribute outside it is not even read out of the -/// footer. A column-subset query over a wide dataset therefore pays nothing for -/// the columns it did not ask for. -/// -/// A value Beacon cannot surface — a `Bool` or list array, a list attribute — -/// is dropped with a `debug` log. A collection can hold a million datasets, so -/// a louder log would be a flood. Such an array is settled from the footer -/// alone, so its segment is never opened. -/// -/// # What this reads -/// -/// Names and element types come from the footer the open already held. A -/// layout and an attribute value live in the variable's own segment, so the -/// first dataset to want one opens it. A segment covers that variable across -/// the whole collection, so every later dataset reuses the same handle. -pub async fn dataset_from_view( - view: Arc, - projected_names: Option<&[String]>, -) -> anyhow::Result { - let included = |name: &str| projected_names.is_none_or(|names| names.iter().any(|n| n == name)); - // Whether any projected column could be an attribute of this array. It - // saves building an attribute map the projection would throw away. - let wants_attrs_of = |array: &str| { - projected_names.is_none_or(|names| { - names - .iter() - .any(|name| compat::is_attr_column_of(name, array)) - }) - }; - let wants_global_attrs = - projected_names.is_none_or(|names| names.iter().any(|name| name.starts_with('.'))); - - // The schema borrows the footer, so it is resolved before the first await - // rather than held across one. - let declared = declared_arrays(&view); - let mut arrays: IndexMap> = IndexMap::with_capacity(declared.len()); - - for (array_name, dtype) in &declared { - if included(array_name) { - match compat::array_dtype_to_nd(dtype) { - // The layout is in the variable's segment, so it is asked for - // only once the dtype says the column can exist at all. - Some(_) => { - let layout = view.array_layout(array_name).await.map_err(|e| { - anyhow::anyhow!( - "Failed to read the layout of atlas array '{array_name}' \ - of dataset '{}': {e}", - view.name() - ) - })?; - match compat::array_to_nd_array(Arc::clone(&view), array_name, dtype, &layout) { - Ok(nd) => { - arrays.insert(array_name.clone(), nd); - } - Err(e) => tracing::debug!( - dataset = %view.name(), - array = %array_name, - "atlas array left out of the dataset: {e}" - ), - } - } - None => tracing::debug!( - dataset = %view.name(), - array = %array_name, - "atlas array left out of the dataset: {} is no Beacon column", - compat::dtype_tag(dtype) - ), - } - } - - if !wants_attrs_of(array_name) { - continue; - } - for (key, value) in attributes_of(&view, Some(array_name)).await? { - let column = compat::array_attr_column(array_name, &key); - if !included(&column) { - continue; - } - match compat::attribute_to_nd_array(&value) { - Ok(nd) => { - arrays.insert(column, nd); - } - Err(e) => tracing::debug!( - dataset = %view.name(), - column = %column, - "atlas attribute left out of the dataset: {e}" - ), - } - } - } - - if wants_global_attrs { - for (key, value) in attributes_of(&view, None).await? { - let column = compat::global_attr_column(&key); - if !included(&column) { - continue; - } - match compat::attribute_to_nd_array(&value) { - Ok(nd) => { - arrays.insert(column, nd); - } - Err(e) => tracing::debug!( - dataset = %view.name(), - column = %column, - "atlas attribute left out of the dataset: {e}" - ), - } - } - } - - arrays.sort_keys(); - - let dataset = Dataset::new(view.name().to_string(), arrays).await; - AnyDataset::try_from_dataset(dataset) - .await - .map_err(|e| anyhow::anyhow!("Failed to wrap atlas dataset '{}': {e}", view.name())) -} - -/// Every array the dataset declares, as owned name and element type. -/// -/// The schema borrows the collection footer. Resolving it up front keeps that -/// borrow off the async path, where the same view is also cloned into a -/// backend. -fn declared_arrays(view: &DatasetView) -> Vec<(String, DType)> { - view.schema() - .iter() - .map(|meta| (meta.name().to_string(), meta.dtype().clone())) - .collect() -} - -/// The attribute values of one scope: `Some(array)` for an array's own, -/// `None` for the dataset's. -/// -/// Values live in a segment, not in the footer, so this reads one. An array -/// with no attribute costs nothing, because the schema settles it first. -async fn attributes_of( - view: &DatasetView, - array: Option<&str>, -) -> anyhow::Result> { - let scope = match array { - Some(array) => view.array_attributes(array).await, - None => view.attributes().await, - }; - scope.map_err(|e| { - let what = array.map_or_else( - || "the dataset attributes".to_string(), - |array| format!("the attributes of array '{array}'"), - ); - anyhow::anyhow!( - "Failed to read {what} of atlas dataset '{}': {e}", - view.name() - ) - }) -} - -/// Narrow `dataset` to `read_dimensions`, or to a broadcast-compatible default -/// when none are given. -/// -/// Without this a `SELECT *` over a dataset whose arrays live on incompatible -/// dimension sets could not broadcast onto one grid. `log_label` names the -/// caller in the auto-selection log; pass `None` from per-dataset code, where -/// schema inference has already logged the choice. -pub fn project_read_dimensions( - dataset: AnyDataset, - read_dimensions: Option>, - log_label: Option<&str>, -) -> anyhow::Result { - match resolve_read_dimensions(&dataset, read_dimensions, log_label) { - Some(dims) => { - let held: Vec = dims - .into_iter() - .filter(|dim| dataset.dataset().dimensions.contains_key(dim)) - .collect(); - dataset - .project(&DatasetProjection::new_with_dimension_projection(held)) - .map_err(|e| { - anyhow::anyhow!("Failed to project the atlas dataset by dimension: {e}") - }) - } - None => Ok(dataset), - } -} - -/// Build one dataset of the collection at `marker`, over `store`. -/// -/// A convenience for a caller holding an object store. The scan opens the -/// collection once and calls [`dataset_from_view`] per dataset instead. -pub async fn open_dataset( - store: Arc, - marker: &ObjectMeta, - dataset: &str, -) -> anyhow::Result { - let atlas = crate::store::get_or_open_atlas(None, store, marker).await?; - let view = atlas - .dataset(dataset) - .map_err(|e| anyhow::anyhow!("Failed to open atlas dataset '{dataset}': {e}"))?; - dataset_from_view(Arc::new(view), None).await -} - -/// The Arrow schema of a whole collection: every live dataset, merged. -/// -/// Atlas reconciles nothing. Two datasets may declare one array name with two -/// dtypes, so the collection's schema is the widening merge of its datasets -/// schemas, under the rule the session carries. -pub async fn collection_schema( - atlas: &Arc, - read_dimensions: Option<&[String]>, - label: &str, - widening: &ArrowTypeWidening, -) -> anyhow::Result { - let mut schemas: Vec = Vec::new(); - - for name in atlas.list_datasets() { - let view = atlas - .dataset(&name) - .map_err(|e| anyhow::anyhow!("Failed to open atlas dataset '{name}': {e}"))?; - - let dataset = dataset_from_view(Arc::new(view), None).await?; - // The same narrowing the scan applies, so the schema states what a - // query can actually return. - let dataset = - project_read_dimensions(dataset, read_dimensions.map(<[String]>::to_vec), None)?; - let schema = any_dataset_to_arrow_schema(&dataset).map_err(|e| { - anyhow::anyhow!("Failed to derive the Arrow schema of atlas dataset '{name}': {e}") - })?; - // The dataset names the schema, so a refused column names both sides. - schemas.push(LabeledSchema::new( - Arc::new(schema), - format!("{label}#{name}"), - )); - } - - if schemas.is_empty() { - // A collection with no live dataset has no column. That is legal, and - // an empty schema is what every other reader answers with. - return Ok(Arc::new(Schema::empty())); - } - - widening - .merge_schemas(&schemas) - .map_err(|e| anyhow::anyhow!("Failed to merge the schemas of the atlas datasets: {e}")) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::test_support; - use arrow::datatypes::DataType; - use beacon_nd_array::{NdArray, datatypes::TimestampNanosecond}; - - async fn view(dir: &std::path::Path, dataset: &str) -> Arc { - let atlas = test_support::open(dir).await; - Arc::new( - atlas - .dataset(dataset) - .expect("the dataset is in the collection"), - ) - } - - fn widening() -> Arc { - ArrowTypeWidening::default_extension() - } - - fn names(dataset: &AnyDataset) -> Vec { - dataset.fields().keys().cloned().collect() - } - - // ── one dataset ───────────────────────────────────────────────────── - - #[tokio::test] - async fn a_dataset_holds_its_arrays_and_its_attributes() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) - .await - .unwrap(); - - assert_eq!(dataset.name(), "winter"); - assert_eq!( - names(&dataset), - vec![ - ".season", - ".year", - "cycle", - "temperature", - "temperature.units", - "time", - ] - ); - } - - /// A dataset attribute takes a leading dot, and an array attribute takes its - /// array's name. That is what netCDF and Zarr do, and it keeps an attribute - /// from colliding with an array of the same name. - #[tokio::test] - async fn attributes_are_named_the_way_every_nd_format_names_them() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) - .await - .unwrap(); - - let season = dataset - .get_array(".season") - .expect("the dataset attribute is a column") - .as_any() - .downcast_ref::>() - .expect("a string column"); - assert!(season.shape().is_empty(), "an attribute has no axis"); - assert_eq!( - season.clone_into_raw_vec().await, - vec!["winter".to_string()] - ); - - let units = dataset - .get_array("temperature.units") - .expect("the array attribute is a column") - .as_any() - .downcast_ref::>() - .expect("a string column"); - assert_eq!( - units.clone_into_raw_vec().await, - vec!["celsius".to_string()] - ); - } - - #[tokio::test] - async fn every_column_keeps_the_type_the_footer_gave_it() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) - .await - .unwrap(); - let schema = any_dataset_to_arrow_schema(&dataset).unwrap(); - let field = |name: &str| schema.field_with_name(name).unwrap().data_type().clone(); - - assert_eq!(field("temperature"), DataType::Float32); - assert_eq!(field("cycle"), DataType::Int32); - assert_eq!( - field("time"), - DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None) - ); - assert_eq!(field(".year"), DataType::Int64); - assert_eq!(field(".season"), DataType::Utf8); - } - - #[tokio::test] - async fn array_values_read_back() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) - .await - .unwrap(); - - let temperature = dataset - .get_array("temperature") - .unwrap() - .as_any() - .downcast_ref::>() - .unwrap(); - assert_eq!( - temperature.clone_into_raw_vec().await, - vec![1.0, 2.0, 3.0, 4.0] - ); - - let time = dataset - .get_array("time") - .unwrap() - .as_any() - .downcast_ref::>() - .unwrap(); - assert_eq!( - time.clone_into_raw_vec().await[0], - TimestampNanosecond(test_support::EPOCH_NANOS) - ); - } - - /// The fill reaches the column, which is what lets the engine null an - /// unwritten cell. - #[tokio::test] - async fn an_arrays_fill_value_reaches_its_column() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) - .await - .unwrap(); - - let cycle = dataset - .get_array("cycle") - .unwrap() - .as_any() - .downcast_ref::>() - .unwrap(); - assert_eq!(cycle.fill_value().await, Some(-1)); - - let temperature = dataset - .get_array("temperature") - .unwrap() - .as_any() - .downcast_ref::>() - .unwrap(); - assert_eq!(temperature.fill_value().await, None, "none was declared"); - } - - /// The chunk shape is the writer's, so a scan cuts the dataset on the grid - /// the file stores rather than on one it invents. - #[tokio::test] - async fn a_column_reports_the_stored_chunk_shape() { - let tmp = tempfile::tempdir().unwrap(); - test_support::chunked_grid(tmp.path()).await; - - let dataset = dataset_from_view(view(tmp.path(), "grid").await, None) - .await - .unwrap(); - let temperature = dataset.get_array("temperature").unwrap(); - assert_eq!(temperature.shape(), vec![4, 6]); - assert_eq!(temperature.chunk_shape(), vec![2, 3]); - assert_eq!( - temperature.dimensions(), - vec!["lat".to_string(), "lon".to_string()] - ); - } - - // ── projection ────────────────────────────────────────────────────── - - #[tokio::test] - async fn a_projection_builds_only_what_it_names() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let wanted = vec!["temperature".to_string(), ".season".to_string()]; - let dataset = dataset_from_view(view(tmp.path(), "winter").await, Some(&wanted)) - .await - .unwrap(); - - assert_eq!(names(&dataset), vec![".season", "temperature"]); - } - - #[tokio::test] - async fn a_projection_may_name_a_column_this_dataset_lacks() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - // `cycle` is winter's alone; summer simply has none of it. - let wanted = vec!["temperature".to_string(), "cycle".to_string()]; - let dataset = dataset_from_view(view(tmp.path(), "summer").await, Some(&wanted)) - .await - .unwrap(); - assert_eq!(names(&dataset), vec!["temperature"]); - } - - #[tokio::test] - async fn an_empty_projection_builds_nothing() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let dataset = dataset_from_view(view(tmp.path(), "winter").await, Some(&[])) - .await - .unwrap(); - assert!(names(&dataset).is_empty(), "COUNT(*) needs no column here"); - } - - // ── values Beacon cannot surface ──────────────────────────────────── - - #[tokio::test] - async fn a_list_attribute_is_dropped_and_the_rest_survives() { - let tmp = tempfile::tempdir().unwrap(); - test_support::skips(tmp.path()).await; - - let dataset = dataset_from_view(view(tmp.path(), "s").await, None) - .await - .unwrap(); - let columns = names(&dataset); - - assert!( - !columns.contains(&"value.range".to_string()), - "a list attribute has no rank-0 form: {columns:?}" - ); - assert!( - !columns.contains(&".tags".to_string()), - "nor does a list dataset attribute: {columns:?}" - ); - assert_eq!( - columns, - vec![".title", "value", "value.units"], - "everything else is kept" - ); - } - - // ── ragged datasets ───────────────────────────────────────────────── - - /// The `{array}.{attr}` naming is what the engine's ragged detection reads, - /// so a CF contiguous ragged collection is recognized without anything - /// atlas-specific. - #[tokio::test] - async fn a_plain_dataset_is_regular() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let dataset = dataset_from_view(view(tmp.path(), "winter").await, None) - .await - .unwrap(); - assert!(!dataset.is_ragged()); - } - - // ── the collection schema ─────────────────────────────────────────── - - #[tokio::test] - async fn a_collections_schema_is_the_union_of_its_datasets() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - let atlas = test_support::open(tmp.path()).await; - - let schema = collection_schema(&atlas, None, "c", &widening()) - .await - .unwrap(); - let columns: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - - for expected in [ - ".season", - ".year", - "cycle", - "temperature", - "temperature.units", - "time", - ] { - assert!( - columns.contains(&expected), - "missing {expected}: {columns:?}" - ); - } - } - - /// Two datasets that give one array two numeric types merge to the type that - /// holds both, by the rule of the session rather than one of atlas's own. - /// - /// `Int16` beside `Float32` gives `Float64`: a `Float32` holds no `Int32`, - /// so a rule that kept it for a narrow integer would make the answer depend - /// on which dataset the merge saw first. See issue #377. - #[tokio::test] - async fn a_shared_array_widens_to_a_type_that_holds_both() { - let tmp = tempfile::tempdir().unwrap(); - test_support::widening(tmp.path()).await; - let atlas = test_support::open(tmp.path()).await; - - let schema = collection_schema(&atlas, None, "c", &widening()) - .await - .unwrap(); - assert_eq!( - schema.field_with_name("value").unwrap().data_type(), - &DataType::Float64, - "Int16 and Float32 widen to Float64" - ); - assert_eq!( - schema.field_with_name("flag").unwrap().data_type(), - &DataType::Int32, - "a column only one dataset declares keeps its own type" - ); - } - - /// A column two datasets type in two families is refused, and the error - /// names both datasets. - /// - /// Atlas reconciles nothing, so a collection can hold this. Beacon settles - /// it the way it settles two files of any other format, and the label makes - /// the offender findable in a collection of a million datasets. - #[tokio::test] - async fn types_that_do_not_widen_are_refused_by_name() { - let tmp = tempfile::tempdir().unwrap(); - test_support::incompatible(tmp.path()).await; - let atlas = test_support::open(tmp.path()).await; - - let error = collection_schema(&atlas, None, "sensor", &widening()) - .await - .expect_err("Utf8 and Int64 are two families") - .to_string(); - - assert!(error.contains("value"), "the column: {error}"); - assert!( - error.contains("Utf8") && error.contains("Int64"), - "both types: {error}" - ); - assert!( - error.contains("sensor#a") && error.contains("sensor#b"), - "and both datasets: {error}" - ); - } - - /// A deployment that reads such a collection anyway sets `keep_first`, and - /// the column then takes the type of the first dataset in listing order. - #[tokio::test] - async fn keep_first_settles_a_conflict_with_the_first_datasets_type() { - use beacon_datafusion_ext::type_widening::{DefaultArrowTypeWidening, TypeConflict}; - - let tmp = tempfile::tempdir().unwrap(); - test_support::incompatible(tmp.path()).await; - let atlas = test_support::open(tmp.path()).await; - - let keep_first = ArrowTypeWidening::new(Arc::new(DefaultArrowTypeWidening { - on_conflict: TypeConflict::KeepFirst, - })); - let schema = collection_schema(&atlas, None, "c", &keep_first) - .await - .unwrap(); - - // `a` is written first and states `value` as a string. - assert_eq!( - schema.field_with_name("value").unwrap().data_type(), - &DataType::Utf8 - ); - } - - #[tokio::test] - async fn a_collection_with_no_dataset_has_no_column() { - let tmp = tempfile::tempdir().unwrap(); - test_support::empty(tmp.path()).await; - let atlas = test_support::open(tmp.path()).await; - - let schema = collection_schema(&atlas, None, "c", &widening()) - .await - .unwrap(); - assert!(schema.fields().is_empty()); - } - - /// Datasets that declare the same arrays share one interned schema, so the - /// derivation runs once however many of them there are. The fleet fixture - /// gives ten datasets one shape, and they differ only in an attribute - /// value. - #[tokio::test] - async fn a_fleet_of_one_shape_merges_to_one_set_of_columns() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 10).await; - let atlas = test_support::open(tmp.path()).await; - - assert_eq!(atlas.interned_schemas(), 1, "the fixture shares its schema"); - - // Ten datasets, one schema. The merge folds the repeats away. - let schema = collection_schema(&atlas, None, "c", &widening()) - .await - .unwrap(); - let columns: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - assert_eq!(columns, vec![".platform", "temperature"]); - } - - // ── dimensions ────────────────────────────────────────────────────── - - #[tokio::test] - async fn read_dimensions_narrow_the_schema_to_the_grid_they_name() { - let tmp = tempfile::tempdir().unwrap(); - test_support::chunked_grid(tmp.path()).await; - let atlas = test_support::open(tmp.path()).await; - - // Both arrays live on `lat` and `lon`; naming only `lat` leaves neither. - let dims = ["lat".to_string()]; - let schema = collection_schema(&atlas, Some(&dims), "c", &widening()) - .await - .unwrap(); - let columns: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - assert!( - !columns.contains(&"temperature"), - "a 2-D array does not fit a 1-D grid: {columns:?}" - ); - } - // ── a wide collection on two dimensions ───────────────────────────── - - /// The shapes every test in this section reads: three datasets whose - /// profile and level counts all differ. - const WIDE: &[(usize, usize)] = &[(3, 4), (5, 4), (2, 6)]; - - /// A dataset per name the writer added, in write order. - #[tokio::test] - async fn a_wide_collection_lists_every_dataset() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - let atlas = test_support::open(tmp.path()).await; - - assert_eq!(atlas.list_datasets(), vec!["set0", "set1", "set2"]); - } - - /// The footer types every array without a segment read, so this is what one - /// open already knows. - #[tokio::test] - async fn every_dataset_of_a_wide_collection_declares_the_same_arrays() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - let atlas = test_support::open(tmp.path()).await; - - for name in atlas.list_datasets() { - let view = atlas.dataset(&name).expect("the dataset is listed"); - assert_eq!(view.schema().len(), 8, "{name} declares eight arrays"); - } - } - - /// The datasets share one interned schema, so the merge derives once. - #[tokio::test] - async fn a_wide_collection_types_every_column_from_the_footer() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - let atlas = test_support::open(tmp.path()).await; - let schema = collection_schema(&atlas, None, "wide", &widening()) - .await - .unwrap(); - - // Eight arrays, four attributes each, and two dataset attributes. - assert_eq!(schema.fields().len(), 42); - - let column = |name: &str| { - schema - .field_with_name(name) - .unwrap_or_else(|_| panic!("{name} is a column of the collection")) - .data_type() - .clone() - }; - // Atlas has a native timestamp, so `time` arrives as one. - assert_eq!( - column("time"), - DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None) - ); - assert_eq!(column("latitude"), DataType::Float64); - assert_eq!(column("temperature"), DataType::Float32); - assert_eq!(column("platform"), DataType::Utf8); - // An attribute is a column under `{array}.{attr}`, and a dataset - // attribute under `.{attr}`. - assert_eq!(column("temperature.units"), DataType::Utf8); - assert_eq!(column("temperature.valid_max"), DataType::Float64); - assert_eq!(column(".title"), DataType::Utf8); - } - - /// `profile` and `level` are the two dimensions. A list that holds only the - /// first keeps the arrays on it and drops the arrays that need both. - #[tokio::test] - async fn narrowing_to_one_dimension_drops_the_two_dimensional_arrays() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - let atlas = test_support::open(tmp.path()).await; - let dims = ["profile".to_string()]; - let schema = collection_schema(&atlas, Some(&dims), "wide", &widening()) - .await - .unwrap(); - - // The four grid arrays are gone; their attributes are scalars and stay. - assert_eq!(schema.fields().len(), 38); - for kept in ["latitude", "longitude", "time", "platform"] { - assert!( - schema.field_with_name(kept).is_ok(), - "{kept} is per profile" - ); - } - for dropped in ["pressure", "temperature", "salinity", "quality"] { - assert!( - schema.field_with_name(dropped).is_err(), - "{dropped} needs level as well" - ); - } - assert!(schema.field_with_name("temperature.units").is_ok()); - } - - /// A projection that names only attributes builds a dataset of scalars. - /// A narrowing by a dimension no surviving array carries must leave it - /// alone, and not refuse a dimension the dataset no longer has. - #[tokio::test] - async fn narrowing_an_attribute_only_dataset_keeps_its_scalars() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - let view = view(tmp.path(), "set0").await; - let dataset = dataset_from_view( - view, - Some(&["temperature.units".to_string(), ".title".to_string()]), - ) - .await - .unwrap(); - - let narrowed = - project_read_dimensions(dataset, Some(vec!["profile".to_string()]), None).unwrap(); - assert_eq!(names(&narrowed), vec![".title", "temperature.units"]); - } - - /// The values themselves, out of the first dataset. - #[tokio::test] - async fn wide_collection_array_values_read_back() { - let tmp = tempfile::tempdir().unwrap(); - test_support::wide_profiles(tmp.path(), WIDE).await; - let view = view(tmp.path(), "set0").await; - let dataset = dataset_from_view( - view, - Some(&["latitude".to_string(), "temperature".to_string()]), - ) - .await - .unwrap(); - - assert_eq!(dataset.get_array("latitude").unwrap().shape(), &[3]); - assert_eq!(dataset.get_array("temperature").unwrap().shape(), &[3, 4]); - } -} diff --git a/beacon-server/beacon-server-config/src/lib.rs b/beacon-server/beacon-server-config/src/lib.rs index 55a06a23..fd0b6bd1 100644 --- a/beacon-server/beacon-server-config/src/lib.rs +++ b/beacon-server/beacon-server-config/src/lib.rs @@ -10,14 +10,13 @@ use error::Result; // Per-format and storage config types are owned by their crates; beacon-config // composes them here and fills them from the environment. -pub use beacon_arrow_atlas::AtlasConfig; pub use beacon_arrow_bbf::datafusion::BbfConfig; pub use beacon_arrow_hdf5::{Hdf5Config, Hdf5Convention}; pub use beacon_arrow_netcdf::datafusion::NetcdfConfig; pub use beacon_arrow_zarr::ZarrConfig; pub use beacon_common::CrawlerConfig; -pub use beacon_datafusion_ext::type_widening::TypeConflict; pub use beacon_common::FileStatsConfig; +pub use beacon_datafusion_ext::type_widening::TypeConflict; #[derive(Debug, Clone)] pub struct Config { @@ -32,7 +31,6 @@ pub struct Config { pub netcdf: NetcdfConfig, pub hdf5: Hdf5Config, pub zarr: ZarrConfig, - pub atlas: AtlasConfig, pub bbf: BbfConfig, pub crawler: CrawlerConfig, pub file_stats: FileStatsConfig, @@ -495,26 +493,6 @@ struct RawConfig { #[envconfig(from = "BEACON_ZARR_ENABLE_STATISTICS", default = "true")] zarr_enable_statistics: bool, - /// Whether a read reuses an opened Atlas collection. - /// - /// A collection is immutable, so a cached handle stays good until its - /// deletion mask changes. Caching saves the footer read and keeps the - /// decompressed blocks of a collection between queries. - #[envconfig(from = "BEACON_ATLAS_USE_READER_CACHE", default = "true")] - atlas_use_reader_cache: bool, - - /// How many opened Atlas collections to keep. - /// - /// Each entry owns 256 MiB of decompressed blocks and 64 MiB of raw slabs, - /// so this bounds memory as well as handles. - #[envconfig(from = "BEACON_ATLAS_READER_CACHE_SIZE", default = "32")] - atlas_reader_cache_size: u64, - - /// Whether a predicate scan drops the Atlas datasets it can rule out from - /// the collection's statistics, before reading them. - #[envconfig(from = "BEACON_ATLAS_USE_PRUNING", default = "true")] - atlas_use_pruning: bool, - /// The batch size for NetCDF reads, in number of rows. This is used for both local and MPIO reads. #[envconfig(from = "BEACON_BATCH_SIZE", default = "64000")] beacon_batch_size: usize, @@ -707,11 +685,6 @@ impl From for Config { zarr: ZarrConfig { enable_statistics: raw.zarr_enable_statistics, }, - atlas: AtlasConfig { - use_reader_cache: raw.atlas_use_reader_cache, - reader_cache_size: raw.atlas_reader_cache_size, - use_pruning: raw.atlas_use_pruning, - }, bbf: BbfConfig { split_streams_slice: raw.bbf_split_streams_slice, }, From 48ac7a67e55f98b8d191ef2653d31b41ed3497f9 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 8 Sep 2026 16:58:11 +0200 Subject: [PATCH 07/16] fix compile --- beacon-server/beacon-server/src/server/mod.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/beacon-server/beacon-server/src/server/mod.rs b/beacon-server/beacon-server/src/server/mod.rs index c3298a80..e9fa8655 100644 --- a/beacon-server/beacon-server/src/server/mod.rs +++ b/beacon-server/beacon-server/src/server/mod.rs @@ -278,7 +278,6 @@ async fn build_runtime( .with_netcdf_config(config.netcdf.clone()) .with_hdf5_config(config.hdf5.clone()) .with_zarr_config(config.zarr.clone()) - .with_atlas_config(config.atlas.clone()) .with_sql_settings(SqlSettings { default_table: config.sql.default_table.clone(), enable_pushdown_projection: config.sql.enable_pushdown_projection, From ee35a3cbc8357817faf1a9e591cd544e7c9dfe64 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 8 Sep 2026 20:37:44 +0200 Subject: [PATCH 08/16] wip --- Cargo.lock | 4 +- Cargo.toml | 2 +- .../beacon-datafusion-ext/src/nd/encoding.rs | 5 +- beacon-db/beacon-datafusion-ext/src/nd/mod.rs | 3 +- .../beacon-arrow-atlas/src/backend.rs | 207 +++--- .../beacon-arrow-atlas/src/compat.rs | 40 +- .../src/datafusion/metrics.rs | 8 - .../src/datafusion/opener.rs | 654 ++++++++++++++++- .../src/datafusion/pruning.rs | 674 ++++++++---------- .../src/datafusion/source.rs | 5 +- .../beacon-arrow-atlas/src/test_support.rs | 34 + 11 files changed, 1098 insertions(+), 538 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 632bd5e5..39a12c1b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1193,9 +1193,9 @@ dependencies = [ [[package]] name = "atlas-rust" -version = "0.17.1" +version = "0.17.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7af7d700a1036f8630d57847273e9a5f438b6ee0c5361ada94991f46f7feb6ef" +checksum = "dbb4b330af5674f3b1ee4e7ac7892d0c21c1719f9f9ec2ca3b38198f16c55466" dependencies = [ "array-format", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index e4c90b4a..307e4a7f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,7 +83,7 @@ object_store = { version = "0.13.1", features = ["aws", "gcp", "azure", "http"] # The atlas container embeds array-format files verbatim, so a change to the # crate's bytes is a change to the format Beacon reads. Pin it exactly. -atlas-rust = "=0.17.1" +atlas-rust = "=0.17.2" oxcdf = { git = "https://github.com/robinskil/oxcdf.git", version = "0.4.0", features = ["async", "object-store", "ndarray"] } diff --git a/beacon-db/beacon-datafusion-ext/src/nd/encoding.rs b/beacon-db/beacon-datafusion-ext/src/nd/encoding.rs index af0802e5..931403c8 100644 --- a/beacon-db/beacon-datafusion-ext/src/nd/encoding.rs +++ b/beacon-db/beacon-datafusion-ext/src/nd/encoding.rs @@ -322,7 +322,10 @@ pub fn decode_nd_record_batch_row(batch: &RecordBatch, row: usize) -> Result Result { +/// +/// A file opener that builds nd batches itself uses the same rule, so a batch +/// it emits and a batch the decoder rebuilds agree on the grid. +pub fn infer_target(columns: &[NdArrowArray]) -> Result { let mut order: Vec = Vec::new(); if let Some(widest) = columns.iter().max_by_key(|c| c.dims().rank()) { order.extend(widest.dims().iter().cloned()); diff --git a/beacon-db/beacon-datafusion-ext/src/nd/mod.rs b/beacon-db/beacon-datafusion-ext/src/nd/mod.rs index a0ea99f0..06feb524 100644 --- a/beacon-db/beacon-datafusion-ext/src/nd/mod.rs +++ b/beacon-db/beacon-datafusion-ext/src/nd/mod.rs @@ -31,7 +31,8 @@ pub use broadcast::BroadcastMap; pub use dimensions::{Dimension, Dimensions}; pub use encoding::{ decode_nd_record_batch, encode_flat_batch_as_nd, encode_nd_record_batch, encoded_schema, - is_nd_encoded, logical_schema, nd_encoded_field, nd_encoded_field_of, nd_encoded_type, + infer_target, is_nd_encoded, logical_schema, nd_encoded_field, nd_encoded_field_of, + nd_encoded_type, }; pub use optimizer::{NdFilterPushdown, NdProjectionPushdown, is_pushable_expr}; diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/backend.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/backend.rs index 53c507c2..5f328698 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/backend.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/backend.rs @@ -1,11 +1,11 @@ //! The lazy array backends the Atlas reader hands to `beacon-nd-array`. //! -//! [`AtlasArrayBackend`] reads a region of one atlas array on demand. -//! [`AttributeBackend`] holds one attribute value as a rank-0 array. +//! [`AtlasArrayBackend`] reads one dataset's entry of a variable's segment on +//! demand. [`AttributeBackend`] holds one attribute value as a rank-0 array. use std::sync::Arc; -use atlas::{DatasetView, FillValue}; +use atlas::{ArrayFile, FillValue}; use beacon_nd_array::{ array::{backend::ArrayBackend, subset::ArraySubset}, datatypes::{NdArrayType, TimestampNanosecond}, @@ -21,10 +21,10 @@ use ndarray::ArrayD; /// entry point, so [`AtlasArrayBackend`] stays generic. #[async_trait::async_trait] pub trait AtlasElement: NdArrayType { - /// Read `shape` elements of `array` from `start`. + /// Read `shape` elements of `dataset`'s entry in `segment` from `start`. async fn read( - view: &DatasetView, - array: &str, + segment: &ArrayFile, + dataset: &str, start: Vec, shape: Vec, ) -> anyhow::Result>; @@ -42,18 +42,17 @@ macro_rules! passthrough { #[async_trait::async_trait] impl AtlasElement for $ty { async fn read( - view: &DatasetView, - array: &str, + segment: &ArrayFile, + dataset: &str, start: Vec, shape: Vec, ) -> anyhow::Result> { - let values = view - .read_array::<$ty>(array, start, shape) + let values = segment + .read_array::<$ty>(dataset, start, shape) .await .map_err(|e| { anyhow::anyhow!( - "Failed to read atlas array '{array}' of dataset '{}': {e}", - view.name() + "Failed to read dataset '{dataset}' from its atlas segment: {e}" ) })?; Ok(values.into_owned()) @@ -86,18 +85,17 @@ passthrough!(Vec); #[async_trait::async_trait] impl AtlasElement for TimestampNanosecond { async fn read( - view: &DatasetView, - array: &str, + segment: &ArrayFile, + dataset: &str, start: Vec, shape: Vec, ) -> anyhow::Result> { - let values = view - .read_array::(array, start, shape) + let values = segment + .read_array::(dataset, start, shape) .await .map_err(|e| { anyhow::anyhow!( - "Failed to read atlas timestamp array '{array}' of dataset '{}': {e}", - view.name() + "Failed to read the timestamps of dataset '{dataset}' from its atlas segment: {e}" ) })?; Ok(values.into_owned().mapv(|ts| TimestampNanosecond(ts.0))) @@ -108,14 +106,16 @@ impl AtlasElement for TimestampNanosecond { } } -/// Reads one atlas array lazily, one requested region at a time. +/// Reads one dataset's entry of an atlas segment lazily, one region at a time. /// -/// The backend holds the [`DatasetView`] rather than the collection and a -/// name. A view is resolved once, when the dataset is built; resolving it per -/// read would cost a linear scan of the collection footer every time. +/// The backend holds the segment itself, not a +/// [`DatasetView`](atlas::DatasetView). A segment holds one variable for every +/// dataset in the collection, keyed by dataset name, so a read is one call on +/// it. A view would resolve the segment through the footer and re-check the +/// element type on every read. pub struct AtlasArrayBackend { - view: Arc, - array: String, + segment: Arc, + dataset: String, shape: Vec, dimensions: Vec, chunk_shape: Vec, @@ -125,8 +125,7 @@ pub struct AtlasArrayBackend { impl std::fmt::Debug for AtlasArrayBackend { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("AtlasArrayBackend") - .field("dataset", &self.view.name()) - .field("array", &self.array) + .field("dataset", &self.dataset) .field("shape", &self.shape) .field("dimensions", &self.dimensions) .field("chunk_shape", &self.chunk_shape) @@ -134,23 +133,31 @@ impl std::fmt::Debug for AtlasArrayBackend { } } -impl AtlasArrayBackend { - pub fn new( - view: Arc, - array: String, - shape: Vec, - dimensions: Vec, - chunk_shape: Vec, - fill_value: Option, - ) -> Self { - Self { - view, - array, +impl AtlasArrayBackend { + /// The backend for `dataset`'s entry in `segment`. + /// + /// The layout comes from the segment, which records the shape, chunking, + /// dimension names and fill value of every entry. The lookup costs no I/O. + /// A dataset the segment does not hold is refused by name. + pub fn try_new(segment: Arc, dataset: String) -> anyhow::Result { + let info = segment.array(&dataset).ok_or_else(|| { + anyhow::anyhow!("dataset '{dataset}' has no entry in this atlas segment") + })?; + let shape = info.shape.iter().map(|&s| s as usize).collect(); + let dimensions = info.dimension_names.clone(); + let chunk_shape = info.chunk_shape.iter().map(|&s| s as usize).collect(); + let fill_value = info + .fill_value + .as_ref() + .map(|fill| T::fill_element(Some(fill))); + Ok(Self { + segment, + dataset, shape, dimensions, chunk_shape, fill_value, - } + }) } } @@ -181,7 +188,7 @@ impl ArrayBackend for AtlasArrayBackend { } async fn read_subset(&self, subset: ArraySubset) -> anyhow::Result> { - T::read(&self.view, &self.array, subset.start, subset.shape).await + T::read(&self.segment, &self.dataset, subset.start, subset.shape).await } } @@ -228,27 +235,25 @@ mod tests { use super::*; use crate::test_support; - /// Open one dataset of a fixture collection. - async fn view(dir: &std::path::Path, dataset: &str) -> Arc { - let atlas = atlas::Atlas::open_path(dir).await.expect("open"); - Arc::new(atlas.dataset(dataset).expect("dataset")) + /// The segment of one variable of a fixture collection. + async fn segment(dir: &std::path::Path, array: &str) -> Arc { + let atlas = test_support::open(dir).await; + Arc::clone(atlas.segment(array).await.expect("segment")) } // ── AtlasArrayBackend ─────────────────────────────────────────────── + /// Shape, dimensions, chunking and fill all come from the segment. #[tokio::test] - async fn the_backend_reports_what_the_footer_holds() { + async fn the_backend_reports_what_the_segment_holds() { let tmp = tempfile::tempdir().unwrap(); test_support::two_datasets(tmp.path()).await; - let backend = AtlasArrayBackend::::new( - view(tmp.path(), "winter").await, - "cycle".to_string(), - vec![4], - vec!["obs".to_string()], - vec![4], - Some(-1), - ); + let backend = AtlasArrayBackend::::try_new( + segment(tmp.path(), "cycle").await, + "winter".to_string(), + ) + .unwrap(); assert_eq!(ArrayBackend::::shape(&backend), vec![4]); assert_eq!( ArrayBackend::::dimensions(&backend), @@ -259,19 +264,32 @@ mod tests { assert_eq!(backend.len(), 4); } + /// A segment holds an entry per dataset that declares the variable. A + /// dataset that does not is refused by name, not read as empty. + #[tokio::test] + async fn a_dataset_the_segment_lacks_is_refused_by_name() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let error = AtlasArrayBackend::::try_new( + segment(tmp.path(), "cycle").await, + "summer".to_string(), + ) + .expect_err("only `winter` declares `cycle`") + .to_string(); + assert!(error.contains("summer"), "{error}"); + } + #[tokio::test] async fn a_full_read_returns_every_value() { let tmp = tempfile::tempdir().unwrap(); test_support::two_datasets(tmp.path()).await; - let backend = AtlasArrayBackend::::new( - view(tmp.path(), "winter").await, - "temperature".to_string(), - vec![4], - vec!["obs".to_string()], - vec![4], - None, - ); + let backend = AtlasArrayBackend::::try_new( + segment(tmp.path(), "temperature").await, + "winter".to_string(), + ) + .unwrap(); let values = backend .read_subset(ArraySubset::new(vec![0], vec![4])) .await @@ -284,14 +302,11 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); test_support::two_datasets(tmp.path()).await; - let backend = AtlasArrayBackend::::new( - view(tmp.path(), "winter").await, - "cycle".to_string(), - vec![4], - vec!["obs".to_string()], - vec![4], - None, - ); + let backend = AtlasArrayBackend::::try_new( + segment(tmp.path(), "cycle").await, + "winter".to_string(), + ) + .unwrap(); let values = backend .read_subset(ArraySubset::new(vec![1], vec![2])) .await @@ -306,14 +321,12 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); test_support::chunked_grid(tmp.path()).await; - let backend = AtlasArrayBackend::::new( - view(tmp.path(), "grid").await, - "temperature".to_string(), - vec![4, 6], - vec!["lat".to_string(), "lon".to_string()], - vec![2, 3], - None, - ); + let backend = AtlasArrayBackend::::try_new( + segment(tmp.path(), "temperature").await, + "grid".to_string(), + ) + .unwrap(); + assert_eq!(ArrayBackend::::chunk_shape(&backend), vec![2, 3]); // Rows 1..3, columns 2..4 of a 4x6 grid whose value is row * 6 + col. // That window straddles all four chunk columns and both chunk rows. let values = backend @@ -333,14 +346,12 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); test_support::chunked_grid(tmp.path()).await; - let backend = AtlasArrayBackend::::new( - view(tmp.path(), "grid").await, - "sparse".to_string(), - vec![4, 6], - vec!["lat".to_string(), "lon".to_string()], - vec![2, 3], - Some(-999.0), - ); + let backend = AtlasArrayBackend::::try_new( + segment(tmp.path(), "sparse").await, + "grid".to_string(), + ) + .unwrap(); + assert_eq!(ArrayBackend::::fill_value(&backend), Some(-999.0)); let values = backend .read_subset(ArraySubset::new(vec![2, 0], vec![1, 3])) .await @@ -353,14 +364,11 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); test_support::two_datasets(tmp.path()).await; - let backend = AtlasArrayBackend::::new( - view(tmp.path(), "winter").await, - "time".to_string(), - vec![4], - vec!["obs".to_string()], - vec![4], - None, - ); + let backend = AtlasArrayBackend::::try_new( + segment(tmp.path(), "time").await, + "winter".to_string(), + ) + .unwrap(); let values = backend .read_subset(ArraySubset::new(vec![0], vec![2])) .await @@ -379,14 +387,11 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); test_support::incompatible(tmp.path()).await; - let backend = AtlasArrayBackend::::new( - view(tmp.path(), "a").await, - "value".to_string(), - vec![2], - vec!["obs".to_string()], - vec![2], - None, - ); + let backend = AtlasArrayBackend::::try_new( + segment(tmp.path(), "value").await, + "a".to_string(), + ) + .unwrap(); let values = backend .read_subset(ArraySubset::new(vec![0], vec![2])) .await diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs index 3dede666..9a5e5767 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/compat.rs @@ -10,13 +10,13 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema}; use arrow::error::ArrowError; -use atlas::{ArrayLayout, Attr, CollectionSchema, DType, DatasetView, FillValue}; +use atlas::{ArrayFile, Attr, CollectionSchema, DType}; use beacon_datafusion_ext::type_widening::{ArrowTypeWidening, LabeledSchema}; use beacon_nd_array::{ NdArray, NdArrayD, datatypes::NdArrayDataType, datatypes::TimestampNanosecond, }; -use crate::backend::{AtlasArrayBackend, AtlasElement, AttributeBackend}; +use crate::backend::{AtlasArrayBackend, AttributeBackend}; // ─── Column names ──────────────────────────────────────────────────────────── @@ -200,41 +200,23 @@ fn merge_types( // ─── Lazy arrays ───────────────────────────────────────────────────────────── -/// Wrap one atlas array as a lazy [`NdArrayD`] over `view`. +/// Wrap one dataset's entry of an atlas segment as a lazy [`NdArrayD`]. /// /// No array data is read here. `dtype` comes from the collection footer, and -/// `layout` from the variable's segment, which one open serves for the whole -/// collection. The values themselves arrive when the engine asks the backend -/// for a subset. +/// the layout from `segment`, which one open serves for the whole collection. +/// The values themselves arrive when the engine asks the backend for a subset. /// /// The chunk shape is the one the writer chose. It is what lets the scan cut a /// dataset on the grid the file actually stores, so one unit of work is one /// stored chunk. pub fn array_to_nd_array( - view: Arc, - array_name: &str, + segment: Arc, + dataset: &str, dtype: &DType, - layout: &ArrayLayout, ) -> anyhow::Result> { - let fill: Option = layout.fill_value().cloned(); - macro_rules! lazy { ($ty:ty) => {{ - let fill = fill - .as_ref() - .map(|value| <$ty as AtlasElement>::fill_element(Some(value))); - let backend = AtlasArrayBackend::<$ty>::new( - view, - array_name.to_string(), - layout.shape().to_vec(), - layout - .dimension_names() - .into_iter() - .map(str::to_string) - .collect(), - layout.chunk_shape().to_vec(), - fill, - ); + let backend = AtlasArrayBackend::<$ty>::try_new(segment, dataset.to_string())?; Ok(Arc::new(NdArray::new_with_backend(backend)?) as Arc) }}; } @@ -254,13 +236,13 @@ pub fn array_to_nd_array( DType::Binary => lazy!(Vec), DType::TimestampNs => lazy!(TimestampNanosecond), DType::Bool => Err(anyhow::anyhow!( - "array '{array_name}' is Bool, which atlas stores no elements of" + "dataset '{dataset}' holds a Bool array, which atlas stores no elements of" )), DType::FixedSizeList { .. } => Err(anyhow::anyhow!( - "array '{array_name}' is a FixedSizeList, which Beacon does not model" + "dataset '{dataset}' holds a FixedSizeList array, which Beacon does not model" )), DType::List { .. } => Err(anyhow::anyhow!( - "array '{array_name}' is a List, which Beacon does not model" + "dataset '{dataset}' holds a List array, which Beacon does not model" )), } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs index 22323d5c..fab10c6b 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs @@ -16,12 +16,6 @@ pub struct AtlasScanMetrics { pub datasets_scanned: Count, /// Datasets it skipped because the collection's statistics ruled them out. pub datasets_pruned: Count, - /// Pruning indexes built. One per collection a predicate scan touches, so a - /// number above the collection count means a partition rebuilt one. - pub index_builds: Count, - /// Datasets those indexes covered, which is what the pruning pass looked - /// at rather than read. - pub index_rows: Count, } impl AtlasScanMetrics { @@ -35,8 +29,6 @@ impl AtlasScanMetrics { .counter("atlas_datasets_scanned", partition), datasets_pruned: MetricBuilder::new(metrics) .counter("atlas_datasets_pruned", partition), - index_builds: MetricBuilder::new(metrics).counter("atlas_index_builds", partition), - index_rows: MetricBuilder::new(metrics).counter("atlas_index_rows", partition), } } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs index 9337af88..ea1c3f0c 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs @@ -1,18 +1,51 @@ +//! One partition's opener: a collection in, nd batches out. +//! +//! A column view says where one column of the scan comes from, for every +//! dataset at once: the variable's segment, or the attribute values keyed by +//! dataset. One dataset then reads as one [`NdRecordBatch`], each column on the +//! axes the dataset stores it on. + use std::sync::Arc; -use arrow::datatypes::SchemaRef; +use arrow::{ + array::{ + ArrayRef, BinaryArray, BooleanArray, Float32Array, Float64Array, Int8Array, Int16Array, + Int32Array, Int64Array, PrimitiveArray, StringArray, UInt8Array, UInt16Array, UInt32Array, + UInt64Array, new_null_array, + }, + buffer::{NullBuffer, ScalarBuffer}, + compute::{cast, kernels::cmp::neq}, + datatypes::{ + ArrowPrimitiveType, Field, FieldRef, Float32Type, Float64Type, Int8Type, Int16Type, + Int32Type, Int64Type, Schema, SchemaRef, TimestampNanosecondType, UInt8Type, UInt16Type, + UInt32Type, UInt64Type, + }, +}; +use atlas::{ + ArrayElement, ArrayFile, Atlas, Attr, DType, FillValue, TimestampNs, array_format::ArrayInfo, +}; +use beacon_datafusion_ext::nd::{ + Dimension, Dimensions, NdArrowArray, NdRecordBatch, encode_nd_record_batch, infer_target, +}; +use beacon_datafusion_ext::type_widening::is_type_conflict; use beacon_nd_array::arrow::metrics::ReadMetrics; use datafusion::{ + common::exec_err, datasource::{ listing::PartitionedFile, physical_plan::{FileOpenFuture, FileOpener}, }, + error::{DataFusionError, Result}, physical_plan::PhysicalExpr, }; -use futures::FutureExt; +use futures::{FutureExt, StreamExt}; +use indexmap::IndexMap; use object_store::ObjectStore; -use crate::{datafusion::metrics::AtlasScanMetrics, store::AtlasReaderCache}; +use crate::{ + datafusion::{metrics::AtlasScanMetrics, pruning::prune_datasets}, + store::{AtlasReaderCache, get_or_open_atlas}, +}; /// One partition's opener: a collection in, its batches out. /// @@ -36,13 +69,620 @@ pub struct AtlasOpener { } impl FileOpener for AtlasOpener { - fn open(&self, file: PartitionedFile) -> datafusion::error::Result { + /// One collection in, one encoded batch per dataset worth reading out. + /// + /// The column views are built once per collection, and every dataset then + /// reads against them. A dataset the deletion mask hides is not read, and + /// neither is one the predicate rules out from the statistics in memory. + fn open(&self, file: PartitionedFile) -> Result { + let store = self.object_store.clone(); + let cache = self.cache.clone(); + let projected_schema = self.projected_schema.clone(); + let logical_schema = self.logical_schema.clone(); + let predicate = self.predicate.clone(); + let scan_metrics = self.scan_metrics.clone(); + let fut = async move { - return Err(datafusion::error::DataFusionError::NotImplemented( - "AtlasOpener::open is not implemented yet".to_string(), - )); + let open_timer = scan_metrics.open_time.timer(); + let atlas = get_or_open_atlas(Some(&cache), store, &file.object_meta) + .await + .map_err(|e| { + DataFusionError::Execution(format!( + "Failed to open atlas collection '{}': {e}", + file.object_meta.location + )) + })?; + let views = Arc::new(column_views(&atlas, &logical_schema).await?); + drop(open_timer); + + let mut datasets = atlas.list_datasets(); + if let Some(predicate) = &predicate { + let prune_timer = scan_metrics.prune_time.timer(); + let listed = datasets.len(); + datasets = prune_datasets(&views, datasets, predicate, &logical_schema).await; + scan_metrics.datasets_pruned.add(listed - datasets.len()); + drop(prune_timer); + } + + let stream = futures::stream::iter(datasets) + .then(move |dataset| { + let views = Arc::clone(&views); + let schema = Arc::clone(&projected_schema); + let metrics = scan_metrics.clone(); + async move { + let nd = fast_view_read_to_record_batch(&views, &dataset).await?; + metrics.datasets_scanned.add(1); + // The encoding names the columns and types the scan + // declared. The scan's schema carries each field's + // marks as well, so the batch takes that schema. + let batch = encode_nd_record_batch(&nd)?.with_schema(schema)?; + Ok::<_, DataFusionError>(batch) + } + }) + .boxed(); + Ok(stream) }; Ok(fut.boxed()) } } + +/// Where each column of the scan comes from, for every dataset at once. +/// +/// One segment open per array, and one attribute sweep per key. Each costs the +/// same however many datasets the collection holds, so a partition pays them +/// once and reads every dataset against the result. A column no dataset +/// declares gets `None`. +pub(crate) async fn column_views( + atlas: &Atlas, + logical_schema: &Schema, +) -> Result>> { + let mut views = IndexMap::with_capacity(logical_schema.fields().len()); + for field in logical_schema.fields() { + let view = if let Some(key) = field.name().strip_prefix('.') { + let map = atlas + .attributes_by_dataset(None, key) + .await + .map_err(external)?; + Some(AtlasColumnView::GlobalAttribute { map }) + } else if let Some((array, key)) = field.name().split_once('.') { + let map = atlas + .attributes_by_dataset(Some(array), key) + .await + .map_err(external)?; + Some(AtlasColumnView::VariableAttribute { + variable: array.to_string(), + map, + }) + } else { + atlas + .try_segment(field.name()) + .await + .map_err(external)? + .map(|segment| AtlasColumnView::Array { + segment: Arc::clone(segment), + }) + }; + views.insert(Arc::clone(field), view); + } + Ok(views) +} + +/// One dataset of the collection as an nd batch, one column per field. +/// +/// A column comes out on the axes the dataset stores it on, and the target +/// grid is their union. A field the dataset lacks is a rank-0 null, which +/// broadcasts to an all-null column: an array no dataset declares, a segment +/// without this dataset's entry, or an attribute nobody set. The decoder makes +/// the same of a null struct row, so the scan sees one thing either way. +async fn fast_view_read_to_record_batch( + views: &IndexMap>, + dataset: &str, +) -> Result { + let mut columns = Vec::with_capacity(views.len()); + for (field, view) in views { + let column = match view { + None => null_scalar(field), + Some(AtlasColumnView::Array { segment }) => match segment.array(dataset) { + Some(info) => array_column(segment, dataset, info, field).await?, + None => null_scalar(field), + }, + Some(AtlasColumnView::GlobalAttribute { map }) + | Some(AtlasColumnView::VariableAttribute { map, .. }) => match map.get(dataset) { + Some(attr) => attr_column(attr, field)?, + None => null_scalar(field), + }, + }; + columns.push(column); + } + let schema = Arc::new(Schema::new(views.keys().cloned().collect::>())); + let target = infer_target(&columns)?; + NdRecordBatch::try_new(schema, columns, target) +} + +/// A rank-0 null. It broadcasts to an all-null column of the target grid. +fn null_scalar(field: &Field) -> NdArrowArray { + NdArrowArray::try_new(new_null_array(field.data_type(), 1), Dimensions::scalar()) + .expect("one element on no axis") +} + +/// `dataset`'s entry of one variable, on the axes the segment records for it. +async fn array_column( + segment: &ArrayFile, + dataset: &str, + info: &ArrayInfo, + field: &Field, +) -> Result { + let dims = Dimensions::try_new( + info.dimension_names + .iter() + .zip(&info.shape) + .map(|(name, &size)| Dimension::new(name.as_str(), size as usize)) + .collect(), + )?; + let values = read_values(segment, dataset, info).await?; + match as_field_type(values, field)? { + Some(values) => NdArrowArray::try_new(values, dims), + None => Ok(null_scalar(field)), + } +} + +/// One attribute value of `dataset`, on no axis. +fn attr_column(attr: &Attr, field: &Field) -> Result { + let values: ArrayRef = match attr { + Attr::Bool(v) => Arc::new(BooleanArray::from(vec![*v])), + Attr::Int8(v) => Arc::new(Int8Array::from(vec![*v])), + Attr::Int16(v) => Arc::new(Int16Array::from(vec![*v])), + Attr::Int32(v) => Arc::new(Int32Array::from(vec![*v])), + Attr::Int64(v) => Arc::new(Int64Array::from(vec![*v])), + Attr::UInt8(v) => Arc::new(UInt8Array::from(vec![*v])), + Attr::UInt16(v) => Arc::new(UInt16Array::from(vec![*v])), + Attr::UInt32(v) => Arc::new(UInt32Array::from(vec![*v])), + Attr::UInt64(v) => Arc::new(UInt64Array::from(vec![*v])), + Attr::Float32(v) => Arc::new(Float32Array::from(vec![*v])), + Attr::Float64(v) => Arc::new(Float64Array::from(vec![*v])), + Attr::String(v) => Arc::new(StringArray::from(vec![v.as_str()])), + Attr::Binary(v) => Arc::new(BinaryArray::from(vec![v.as_slice()])), + // A list has no rank-0 form, and the schema holds no list column. A + // dataset that stores a list under a scalar column's key reads as null. + _ => return Ok(null_scalar(field)), + }; + match as_field_type(values, field)? { + Some(values) => NdArrowArray::try_new(values, Dimensions::scalar()), + None => Ok(null_scalar(field)), + } +} + +/// `values` in the type the table declares for `field`, or `None` for values +/// the table cannot hold. +/// +/// A dataset may store a column narrower than the merged type, and the merge +/// widened it: that is a cast. A column the merge could not join is marked, +/// and a dataset of the other family then reads as null. That is what the mark +/// promises the scan. +fn as_field_type(values: ArrayRef, field: &Field) -> Result> { + if values.data_type() == field.data_type() { + return Ok(Some(values)); + } + match cast(&values, field.data_type()) { + Ok(values) => Ok(Some(values)), + Err(_) if is_type_conflict(field) => Ok(None), + Err(error) => Err(error.into()), + } +} + +/// The flat values of `dataset`'s entry in `segment`, its fill read as null. +async fn read_values(segment: &ArrayFile, dataset: &str, info: &ArrayInfo) -> Result { + let fill = info.fill_value.as_ref(); + match info.dtype { + DType::Int8 => primitive_values::(segment, dataset, fill).await, + DType::Int16 => primitive_values::(segment, dataset, fill).await, + DType::Int32 => primitive_values::(segment, dataset, fill).await, + DType::Int64 => primitive_values::(segment, dataset, fill).await, + DType::UInt8 => primitive_values::(segment, dataset, fill).await, + DType::UInt16 => primitive_values::(segment, dataset, fill).await, + DType::UInt32 => primitive_values::(segment, dataset, fill).await, + DType::UInt64 => primitive_values::(segment, dataset, fill).await, + DType::Float32 => primitive_values::(segment, dataset, fill).await, + DType::Float64 => primitive_values::(segment, dataset, fill).await, + DType::TimestampNs => timestamp_values(segment, dataset, fill).await, + DType::String => text_values(segment, dataset, fill).await, + DType::Binary => binary_values(segment, dataset, fill).await, + DType::Bool | DType::List { .. } | DType::FixedSizeList { .. } => exec_err!( + "dataset '{dataset}' holds a {:?} array, which Beacon does not read", + info.dtype + ), + } +} + +/// The whole entry of `dataset` in `segment`, flat in row-major order. +async fn read_flat(segment: &ArrayFile, dataset: &str) -> Result> { + let values = segment + .read_array::(dataset, vec![], vec![]) + .await + .map_err(external)?; + Ok(values.into_owned().into_raw_vec_and_offset().0) +} + +async fn primitive_values( + segment: &ArrayFile, + dataset: &str, + fill: Option<&FillValue>, +) -> Result +where + A: ArrowPrimitiveType, + A::Native: ArrayElement, +{ + let values = read_flat::(segment, dataset).await?; + let fill = fill.map(|fill| ::fill_element(Some(fill))); + mask_fill( + PrimitiveArray::::new(ScalarBuffer::from(values), None), + fill, + ) +} + +/// Both types are `#[repr(transparent)]` over `i64`. The rename is done +/// element by element all the same, on the type system rather than on layout. +async fn timestamp_values( + segment: &ArrayFile, + dataset: &str, + fill: Option<&FillValue>, +) -> Result { + let values: Vec = read_flat::(segment, dataset) + .await? + .into_iter() + .map(|ts| ts.0) + .collect(); + let fill = fill.map(|fill| ::fill_element(Some(fill)).0); + mask_fill( + PrimitiveArray::::from(values), + fill, + ) +} + +async fn text_values( + segment: &ArrayFile, + dataset: &str, + fill: Option<&FillValue>, +) -> Result { + let values = read_flat::(segment, dataset).await?; + let fill = fill.map(|fill| ::fill_element(Some(fill))); + Ok(Arc::new(StringArray::from_iter(values.into_iter().map( + |value| (fill.as_ref() != Some(&value)).then_some(value), + )))) +} + +async fn binary_values( + segment: &ArrayFile, + dataset: &str, + fill: Option<&FillValue>, +) -> Result { + let values = read_flat::>(segment, dataset).await?; + let fill = fill.map(|fill| as ArrayElement>::fill_element(Some(fill))); + Ok(Arc::new(BinaryArray::from_iter(values.into_iter().map( + |value| (fill.as_ref() != Some(&value)).then_some(value), + )))) +} + +/// `array`, with every element equal to `fill` read as null. +/// +/// One vectorised compare gives the validity. A `NaN` fill masks nothing, as +/// `NaN` equals no value, and the same holds in every other nd format. +fn mask_fill( + array: PrimitiveArray, + fill: Option, +) -> Result { + let Some(fill) = fill else { + return Ok(Arc::new(array)); + }; + let fill = PrimitiveArray::::new_scalar(fill); + let kept = neq(&array, &fill)?; + let nulls = NullBuffer::new(kept.values().clone()); + Ok(Arc::new(PrimitiveArray::::new( + array.values().clone(), + Some(nulls), + ))) +} + +/// An atlas error, as the scan reports it. +fn external(error: impl std::error::Error + Send + Sync + 'static) -> DataFusionError { + DataFusionError::External(Box::new(error)) +} + +/// Where one column of the scan comes from, for every dataset of a collection. +/// +/// The scan reads through it, and pruning judges through it, so both see one +/// resolution of a column name. +pub(crate) enum AtlasColumnView { + /// The variable's segment. It holds the array for every dataset that + /// declares it, keyed by dataset name. + Array { segment: Arc }, + /// A dataset-level attribute, its value per dataset. + GlobalAttribute { map: IndexMap }, + /// An attribute of one array, its value per dataset. + VariableAttribute { + variable: String, + map: IndexMap, + }, +} + +#[cfg(test)] +mod tests { + use arrow::array::{Array, AsArray, RecordBatch}; + use arrow::datatypes::{Float32Type, Float64Type, Int32Type, Int64Type}; + use beacon_datafusion_ext::type_widening::ArrowTypeWidening; + + use super::*; + use crate::{compat, test_support}; + use std::path::Path; + + use beacon_datafusion_ext::nd::{decode_nd_record_batch, encoded_schema}; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Column as ColumnExpr, Literal}; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; + use datafusion::scalar::ScalarValue; + use futures::TryStreamExt; + + /// The column views of a fixture, over the schema `infer_schema` derives. + async fn views(dir: &std::path::Path) -> IndexMap> { + let atlas = test_support::open(dir).await; + let schema = compat::collection_arrow_schema( + &atlas.footer().collection_schema(), + &ArrowTypeWidening::default_extension(), + ) + .unwrap(); + column_views(&atlas, &schema).await.unwrap() + } + + async fn read(dir: &std::path::Path, dataset: &str) -> (NdRecordBatch, RecordBatch) { + let views = views(dir).await; + let nd = fast_view_read_to_record_batch(&views, dataset) + .await + .unwrap(); + let batch = nd.materialize().unwrap(); + (nd, batch) + } + + fn column<'a>(batch: &'a RecordBatch, name: &str) -> &'a ArrayRef { + batch + .column_by_name(name) + .unwrap_or_else(|| panic!("no column {name}")) + } + + /// Every column comes out on the dataset's grid. An attribute has no axis + /// of its own, so it repeats on every row. + #[tokio::test] + async fn a_dataset_reads_every_column_on_its_own_grid() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let (nd, batch) = read(tmp.path(), "winter").await; + + assert_eq!(nd.target().shape(), vec![4]); + assert_eq!(batch.num_rows(), 4); + assert_eq!( + column(&batch, "temperature") + .as_primitive::() + .values() + .to_vec(), + vec![1.0, 2.0, 3.0, 4.0] + ); + assert_eq!( + column(&batch, "cycle") + .as_primitive::() + .values() + .to_vec(), + vec![10, 20, 30, 40] + ); + let season = column(&batch, ".season").as_string::(); + assert!( + (0..4).all(|row| season.value(row) == "winter"), + "a rank-0 attribute repeats on every row" + ); + assert_eq!( + column(&batch, ".year") + .as_primitive::() + .values() + .to_vec(), + vec![2024; 4] + ); + assert_eq!( + column(&batch, "temperature.units") + .as_string::() + .value(3), + "celsius" + ); + } + + /// `summer` declares neither `cycle` nor `time`, sets no `year`, and has no + /// `units` on `temperature`. Each is a column of nulls on summer's grid. + #[tokio::test] + async fn a_column_the_dataset_lacks_is_all_null() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let (_, batch) = read(tmp.path(), "summer").await; + + assert_eq!(batch.num_rows(), 3); + for missing in ["cycle", "time", ".year", "temperature.units"] { + assert_eq!(column(&batch, missing).null_count(), 3, "{missing}"); + } + assert_eq!( + column(&batch, "temperature") + .as_primitive::() + .values() + .to_vec(), + vec![20.0, 21.0, 22.0] + ); + assert_eq!( + column(&batch, ".season").as_string::().value(2), + "summer" + ); + } + + /// A 2-D array keeps both axes, and a cell nobody wrote reads as null. + #[tokio::test] + async fn a_fill_value_reads_as_null_on_a_two_dimensional_grid() { + let tmp = tempfile::tempdir().unwrap(); + test_support::chunked_grid(tmp.path()).await; + + let (nd, batch) = read(tmp.path(), "grid").await; + + assert_eq!(nd.target().shape(), vec![4, 6]); + assert_eq!(batch.num_rows(), 24); + let temperature = column(&batch, "temperature").as_primitive::(); + assert_eq!(temperature.value(7), 7.0, "row 1, column 1 of a 4x6 grid"); + let sparse = column(&batch, "sparse"); + assert_eq!(sparse.null_count(), 12, "two of four rows were written"); + assert!(sparse.is_valid(0)); + assert!(sparse.is_null(23)); + } + + /// `a` stores `value` as `Int16` and `b` as `Float32`. The table declares + /// `Float64`, so each dataset casts up to it. `flag` is `a`'s alone. + #[tokio::test] + async fn a_narrower_dataset_casts_to_the_merged_type() { + let tmp = tempfile::tempdir().unwrap(); + test_support::widening(tmp.path()).await; + + let (_, a) = read(tmp.path(), "a").await; + let (_, b) = read(tmp.path(), "b").await; + + assert_eq!( + column(&a, "value") + .as_primitive::() + .values() + .to_vec(), + vec![1.0, 2.0] + ); + assert_eq!( + column(&b, "value") + .as_primitive::() + .values() + .to_vec(), + vec![3.5, 4.5] + ); + assert_eq!( + column(&a, "flag") + .as_primitive::() + .values() + .to_vec(), + vec![7, 8] + ); + assert_eq!(column(&b, "flag").null_count(), 2); + } + + // ── the opener ────────────────────────────────────────────────────── + + /// An opener over a fixture, built the way `AtlasSource` builds one. + async fn opener(dir: &Path) -> (AtlasOpener, PartitionedFile) { + let atlas = test_support::open(dir).await; + let logical_schema = Arc::new( + compat::collection_arrow_schema( + &atlas.footer().collection_schema(), + &ArrowTypeWidening::default_extension(), + ) + .unwrap(), + ); + let projected_schema = Arc::new(encoded_schema(&logical_schema)); + let (store, marker) = test_support::store_and_marker(dir); + let metrics = ExecutionPlanMetricsSet::new(); + let opener = AtlasOpener { + object_store: store, + cache: AtlasReaderCache::new(4), + projected_schema, + logical_schema, + read_dimensions: None, + batch_size: 8192, + predicate: None, + read_metrics: ReadMetrics::new(&metrics, 0), + scan_metrics: AtlasScanMetrics::new(&metrics, 0), + }; + (opener, PartitionedFile::from(marker)) + } + + async fn stream(opener: &AtlasOpener, file: PartitionedFile) -> Vec { + opener + .open(file) + .unwrap() + .await + .unwrap() + .try_collect() + .await + .unwrap() + } + + /// One encoded batch per dataset, in write order, each on the scan's own + /// schema. + #[tokio::test] + async fn the_opener_streams_one_encoded_batch_per_dataset() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let (opener, file) = opener(tmp.path()).await; + + let batches = stream(&opener, file).await; + + assert_eq!(batches.len(), 2); + for batch in &batches { + assert_eq!( + batch.schema(), + opener.projected_schema, + "the scan's schema, marks and all" + ); + } + let rows: Vec = batches + .iter() + .map(|batch| decode_nd_record_batch(batch).unwrap().num_rows()) + .collect(); + assert_eq!(rows, vec![4, 3], "winter, then summer"); + assert_eq!(opener.scan_metrics.datasets_scanned.value(), 2); + } + + /// The deletion mask hides a dataset from the scan, though not from the + /// schema. + #[tokio::test] + async fn a_deleted_dataset_is_not_streamed() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + test_support::open(tmp.path()) + .await + .delete_dataset("winter") + .await + .unwrap(); + let (opener, file) = opener(tmp.path()).await; + + let batches = stream(&opener, file).await; + + assert_eq!(batches.len(), 1); + let summer = decode_nd_record_batch(&batches[0]) + .unwrap() + .materialize() + .unwrap(); + assert_eq!(summer.num_rows(), 3); + assert_eq!( + column(&summer, "cycle").null_count(), + 3, + "winter's column, summer's nulls" + ); + } + + /// A predicate the statistics can judge skips the datasets it rules out + /// before any of them is read. + #[tokio::test] + async fn a_predicate_prunes_datasets_before_the_read() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 10).await; + let (mut opener, file) = opener(tmp.path()).await; + opener.predicate = Some(Arc::new(BinaryExpr::new( + Arc::new(ColumnExpr::new("temperature", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Float32(Some(45.0)))), + ))); + + let batches = stream(&opener, file).await; + + assert_eq!(batches.len(), 5, "d5 to d9 reach past 45"); + assert_eq!(opener.scan_metrics.datasets_pruned.value(), 5); + assert_eq!(opener.scan_metrics.datasets_scanned.value(), 5); + } +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs index d1336804..0dd3d158 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs @@ -1,36 +1,45 @@ -//! Dropping the datasets a predicate cannot match, from the collection footer. +//! Dropping the datasets a predicate cannot match, from what is in memory. //! //! # One index, not a decision per dataset //! //! A collection can hold millions of datasets. Evaluating a predicate against -//! each one in turn would cost millions of evaluations, and each would open the -//! dataset to get its numbers. +//! each one in turn would cost millions of evaluations. Instead the opener +//! builds one [`PruningIndex`] over the collection: one row per live dataset, +//! and one column of typed Arrow statistics per column the predicate names. +//! DataFusion's [`PruningPredicate`] then judges the whole collection in one +//! vectorised pass, and the result is one bit per dataset. //! -//! Instead the first opener that reaches a collection builds one -//! [`PruningIndex`] over it: one row per live dataset, and one column of typed -//! Arrow statistics per column the predicate names. DataFusion's -//! [`PruningPredicate`] then evaluates the whole collection in one vectorised -//! pass, and the result is a bit per dataset that every partition reads. +//! The inputs are the opener's own column views. A variable's segment records +//! the statistics of every dataset that wrote it, and an attribute view holds +//! every dataset's value. Both are in memory once the views exist, so the +//! index costs no I/O. Reading the views rather than asking the collection +//! again also keeps pruning on the columns the scan reads: a column resolves +//! one way, in [`column_views`](super::opener::column_views). //! -//! One column is one request. Atlas stores a variable in one segment, so -//! [`Atlas::array_stats_by_dataset`] and [`Atlas::attributes_by_dataset`] each -//! return every live dataset's value from a single open — however many datasets -//! there are. The index costs one open per column the predicate names, and no -//! array data at all. +//! # A column the dataset lacks +//! +//! The scan reads such a column as nulls, so the index says so: +//! `null_count == row_count`. DataFusion then drops the dataset for `x > 5` +//! and for `x IS NOT NULL`, and keeps it for `x IS NULL`. +//! +//! The writer counts a cell nobody wrote as null too. With a fill value that is +//! what the scan reads, and the counts hold. Without one the scan reads zeros, +//! so an entry with unwritten cells and no fill value says nothing about its +//! values. Its counts stay unknown, and its dataset stays in. //! //! # Pruning is only ever an optimization //! -//! Every path here fails open: an error, a predicate the engine cannot use, a -//! column with no statistics, or a bound that will not cast all leave the -//! datasets in. A dataset that survives is still filtered row by row above the -//! scan, so a hiccup here costs time and never a row. +//! Every path here fails open: an error, a predicate the engine cannot use, or +//! a bound that will not cast all leave the datasets in. A dataset that +//! survives is still filtered row by row above the scan, so a hiccup here +//! costs time and never a row. use std::collections::HashMap; use std::sync::Arc; -use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use atlas::{ArrayStats, Atlas, Attr, StatValue}; +use arrow::array::{ArrayRef, UInt64Array, new_null_array}; +use arrow::datatypes::{DataType, FieldRef, SchemaRef}; +use atlas::{ArrayFile, Attr, StatValue}; use datafusion::common::Column; use datafusion::common::pruning::PruningStatistics; use datafusion::physical_expr::PhysicalExpr; @@ -39,57 +48,66 @@ use datafusion::physical_optimizer::pruning::PruningPredicate; use datafusion::scalar::ScalarValue; use indexmap::IndexMap; -// ─── What a scan does with the answer ──────────────────────────────────────── - -/// Which datasets of one collection a predicate could still match. -#[derive(Debug)] -pub enum CandidateFilter { - /// Pruning did not apply. Every dataset is read. - KeepAll, - /// One bit per dataset, in the order the plan listed them. - Rows { - kept: BooleanArray, - /// The dataset at each row, so a listing that changed under the plan is - /// detected rather than mis-indexed. - names: Vec, - }, -} +use super::opener::AtlasColumnView; -impl CandidateFilter { - /// Whether the dataset at `position` is worth reading. - /// - /// The name is checked against the row it indexes. A collection is - /// immutable, but its deletion mask is not, so a delete between the plan - /// and the open would shift every row after it. A mismatch keeps the - /// dataset: the filter above the scan decides it either way. - pub fn keeps(&self, position: usize, dataset: &str) -> bool { - match self { - Self::KeepAll => true, - Self::Rows { kept, names } => match names.get(position) { - Some(name) if name == dataset => kept.value(position), - _ => true, - }, - } - } - - /// Whether an index was built at all, as opposed to pruning not applying. - pub fn is_index(&self) -> bool { - matches!(self, Self::Rows { .. }) +/// The datasets of `names` that `predicate` could still match, in order. +/// +/// `logical_schema` must type every column the predicate names, which the +/// scan's own projected schema does: a filter that stays above the scan forces +/// its columns into the projection. `views` must resolve those columns the way +/// the scan reads them. +/// +/// Fails open to `names` on anything it cannot prove. +pub(crate) async fn prune_datasets( + views: &Arc>>, + names: Vec, + predicate: &Arc, + logical_schema: &SchemaRef, +) -> Vec { + let Ok(pruning) = PruningPredicate::try_new(Arc::clone(predicate), Arc::clone(logical_schema)) + else { + // The engine cannot use this predicate shape. + return names; + }; + let referenced = collect_columns(pruning.orig_expr()); + if referenced.is_empty() || names.is_empty() { + return names; } - /// How many datasets this filter drops. For diagnostics. - pub fn pruned(&self) -> usize { - match self { - Self::KeepAll => 0, - Self::Rows { kept, .. } => kept.len() - kept.true_count(), - } - } + // A column the predicate names, with the type the table gives it. + let wanted: Vec<(String, DataType)> = referenced + .iter() + .filter_map(|column| { + let (field, _) = views + .iter() + .find(|(field, _)| field.name() == column.name())?; + Some((column.name().to_string(), field.data_type().clone())) + }) + .collect(); + if wanted.is_empty() { + // Nothing the predicate names is a column of the scan. + return names; + } + + // The pivot is pure CPU over what is in memory, and a million rows is + // real work, so it does not run on the async runtime. + let names: Arc<[String]> = names.into(); + let (views, rows) = (Arc::clone(views), Arc::clone(&names)); + let built = tokio::task::spawn_blocking(move || build_index(&views, &rows, &wanted)).await; + let Ok(index) = built else { + return names.to_vec(); + }; - /// How many rows the index behind this filter holds. - pub fn rows(&self) -> usize { - match self { - Self::KeepAll => 0, - Self::Rows { kept, .. } => kept.len(), + match pruning.prune(&index) { + Ok(kept) => names + .iter() + .zip(kept) + .filter(|(_, keep)| *keep) + .map(|(name, _)| name.clone()) + .collect(), + Err(e) => { + tracing::debug!("atlas pruning fell back to reading every dataset: {e}"); + names.to_vec() } } } @@ -139,7 +157,7 @@ impl PruningStatistics for PruningIndex { &self, _column: &Column, _values: &std::collections::HashSet, - ) -> Option { + ) -> Option { // An attribute's value is exact, so an `IN` list could prune on one. // Not yet: every column here reports a range, and a range says nothing // about membership. @@ -149,214 +167,70 @@ impl PruningStatistics for PruningIndex { // ─── Building it ───────────────────────────────────────────────────────────── -/// The logical schema behind an nd-encoded one. -/// -/// A scan carries its columns as `beacon.nd` structs, and a predicate is -/// written against the values inside them. A field whose type does not decode -/// keeps its own type, which simply leaves it unprunable. -pub fn logical_schema(encoded: &Schema) -> SchemaRef { - let fields: Vec> = encoded - .fields() - .iter() - .map(|field| { - let value_type = beacon_datafusion_ext::nd::encoding::nd_value_type(field.data_type()) - .unwrap_or_else(|_| field.data_type().clone()); - Arc::new(Field::new(field.name(), value_type, true)) - }) - .collect(); - Arc::new(Schema::new(fields)) -} - -/// Which datasets of `atlas` could satisfy `predicate`. -/// -/// `logical_schema` must type every column the predicate names, which the -/// scan's own projected schema does: a filter that stays above the scan forces -/// its columns into the projection. -/// -/// Fails open to [`CandidateFilter::KeepAll`] on anything it cannot prove. -pub async fn candidate_filter( - atlas: &Arc, - predicate: &Arc, - logical_schema: &SchemaRef, -) -> CandidateFilter { - let Ok(pruning) = PruningPredicate::try_new(Arc::clone(predicate), Arc::clone(logical_schema)) - else { - // The engine cannot use this predicate shape. - return CandidateFilter::KeepAll; - }; - - let referenced = collect_columns(pruning.orig_expr()); - if referenced.is_empty() { - return CandidateFilter::KeepAll; - } - - let names = atlas.list_datasets(); - if names.is_empty() { - return CandidateFilter::KeepAll; - } - - // Fetching is one request per column, whatever the dataset count. The - // pivot after it is pure CPU over what is then in memory, and a million - // rows is real work, so it does not run on the async runtime. - let wanted: Vec<(String, DataType)> = referenced +/// Pivot the views into one [`StatColumn`] per wanted column. +fn build_index( + views: &IndexMap>, + names: &[String], + wanted: &[(String, DataType)], +) -> PruningIndex { + let columns = wanted .iter() - .filter_map(|column| { - let field = schema_field(logical_schema, column.name())?; - Some((column.name().to_string(), field)) - }) - .collect(); - - let mut fetched: Vec<(String, DataType, Measured)> = Vec::with_capacity(wanted.len()); - let arrays = atlas.list_arrays(); - for (column, target) in wanted { - let measured = if arrays.iter().any(|array| array == &column) { - match atlas.array_stats_by_dataset(&column).await { - Ok(stats) => Some(Measured::Array(stats)), - Err(e) => { - tracing::debug!(column, "atlas statistics unavailable for pruning: {e}"); - None + .filter_map(|(column, target)| { + let (_, view) = views.iter().find(|(field, _)| field.name() == column)?; + let packed = match view { + // No dataset declares the column. The scan reads nulls. + None => all_null_column(names.len(), target), + Some(AtlasColumnView::Array { segment }) => { + pack_array_column(segment, names, target) } - } - } else { - attribute_values(atlas, &column) - .await - .map(Measured::Attribute) - }; - if let Some(measured) = measured { - fetched.push((column, target, measured)); - } - } - - if fetched.is_empty() { - // Nothing the predicate names has statistics, so nothing can be ruled - // out. - return CandidateFilter::KeepAll; - } - - let built = tokio::task::spawn_blocking(move || { - let index = build_index(&names, fetched); - (names, index) - }) - .await; - - let Ok((names, index)) = built else { - return CandidateFilter::KeepAll; - }; - - match pruning.prune(&index) { - Ok(kept) => CandidateFilter::Rows { - kept: BooleanArray::from(kept), - names, - }, - Err(e) => { - tracing::debug!("atlas pruning fell back to reading every dataset: {e}"); - CandidateFilter::KeepAll - } - } -} - -/// The type a column carries in the scan's own schema, or `None` when it holds -/// no such column and nothing can be typed against it. -fn schema_field(schema: &SchemaRef, column: &str) -> Option { - schema - .field_with_name(column) - .ok() - .map(|field| field.data_type().clone()) -} - -/// What one request measured about one column, across every live dataset. -enum Measured { - /// An array's statistics. A dataset that wrote the array has an entry, and - /// [`ArrayStats::name`] names it. - Array(Vec), - /// An attribute's value, keyed by dataset. - Attribute(IndexMap), -} - -/// Every live dataset's value for one attribute column, or `None` when the -/// collection carries no such attribute. -/// -/// A column is `.key` at dataset scope, or `array.key` at array scope. An array -/// name and an attribute key may both hold dots, so every split of the latter -/// is a candidate and the first that finds a value wins. -async fn attribute_values(atlas: &Atlas, column: &str) -> Option> { - let found = |values: IndexMap| (!values.is_empty()).then_some(values); - - if let Some(key) = column.strip_prefix('.') { - return found(atlas.attributes_by_dataset(None, key).await.ok()?); - } - for (index, character) in column.char_indices() { - if character != '.' { - continue; - } - let (array, rest) = column.split_at(index); - if let Ok(values) = atlas.attributes_by_dataset(Some(array), &rest[1..]).await - && let Some(values) = found(values) - { - return Some(values); - } - } - None -} - -/// Pivot what was fetched into one [`StatColumn`] per column. -fn build_index(names: &[String], fetched: Vec<(String, DataType, Measured)>) -> PruningIndex { - // Where each dataset sits, so a pass in write order can be scattered into - // rows without a search. - let row_of: HashMap<&str, usize> = names - .iter() - .enumerate() - .map(|(row, name)| (name.as_str(), row)) + Some(AtlasColumnView::GlobalAttribute { map }) + | Some(AtlasColumnView::VariableAttribute { map, .. }) => { + pack_attribute_column(map, names, target) + } + }; + Some((column.clone(), packed)) + }) .collect(); - let mut columns = HashMap::new(); - for (column, target, measured) in fetched { - let packed = match measured { - Measured::Array(stats) => { - Some(pack_array_column(&stats, names.len(), &row_of, &target)) - } - Measured::Attribute(values) => { - pack_attribute_column(&values, names.len(), &row_of, &target) - } - }; - if let Some(packed) = packed { - columns.insert(column, packed); - } - } - PruningIndex { rows: names.len(), columns, } } -/// One array column. +/// One array column, from the segment that holds the variable. /// -/// A dataset with no entry for the array keeps a null bound and unknown counts. -/// That is what a dataset which does not declare the array looks like, and it -/// is also what one that declared it and never wrote it looks like — both must -/// stay in, and a null does exactly that. -fn pack_array_column( - stats: &[ArrayStats], - rows: usize, - row_of: &HashMap<&str, usize>, - target: &DataType, -) -> StatColumn { - let null = ScalarValue::try_from(target).unwrap_or(ScalarValue::Null); +/// A dataset the segment has no entry for does not declare the array. The scan +/// reads it as nulls, and `null_count == row_count` says so. An entry without +/// statistics, or one whose unwritten cells read as zeros rather than as the +/// nulls the writer counted, says nothing, and that dataset stays in. +fn pack_array_column(segment: &ArrayFile, names: &[String], target: &DataType) -> StatColumn { + let rows = names.len(); + let null = null_of(target); let mut mins = vec![null.clone(); rows]; let mut maxes = vec![null.clone(); rows]; let mut null_counts: Vec> = vec![None; rows]; let mut row_counts: Vec> = vec![None; rows]; - for entry in stats { - // A per-dataset entry names its dataset, not its array. - let Some(&row) = row_of.get(entry.name.as_str()) else { + for (row, name) in names.iter().enumerate() { + let Some(info) = segment.array(name) else { + null_counts[row] = Some(1); + row_counts[row] = Some(1); + continue; + }; + let Some(stats) = info.stats.as_ref() else { continue; }; - mins[row] = stat_to_scalar(entry.min.as_ref(), target, &null); - maxes[row] = stat_to_scalar(entry.max.as_ref(), target, &null); - null_counts[row] = Some(entry.null_count); - row_counts[row] = Some(entry.row_count); + if info.fill_value.is_none() && stats.null_count > 0 { + // The nulls the writer counted are cells nobody wrote. Without a + // fill value the scan reads them as zeros, which the bounds do + // not cover either. + continue; + } + mins[row] = stat_to_scalar(stats.min.as_ref(), target, &null); + maxes[row] = stat_to_scalar(stats.max.as_ref(), target, &null); + null_counts[row] = Some(stats.null_count); + row_counts[row] = Some(stats.row_count); } StatColumn { @@ -367,55 +241,66 @@ fn pack_array_column( } } -/// One attribute column, or `None` when no value of it casts. +/// One attribute column. /// /// An attribute's value is exact, so it is both the minimum and the maximum of -/// its dataset. That prunes an equality on a dataset-level attribute — the -/// platform a file came from, say — from one request. +/// its dataset, on the one cell the scan reads. That prunes an equality on a +/// dataset-level attribute, the platform a file came from, say. A dataset +/// without the key reads as null, and the counts say so. A list, a `NaN`, or a +/// value that will not cast bounds nothing, and that dataset stays in. fn pack_attribute_column( - found: &IndexMap, - rows: usize, - row_of: &HashMap<&str, usize>, + values: &IndexMap, + names: &[String], target: &DataType, -) -> Option { - let null = ScalarValue::try_from(target).unwrap_or(ScalarValue::Null); - let mut values = vec![null.clone(); rows]; +) -> StatColumn { + let rows = names.len(); + let null = null_of(target); + let mut bounds = vec![null.clone(); rows]; let mut null_counts: Vec> = vec![None; rows]; - let mut seen = false; + let mut row_counts: Vec> = vec![None; rows]; - for (dataset, attr) in found { - let Some(&row) = row_of.get(dataset.as_str()) else { + for (row, name) in names.iter().enumerate() { + let Some(attr) = values.get(name) else { + null_counts[row] = Some(1); + row_counts[row] = Some(1); continue; }; - let Some(scalar) = attr_to_scalar(attr) else { + let Some(scalar) = attr_to_scalar(attr).and_then(|scalar| scalar.cast_to(target).ok()) + else { continue; }; - values[row] = scalar.cast_to(target).unwrap_or_else(|_| null.clone()); - // One value, and it is not the fill of anything. + bounds[row] = scalar; null_counts[row] = Some(0); - seen = true; + row_counts[row] = Some(1); } - if !seen { - return None; + StatColumn { + min: scalars_to_array(bounds.clone(), rows, target), + max: scalars_to_array(bounds, rows, target), + null_count: Arc::new(UInt64Array::from(null_counts)), + row_count: Arc::new(UInt64Array::from(row_counts)), } +} - let min = scalars_to_array(values.clone(), rows, target); - let max = scalars_to_array(values, rows, target); - Some(StatColumn { - min, - max, - null_count: Arc::new(UInt64Array::from(null_counts)), - // An attribute is one value broadcast over whatever grid the dataset - // has, so its row count is not the dataset's. Unknown is honest. - row_count: Arc::new(UInt64Array::from(vec![None::; rows])), - }) +/// A column every dataset reads as null. +fn all_null_column(rows: usize, target: &DataType) -> StatColumn { + let ones = || Arc::new(UInt64Array::from(vec![1u64; rows])) as ArrayRef; + StatColumn { + min: new_null_array(target, rows), + max: new_null_array(target, rows), + null_count: ones(), + row_count: ones(), + } +} + +/// The null of `target`, or the untyped null for a type that has none. +fn null_of(target: &DataType) -> ScalarValue { + ScalarValue::try_from(target).unwrap_or(ScalarValue::Null) } /// Pack scalars into one typed array, or a column of nulls when they will not. fn scalars_to_array(values: Vec, rows: usize, target: &DataType) -> ArrayRef { - ScalarValue::iter_to_array(values) - .unwrap_or_else(|_| arrow::array::new_null_array(target, rows)) + ScalarValue::iter_to_array(values).unwrap_or_else(|_| new_null_array(target, rows)) } /// An atlas statistic as a scalar of the table's own type. @@ -463,10 +348,16 @@ fn attr_to_scalar(attr: &Attr) -> Option { #[cfg(test)] mod tests { + use arrow::datatypes::{Field, Schema}; + use atlas::Atlas; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{ + BinaryExpr, Column as ColumnExpr, IsNullExpr, Literal, + }; + use super::*; + use crate::datafusion::opener::column_views; use crate::test_support; - use datafusion::logical_expr::Operator; - use datafusion::physical_expr::expressions::{BinaryExpr, Column as ColumnExpr, Literal}; fn schema(name: &str, data_type: DataType) -> SchemaRef { Arc::new(Schema::new(vec![Field::new(name, data_type, true)])) @@ -480,20 +371,19 @@ mod tests { )) } - /// The datasets a predicate leaves in, in listing order. + fn is_null(column: &str) -> Arc { + Arc::new(IsNullExpr::new(Arc::new(ColumnExpr::new(column, 0)))) + } + + /// The datasets a predicate leaves in, in listing order. The views are the + /// ones the scan would read. async fn kept( atlas: &Arc, predicate: Arc, schema: SchemaRef, ) -> Vec { - let filter = candidate_filter(atlas, &predicate, &schema).await; - atlas - .list_datasets() - .into_iter() - .enumerate() - .filter(|(position, name)| filter.keeps(*position, name)) - .map(|(_, name)| name) - .collect() + let views = Arc::new(column_views(atlas, &schema).await.unwrap()); + prune_datasets(&views, atlas.list_datasets(), &predicate, &schema).await } // ── the index over array statistics ───────────────────────────────── @@ -505,7 +395,6 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); test_support::ranged(tmp.path(), 10).await; let atlas = test_support::open(tmp.path()).await; - let schema = schema("temperature", DataType::Float32); let survivors = kept( &atlas, @@ -514,7 +403,7 @@ mod tests { Operator::Gt, ScalarValue::Float32(Some(45.0)), ), - schema, + schema("temperature", DataType::Float32), ) .await; assert_eq!(survivors, vec!["d5", "d6", "d7", "d8", "d9"]); @@ -558,56 +447,26 @@ mod tests { assert_eq!(survivors, atlas.list_datasets()); } - /// The index holds one row per live dataset, in listing order, and the - /// filter reads it by position. - #[tokio::test] - async fn the_index_holds_one_row_per_live_dataset() { - let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 10).await; - let atlas = test_support::open(tmp.path()).await; - - let filter = candidate_filter( - &atlas, - &binary( - "temperature", - Operator::Gt, - ScalarValue::Float32(Some(45.0)), - ), - &schema("temperature", DataType::Float32), - ) - .await; - assert_eq!(filter.rows(), 10); - assert_eq!(filter.pruned(), 5); - } - - /// A deleted dataset has no row at all, and the rows after it shift up. - /// That is why the filter checks the name it was given against the row it - /// indexes. + /// A deleted dataset is not in the list, so it is neither judged nor read. + /// The segment still holds its entry, and that entry is never looked at. #[tokio::test] - async fn a_delete_shifts_the_rows_and_the_name_check_catches_it() { + async fn a_deleted_dataset_is_neither_judged_nor_read() { let tmp = tempfile::tempdir().unwrap(); test_support::ranged(tmp.path(), 6).await; let atlas = test_support::open(tmp.path()).await; atlas.delete_dataset("d0").await.unwrap(); - let filter = candidate_filter( + let survivors = kept( &atlas, - &binary( + binary( "temperature", - Operator::Gt, - ScalarValue::Float32(Some(45.0)), + Operator::GtEq, + ScalarValue::Float32(Some(0.0)), ), - &schema("temperature", DataType::Float32), + schema("temperature", DataType::Float32), ) .await; - assert_eq!(filter.rows(), 5, "the deleted dataset has no row"); - - // Row 0 is now d1. A plan made before the delete would ask about d0 - // there, and that must not read as d1's answer. - assert!( - filter.keeps(0, "d0"), - "a name that does not match its row is kept, not mis-indexed" - ); + assert_eq!(survivors, vec!["d1", "d2", "d3", "d4", "d5"]); } // ── mixed and awkward types ───────────────────────────────────────── @@ -652,8 +511,8 @@ mod tests { ); } - /// A dataset-level attribute is exact, so an equality on it prunes from the - /// footer alone. + /// A dataset-level attribute is exact, so an equality on it prunes from + /// what is in memory alone. #[tokio::test] async fn an_attribute_predicate_prunes() { let tmp = tempfile::tempdir().unwrap(); @@ -673,54 +532,117 @@ mod tests { assert_eq!(survivors, vec!["d3"]); } - // ── failing open ──────────────────────────────────────────────────── + // ── a column the dataset lacks ────────────────────────────────────── + /// `summer` never set `year`. The scan reads the column as null for it, so + /// an equality drops it and an `IS NULL` keeps it alone. #[tokio::test] - async fn a_column_with_no_statistics_prunes_nothing() { + async fn a_missing_attribute_reads_as_null_and_prunes_as_null() { let tmp = tempfile::tempdir().unwrap(); - test_support::ranged(tmp.path(), 4).await; + test_support::two_datasets(tmp.path()).await; let atlas = test_support::open(tmp.path()).await; + let schema = schema(".year", DataType::Int64); - let survivors = kept( - &atlas, - binary("ghost", Operator::Gt, ScalarValue::Float32(Some(0.0))), - schema("ghost", DataType::Float32), - ) - .await; - assert_eq!(survivors, atlas.list_datasets()); + assert_eq!( + kept( + &atlas, + binary(".year", Operator::Eq, ScalarValue::Int64(Some(2024))), + Arc::clone(&schema) + ) + .await, + vec!["winter"] + ); + assert_eq!(kept(&atlas, is_null(".year"), schema).await, vec!["summer"]); } - /// A column one dataset declares and another does not: the one without it - /// has no bound, so it stays in and its rows are decided above the scan. + /// Only `a` declares `flag`, and it holds [7, 8]. `b` reads the column as + /// null, so a comparison drops it and an `IS NULL` keeps it alone. #[tokio::test] - async fn a_dataset_that_lacks_the_column_is_never_pruned_on_it() { + async fn a_dataset_that_lacks_the_array_reads_as_null_and_prunes_as_null() { let tmp = tempfile::tempdir().unwrap(); test_support::widening(tmp.path()).await; let atlas = test_support::open(tmp.path()).await; + let schema = schema("flag", DataType::Int32); + + assert_eq!( + kept( + &atlas, + binary("flag", Operator::Gt, ScalarValue::Int32(Some(5))), + Arc::clone(&schema) + ) + .await, + vec!["a"] + ); + assert!( + kept( + &atlas, + binary("flag", Operator::Gt, ScalarValue::Int32(Some(100))), + Arc::clone(&schema) + ) + .await + .is_empty(), + "a is ruled out by its range, b by its nulls" + ); + assert_eq!(kept(&atlas, is_null("flag"), schema).await, vec!["b"]); + } + + /// A column no dataset declares is null everywhere the scan looks. + #[tokio::test] + async fn a_column_no_dataset_declares_is_null_everywhere() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 4).await; + let atlas = test_support::open(tmp.path()).await; + let schema = schema("ghost", DataType::Float32); + + assert!( + kept( + &atlas, + binary("ghost", Operator::Gt, ScalarValue::Float32(Some(0.0))), + Arc::clone(&schema) + ) + .await + .is_empty() + ); + assert_eq!( + kept(&atlas, is_null("ghost"), schema).await, + atlas.list_datasets() + ); + } + + /// `d` declares `value` with no fill value and never writes it. The writer + /// counted both cells as null, yet the scan reads them as zeros, so the + /// index must not trust the count. `d` stays in. `w` wrote [5, 6], and + /// its statistics rule it out. + #[tokio::test] + async fn a_declared_array_nobody_wrote_is_never_pruned() { + let tmp = tempfile::tempdir().unwrap(); + test_support::declared_unwritten(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; - // Only `a` declares `flag`, and it holds [7, 8]. let survivors = kept( &atlas, - binary("flag", Operator::Gt, ScalarValue::Int32(Some(100))), - schema("flag", DataType::Int32), + binary("value", Operator::Eq, ScalarValue::Int32(Some(0))), + schema("value", DataType::Int32), ) .await; - assert_eq!(survivors, vec!["b"], "a is ruled out, b cannot be"); + assert_eq!(survivors, vec!["d"]); } + // ── failing open ──────────────────────────────────────────────────── + #[tokio::test] async fn a_collection_with_no_datasets_prunes_nothing() { let tmp = tempfile::tempdir().unwrap(); test_support::empty(tmp.path()).await; let atlas = test_support::open(tmp.path()).await; - let filter = candidate_filter( + let survivors = kept( &atlas, - &binary("temperature", Operator::Gt, ScalarValue::Float32(Some(0.0))), - &schema("temperature", DataType::Float32), + binary("temperature", Operator::Gt, ScalarValue::Float32(Some(0.0))), + schema("temperature", DataType::Float32), ) .await; - assert!(matches!(filter, CandidateFilter::KeepAll)); + assert!(survivors.is_empty(), "nothing in, nothing out"); } // ── the pieces ────────────────────────────────────────────────────── @@ -759,16 +681,13 @@ mod tests { assert!(attr_to_scalar(&Attr::Int32List(vec![1, 2])).is_none()); assert!(attr_to_scalar(&Attr::Float64(f64::NAN)).is_none()); assert_eq!( - attr_to_scalar(&Attr::Int64(7)), - Some(ScalarValue::Int64(Some(7))) + attr_to_scalar(&Attr::String("p1".into())), + Some(ScalarValue::Utf8(Some("p1".into()))) ); } - /// The whole point of an index: a collection of hundreds of thousands of - /// datasets is judged in one vectorised pass. - /// /// The index is built by hand here. Writing that many real datasets would - /// take minutes and prove nothing extra — what this pins is that the + /// take minutes and prove nothing extra. What this pins is that the /// evaluation is one pass over Arrow arrays rather than a decision per /// dataset. #[test] @@ -815,21 +734,4 @@ mod tests { let expected = ROWS - THRESHOLD as usize; assert_eq!(kept.iter().filter(|keep| **keep).count(), expected); } - - /// The scan's schema is nd-encoded; a predicate is written against the - /// values inside it. - #[test] - fn the_logical_schema_unwraps_the_encoding() { - let logical = Schema::new(vec![Field::new("temperature", DataType::Float32, true)]); - let encoded = beacon_datafusion_ext::nd::encoded_schema(&logical); - assert_ne!( - encoded.field(0).data_type(), - &DataType::Float32, - "the encoded form is a struct" - ); - assert_eq!( - logical_schema(&encoded).field(0).data_type(), - &DataType::Float32 - ); - } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs index f73b3010..633a564a 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs @@ -51,8 +51,9 @@ use datafusion::{ use futures::{FutureExt, StreamExt, TryStreamExt}; use object_store::ObjectStore; +use beacon_datafusion_ext::nd::logical_schema; + use crate::datafusion::metrics::AtlasScanMetrics; -use crate::datafusion::pruning::{CandidateFilter, candidate_filter, logical_schema}; use crate::store::{AtlasReaderCache, get_or_open_atlas}; use crate::{compat, datafusion::opener::AtlasOpener}; @@ -115,7 +116,7 @@ impl FileSource for AtlasSource { cache: self.cache.clone(), // A predicate is written against the values, not the encoding the // scan carries them in. - logical_schema: logical_schema(&projected_schema), + logical_schema: logical_schema(&projected_schema)?, projected_schema, read_dimensions: self.read_dimensions.clone(), batch_size: self.batch_size, diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs index 66122372..d6c52fe9 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs @@ -352,6 +352,40 @@ pub async fn empty(dir: &Path) { writer.finish().await.expect("finish the collection"); } +/// One dataset declares an array and never writes it, beside one that does. +/// +/// - `d`: `value: Int32[2]`, defined with no fill and never written. It reads +/// as zeros, while its statistics count both cells as null. +/// - `w`: `value: Int32[2] = [5, 6]`. +/// +/// A predicate on `value` can rule `w` out and must leave `d` in. +pub async fn declared_unwritten(dir: &Path) { + let writer = AtlasWriter::create_path(dir, WriterConfig::default()) + .await + .expect("create the collection"); + + { + let mut d = writer.add_dataset("d").await.expect("add d"); + d.define_array::("value", vec!["obs".into()], vec![2], None, None) + .await + .expect("define value"); + d.finish().await.expect("finish d"); + } + + { + let mut w = writer.add_dataset("w").await.expect("add w"); + w.define_array::("value", vec!["obs".into()], vec![2], None, None) + .await + .expect("define value"); + w.write_array("value", vec![0], arr1(&[5i32, 6]).into_dyn().view()) + .await + .expect("write value"); + w.finish().await.expect("finish w"); + } + + writer.finish().await.expect("finish the collection"); +} + /// A wide collection on two dimensions, `profile` and `level`. /// /// Dataset `i` is named `set{i}` and holds `shapes[i]` as From 1e70efd3b25b286d7b81f584595950484dd5ad58 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 8 Sep 2026 23:38:14 +0200 Subject: [PATCH 09/16] wip --- .../src/datafusion/opener.rs | 414 ++++++++---------- .../beacon-nd-array/src/dataset/default.rs | 392 +++++++++++++++++ .../beacon-nd-array/src/dataset/mod.rs | 66 ++- .../beacon-nd-array/src/dataset/source.rs | 72 +++ 4 files changed, 685 insertions(+), 259 deletions(-) create mode 100644 beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs create mode 100644 beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/source.rs diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs index ea1c3f0c..257c717b 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs @@ -2,35 +2,26 @@ //! //! A column view says where one column of the scan comes from, for every //! dataset at once: the variable's segment, or the attribute values keyed by -//! dataset. One dataset then reads as one [`NdRecordBatch`], each column on the -//! axes the dataset stores it on. +//! dataset. One dataset is then a lazy [`Dataset`] over those views. It reads +//! one stored chunk at a time as an [`NdRecordBatch`], each column on the axes +//! the dataset stores it on. -use std::sync::Arc; +use std::{any::Any, sync::Arc}; use arrow::{ - array::{ - ArrayRef, BinaryArray, BooleanArray, Float32Array, Float64Array, Int8Array, Int16Array, - Int32Array, Int64Array, PrimitiveArray, StringArray, UInt8Array, UInt16Array, UInt32Array, - UInt64Array, new_null_array, - }, - buffer::{NullBuffer, ScalarBuffer}, - compute::{cast, kernels::cmp::neq}, - datatypes::{ - ArrowPrimitiveType, Field, FieldRef, Float32Type, Float64Type, Int8Type, Int16Type, - Int32Type, Int64Type, Schema, SchemaRef, TimestampNanosecondType, UInt8Type, UInt16Type, - UInt32Type, UInt64Type, - }, -}; -use atlas::{ - ArrayElement, ArrayFile, Atlas, Attr, DType, FillValue, TimestampNs, array_format::ArrayInfo, -}; -use beacon_datafusion_ext::nd::{ - Dimension, Dimensions, NdArrowArray, NdRecordBatch, encode_nd_record_batch, infer_target, + array::{ArrayRef, RecordBatch, new_null_array}, + compute::cast, + datatypes::{Field, FieldRef, Schema, SchemaRef}, }; +use atlas::{ArrayFile, Atlas, Attr}; +use beacon_datafusion_ext::nd::{Dimensions, NdArrowArray, NdRecordBatch, encode_nd_record_batch}; use beacon_datafusion_ext::type_widening::is_type_conflict; -use beacon_nd_array::arrow::metrics::ReadMetrics; +use beacon_nd_array::{ + NdArrayD, + arrow::metrics::ReadMetrics, + dataset::{default::Dataset, source::DatasetSource}, +}; use datafusion::{ - common::exec_err, datasource::{ listing::PartitionedFile, physical_plan::{FileOpenFuture, FileOpener}, @@ -38,11 +29,12 @@ use datafusion::{ error::{DataFusionError, Result}, physical_plan::PhysicalExpr, }; -use futures::{FutureExt, StreamExt}; +use futures::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream}; use indexmap::IndexMap; use object_store::ObjectStore; use crate::{ + compat, datafusion::{metrics::AtlasScanMetrics, pruning::prune_datasets}, store::{AtlasReaderCache, get_or_open_atlas}, }; @@ -69,7 +61,8 @@ pub struct AtlasOpener { } impl FileOpener for AtlasOpener { - /// One collection in, one encoded batch per dataset worth reading out. + /// One collection in, one encoded batch per stored chunk of every dataset + /// worth reading out. /// /// The column views are built once per collection, and every dataset then /// reads against them. A dataset the deletion mask hides is not read, and @@ -105,20 +98,15 @@ impl FileOpener for AtlasOpener { } let stream = futures::stream::iter(datasets) - .then(move |dataset| { - let views = Arc::clone(&views); - let schema = Arc::clone(&projected_schema); - let metrics = scan_metrics.clone(); - async move { - let nd = fast_view_read_to_record_batch(&views, &dataset).await?; - metrics.datasets_scanned.add(1); - // The encoding names the columns and types the scan - // declared. The scan's schema carries each field's - // marks as well, so the batch takes that schema. - let batch = encode_nd_record_batch(&nd)?.with_schema(schema)?; - Ok::<_, DataFusionError>(batch) - } + .map(move |dataset| { + dataset_stream( + Arc::clone(&views), + dataset, + Arc::clone(&projected_schema), + scan_metrics.clone(), + ) }) + .flatten() .boxed(); Ok(stream) }; @@ -168,36 +156,122 @@ pub(crate) async fn column_views( Ok(views) } -/// One dataset of the collection as an nd batch, one column per field. +/// One dataset's batches: one encoded nd batch per stored chunk, in C order. /// -/// A column comes out on the axes the dataset stores it on, and the target -/// grid is their union. A field the dataset lacks is a rank-0 null, which -/// broadcasts to an all-null column: an array no dataset declares, a segment -/// without this dataset's entry, or an attribute nobody set. The decoder makes -/// the same of a null struct row, so the scan sees one thing either way. -async fn fast_view_read_to_record_batch( - views: &IndexMap>, - dataset: &str, -) -> Result { - let mut columns = Vec::with_capacity(views.len()); - for (field, view) in views { - let column = match view { - None => null_scalar(field), - Some(AtlasColumnView::Array { segment }) => match segment.array(dataset) { - Some(info) => array_column(segment, dataset, info, field).await?, - None => null_scalar(field), - }, - Some(AtlasColumnView::GlobalAttribute { map }) - | Some(AtlasColumnView::VariableAttribute { map, .. }) => match map.get(dataset) { - Some(attr) => attr_column(attr, field)?, +/// The dataset is built when the stream reaches it, and each chunk is read +/// when the stream reaches that, so the next dataset is not touched before +/// this one is drained. +fn dataset_stream( + views: Arc>>, + dataset: String, + schema: SchemaRef, + metrics: AtlasScanMetrics, +) -> BoxStream<'static, Result> { + futures::stream::once(async move { + let read = Arc::new(DatasetRead::build(&views, &dataset)?); + metrics.datasets_scanned.add(1); + let chunks = read.dataset.chunks(); + Ok::<_, DataFusionError>(futures::stream::iter(chunks).then(move |chunk| { + let read = Arc::clone(&read); + let schema = Arc::clone(&schema); + async move { + let nd = read.chunk(chunk).await?; + // The encoding names the columns and types the scan declared. + // The scan's schema carries each field's marks as well, so the + // batch takes that schema. + let batch = encode_nd_record_batch(&nd)?.with_schema(schema)?; + Ok::<_, DataFusionError>(batch) + } + })) + }) + .try_flatten() + .boxed() +} + +/// One dataset of the collection as a lazy nd dataset, under the scan's fields. +/// +/// The dataset holds an array for every field it has: the variable's entry in +/// its segment, read on demand through the atlas backend, or an attribute +/// value on no axis. A field it lacks has no array, and reads as a rank-0 +/// null. +struct DatasetRead { + /// The scan's fields, in order. + fields: Vec, + /// The lazy dataset, its arrays keyed by field name. + dataset: Dataset, +} + +impl DatasetRead { + /// No array data is read here. The dataset's layout comes from the + /// segments, and its chunk grid is the one the writer chose. + fn build(views: &IndexMap>, dataset: &str) -> Result { + let mut arrays: IndexMap> = IndexMap::new(); + for (field, view) in views { + let array = match view { + None => None, + Some(AtlasColumnView::Array { segment }) => match segment.array(dataset) { + Some(info) => Some( + compat::array_to_nd_array(Arc::clone(segment), dataset, &info.dtype) + .map_err(execution)?, + ), + None => None, + }, + Some(AtlasColumnView::GlobalAttribute { map }) + | Some(AtlasColumnView::VariableAttribute { map, .. }) => { + // A list has no rank-0 form, and the schema holds no list + // column. A dataset that stores a list under a scalar + // column's key reads as null. + map.get(dataset) + .and_then(|attr| compat::attribute_to_nd_array(attr).ok()) + } + }; + if let Some(array) = array { + arrays.insert(field.name().clone(), array); + } + } + let dataset = Dataset::new(dataset.to_string(), arrays).map_err(execution)?; + Ok(Self { + fields: views.keys().cloned().collect(), + dataset, + }) + } + + /// One chunk of the dataset, under the scan's fields. + /// + /// A column comes out under the array's own type, and the table may + /// declare a wider one: that is a cast. A field the dataset lacks is a + /// rank-0 null, which broadcasts to an all-null column. The decoder makes + /// the same of a null struct row, so the scan sees one thing either way. + async fn chunk(&self, chunk: Arc) -> Result { + let nd = self + .dataset + .poll_next(chunk) + .await + .map_err(execution)? + .ok_or_else(|| { + DataFusionError::Execution(format!( + "dataset '{}' read no batch for a chunk of its own grid", + self.dataset.name + )) + })?; + + let mut columns = Vec::with_capacity(self.fields.len()); + for field in &self.fields { + let column = match nd.schema().column_with_name(field.name()) { + Some((index, _)) => { + let column = nd.column(index); + match as_field_type(Arc::clone(column.values()), field)? { + Some(values) => NdArrowArray::try_new(values, column.dims().clone())?, + None => null_scalar(field), + } + } None => null_scalar(field), - }, - }; - columns.push(column); + }; + columns.push(column); + } + let schema = Arc::new(Schema::new(self.fields.clone())); + NdRecordBatch::try_new(schema, columns, nd.target().clone()) } - let schema = Arc::new(Schema::new(views.keys().cloned().collect::>())); - let target = infer_target(&columns)?; - NdRecordBatch::try_new(schema, columns, target) } /// A rank-0 null. It broadcasts to an all-null column of the target grid. @@ -206,53 +280,6 @@ fn null_scalar(field: &Field) -> NdArrowArray { .expect("one element on no axis") } -/// `dataset`'s entry of one variable, on the axes the segment records for it. -async fn array_column( - segment: &ArrayFile, - dataset: &str, - info: &ArrayInfo, - field: &Field, -) -> Result { - let dims = Dimensions::try_new( - info.dimension_names - .iter() - .zip(&info.shape) - .map(|(name, &size)| Dimension::new(name.as_str(), size as usize)) - .collect(), - )?; - let values = read_values(segment, dataset, info).await?; - match as_field_type(values, field)? { - Some(values) => NdArrowArray::try_new(values, dims), - None => Ok(null_scalar(field)), - } -} - -/// One attribute value of `dataset`, on no axis. -fn attr_column(attr: &Attr, field: &Field) -> Result { - let values: ArrayRef = match attr { - Attr::Bool(v) => Arc::new(BooleanArray::from(vec![*v])), - Attr::Int8(v) => Arc::new(Int8Array::from(vec![*v])), - Attr::Int16(v) => Arc::new(Int16Array::from(vec![*v])), - Attr::Int32(v) => Arc::new(Int32Array::from(vec![*v])), - Attr::Int64(v) => Arc::new(Int64Array::from(vec![*v])), - Attr::UInt8(v) => Arc::new(UInt8Array::from(vec![*v])), - Attr::UInt16(v) => Arc::new(UInt16Array::from(vec![*v])), - Attr::UInt32(v) => Arc::new(UInt32Array::from(vec![*v])), - Attr::UInt64(v) => Arc::new(UInt64Array::from(vec![*v])), - Attr::Float32(v) => Arc::new(Float32Array::from(vec![*v])), - Attr::Float64(v) => Arc::new(Float64Array::from(vec![*v])), - Attr::String(v) => Arc::new(StringArray::from(vec![v.as_str()])), - Attr::Binary(v) => Arc::new(BinaryArray::from(vec![v.as_slice()])), - // A list has no rank-0 form, and the schema holds no list column. A - // dataset that stores a list under a scalar column's key reads as null. - _ => return Ok(null_scalar(field)), - }; - match as_field_type(values, field)? { - Some(values) => NdArrowArray::try_new(values, Dimensions::scalar()), - None => Ok(null_scalar(field)), - } -} - /// `values` in the type the table declares for `field`, or `None` for values /// the table cannot hold. /// @@ -271,124 +298,16 @@ fn as_field_type(values: ArrayRef, field: &Field) -> Result> { } } -/// The flat values of `dataset`'s entry in `segment`, its fill read as null. -async fn read_values(segment: &ArrayFile, dataset: &str, info: &ArrayInfo) -> Result { - let fill = info.fill_value.as_ref(); - match info.dtype { - DType::Int8 => primitive_values::(segment, dataset, fill).await, - DType::Int16 => primitive_values::(segment, dataset, fill).await, - DType::Int32 => primitive_values::(segment, dataset, fill).await, - DType::Int64 => primitive_values::(segment, dataset, fill).await, - DType::UInt8 => primitive_values::(segment, dataset, fill).await, - DType::UInt16 => primitive_values::(segment, dataset, fill).await, - DType::UInt32 => primitive_values::(segment, dataset, fill).await, - DType::UInt64 => primitive_values::(segment, dataset, fill).await, - DType::Float32 => primitive_values::(segment, dataset, fill).await, - DType::Float64 => primitive_values::(segment, dataset, fill).await, - DType::TimestampNs => timestamp_values(segment, dataset, fill).await, - DType::String => text_values(segment, dataset, fill).await, - DType::Binary => binary_values(segment, dataset, fill).await, - DType::Bool | DType::List { .. } | DType::FixedSizeList { .. } => exec_err!( - "dataset '{dataset}' holds a {:?} array, which Beacon does not read", - info.dtype - ), - } -} - -/// The whole entry of `dataset` in `segment`, flat in row-major order. -async fn read_flat(segment: &ArrayFile, dataset: &str) -> Result> { - let values = segment - .read_array::(dataset, vec![], vec![]) - .await - .map_err(external)?; - Ok(values.into_owned().into_raw_vec_and_offset().0) -} - -async fn primitive_values( - segment: &ArrayFile, - dataset: &str, - fill: Option<&FillValue>, -) -> Result -where - A: ArrowPrimitiveType, - A::Native: ArrayElement, -{ - let values = read_flat::(segment, dataset).await?; - let fill = fill.map(|fill| ::fill_element(Some(fill))); - mask_fill( - PrimitiveArray::::new(ScalarBuffer::from(values), None), - fill, - ) -} - -/// Both types are `#[repr(transparent)]` over `i64`. The rename is done -/// element by element all the same, on the type system rather than on layout. -async fn timestamp_values( - segment: &ArrayFile, - dataset: &str, - fill: Option<&FillValue>, -) -> Result { - let values: Vec = read_flat::(segment, dataset) - .await? - .into_iter() - .map(|ts| ts.0) - .collect(); - let fill = fill.map(|fill| ::fill_element(Some(fill)).0); - mask_fill( - PrimitiveArray::::from(values), - fill, - ) -} - -async fn text_values( - segment: &ArrayFile, - dataset: &str, - fill: Option<&FillValue>, -) -> Result { - let values = read_flat::(segment, dataset).await?; - let fill = fill.map(|fill| ::fill_element(Some(fill))); - Ok(Arc::new(StringArray::from_iter(values.into_iter().map( - |value| (fill.as_ref() != Some(&value)).then_some(value), - )))) -} - -async fn binary_values( - segment: &ArrayFile, - dataset: &str, - fill: Option<&FillValue>, -) -> Result { - let values = read_flat::>(segment, dataset).await?; - let fill = fill.map(|fill| as ArrayElement>::fill_element(Some(fill))); - Ok(Arc::new(BinaryArray::from_iter(values.into_iter().map( - |value| (fill.as_ref() != Some(&value)).then_some(value), - )))) -} - -/// `array`, with every element equal to `fill` read as null. -/// -/// One vectorised compare gives the validity. A `NaN` fill masks nothing, as -/// `NaN` equals no value, and the same holds in every other nd format. -fn mask_fill( - array: PrimitiveArray, - fill: Option, -) -> Result { - let Some(fill) = fill else { - return Ok(Arc::new(array)); - }; - let fill = PrimitiveArray::::new_scalar(fill); - let kept = neq(&array, &fill)?; - let nulls = NullBuffer::new(kept.values().clone()); - Ok(Arc::new(PrimitiveArray::::new( - array.values().clone(), - Some(nulls), - ))) -} - /// An atlas error, as the scan reports it. fn external(error: impl std::error::Error + Send + Sync + 'static) -> DataFusionError { DataFusionError::External(Box::new(error)) } +/// A read error, as the scan reports it. +fn execution(error: impl std::fmt::Display) -> DataFusionError { + DataFusionError::Execution(error.to_string()) +} + /// Where one column of the scan comes from, for every dataset of a collection. /// /// The scan reads through it, and pruning judges through it, so both see one @@ -434,13 +353,19 @@ mod tests { column_views(&atlas, &schema).await.unwrap() } - async fn read(dir: &std::path::Path, dataset: &str) -> (NdRecordBatch, RecordBatch) { + /// Every chunk of `dataset`, read as the opener reads it, and the rows of + /// all of them in chunk order. + async fn read(dir: &std::path::Path, dataset: &str) -> (Vec, RecordBatch) { let views = views(dir).await; - let nd = fast_view_read_to_record_batch(&views, dataset) - .await - .unwrap(); - let batch = nd.materialize().unwrap(); - (nd, batch) + let read = DatasetRead::build(&views, dataset).unwrap(); + let mut chunks = Vec::new(); + for chunk in read.dataset.chunks() { + chunks.push(read.chunk(chunk).await.unwrap()); + } + let batches: Vec = chunks.iter().map(|nd| nd.materialize().unwrap()).collect(); + let schema = Arc::new(Schema::new(views.keys().cloned().collect::>())); + let batch = arrow::compute::concat_batches(&schema, &batches).unwrap(); + (chunks, batch) } fn column<'a>(batch: &'a RecordBatch, name: &str) -> &'a ArrayRef { @@ -458,7 +383,8 @@ mod tests { let (nd, batch) = read(tmp.path(), "winter").await; - assert_eq!(nd.target().shape(), vec![4]); + assert_eq!(nd.len(), 1, "an unchunked array is one chunk"); + assert_eq!(nd[0].target().shape(), vec![4]); assert_eq!(batch.num_rows(), 4); assert_eq!( column(&batch, "temperature") @@ -520,7 +446,8 @@ mod tests { ); } - /// A 2-D array keeps both axes, and a cell nobody wrote reads as null. + /// A 2-D array reads one stored chunk at a time, keeps both axes, and a + /// cell nobody wrote reads as null. #[tokio::test] async fn a_fill_value_reads_as_null_on_a_two_dimensional_grid() { let tmp = tempfile::tempdir().unwrap(); @@ -528,14 +455,27 @@ mod tests { let (nd, batch) = read(tmp.path(), "grid").await; - assert_eq!(nd.target().shape(), vec![4, 6]); + assert_eq!(nd.len(), 4, "a [4, 6] grid chunked [2, 3]"); + for chunk in &nd { + assert_eq!(chunk.target().shape(), vec![2, 3]); + } assert_eq!(batch.num_rows(), 24); let temperature = column(&batch, "temperature").as_primitive::(); - assert_eq!(temperature.value(7), 7.0, "row 1, column 1 of a 4x6 grid"); + assert_eq!( + temperature.value(4), + 7.0, + "row 1, column 1 of the grid: the fifth cell of the first chunk" + ); + let mut cells = temperature.values().to_vec(); + cells.sort_by(|a, b| a.partial_cmp(b).unwrap()); + assert_eq!(cells, (0..24).map(f64::from).collect::>()); let sparse = column(&batch, "sparse"); assert_eq!(sparse.null_count(), 12, "two of four rows were written"); - assert!(sparse.is_valid(0)); - assert!(sparse.is_null(23)); + assert!( + sparse.is_valid(0), + "the first chunk lies in the written rows" + ); + assert!(sparse.is_null(23), "the last chunk lies outside them"); } /// `a` stores `value` as `Int16` and `b` as `Float32`. The table declares diff --git a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs new file mode 100644 index 00000000..fccb4ea6 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs @@ -0,0 +1,392 @@ +use std::sync::Arc; + +use beacon_datafusion_ext::nd::NdRecordBatch; +use indexmap::IndexMap; + +use crate::{ + NdArrayD, + array::subset::ArraySubset, + arrow::{ + batch::{build_dataset_schema, generate_chunk_subsets}, + nd_provider::read_nd_chunk, + }, + dataset::source::DatasetSource, +}; + +#[derive(Debug, Clone)] +pub struct Dataset { + pub name: String, + pub dimensions: Vec, + pub shape: Vec, + pub chunk_shape: Vec, + pub arrays: IndexMap>, +} + +impl Dataset { + pub fn new(name: String, arrays: IndexMap>) -> anyhow::Result { + let mut dimensions = Vec::new(); + let mut shape = Vec::new(); + let mut chunk_shape = Vec::new(); + for (_, array) in &arrays { + let array_dims = array.dimensions(); + let array_shape = array.shape(); + let array_chunk_shape = array.chunk_shape(); + + if array_dims.len() > dimensions.len() { + // Array dims should contain all the dimensions of the dataset, in order. + for (i, dim) in dimensions.iter().enumerate() { + let array_dim = array_dims.get(i); + if array_dim != Some(dim) { + return Err(anyhow::anyhow!( + "Array dimensions {:?} has incompatible dimension {:?} at index {}", + array_dims, + dim, + i + )); + } + } + dimensions = array_dims.clone(); + shape = array_shape.clone(); + chunk_shape = array_chunk_shape.clone(); + } else { + // Array dims should be a prefix of the dataset dimensions. + for (i, array_dim) in array_dims.iter().enumerate() { + let dim = dimensions.get(i); + if dim != Some(array_dim) { + return Err(anyhow::anyhow!( + "Array dimensions {:?} has incompatible dimension {:?} at index {}", + array_dims, + array_dim, + i + )); + } + } + } + } + + Ok(Self { + name, + dimensions, + shape, + chunk_shape, + arrays, + }) + } +} + +#[async_trait::async_trait] +impl DatasetSource for Dataset { + /// Every chunk of the dataset grid, in C order, as an [`ArraySubset`]. + /// + /// The cut comes from [`generate_chunk_subsets`], so a boundary chunk + /// shrinks to fit. A scalar dataset has one empty chunk. A dataset with an + /// empty axis has none. + fn chunks(&self) -> Vec> { + generate_chunk_subsets(&self.shape, &self.chunk_shape) + .into_iter() + .map(|subset| Arc::new(subset) as Arc) + .collect() + } + + /// Read one chunk into an un-broadcast [`NdRecordBatch`]. + /// + /// `chunk` is one of [`DatasetSource::chunks`]. Each array is sliced on + /// its own axes, so an array of lower rank reads only the axes it has. The + /// broadcast onto the chunk grid happens above the scan. + async fn poll_next( + &self, + chunk: Arc, + ) -> anyhow::Result> { + let subset = chunk + .downcast_ref::() + .ok_or_else(|| anyhow::anyhow!("chunk is not an ArraySubset"))?; + + let schema = build_dataset_schema(&self.arrays); + let batch = read_nd_chunk(&self.arrays, &self.dimensions, schema, subset.clone()).await?; + Ok(Some(batch)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + NdArray, + array::{backend::ArrayBackend, backend::mem::InMemoryArrayBackend, subset::ArraySubset}, + }; + use ndarray::ArrayD; + + /// An in-memory array that reports the chunk shape a test hands it. + #[derive(Debug)] + struct Chunked { + inner: InMemoryArrayBackend, + chunk_shape: Vec, + } + + #[async_trait::async_trait] + impl ArrayBackend for Chunked { + fn len(&self) -> usize { + self.inner.len() + } + fn shape(&self) -> Vec { + self.inner.shape() + } + fn chunk_shape(&self) -> Vec { + self.chunk_shape.clone() + } + fn dimensions(&self) -> Vec { + self.inner.dimensions() + } + async fn read_subset(&self, subset: ArraySubset) -> anyhow::Result> { + self.inner.read_subset(subset).await + } + } + + fn arr(dims: &[&str], shape: &[usize], chunk: &[usize]) -> Arc { + let len = shape.iter().product(); + let values = ArrayD::from_shape_vec(shape.to_vec(), vec![0.0; len]).unwrap(); + let inner = InMemoryArrayBackend::new( + values, + shape.to_vec(), + dims.iter().map(|d| d.to_string()).collect(), + None, + ); + let backend = Chunked { + inner, + chunk_shape: chunk.to_vec(), + }; + Arc::new(NdArray::new_with_backend(backend).unwrap()) + } + + fn dataset(arrays: Vec<(&str, Arc)>) -> anyhow::Result { + Dataset::new( + "test".to_string(), + arrays + .into_iter() + .map(|(name, array)| (name.to_string(), array)) + .collect(), + ) + } + + fn names(dims: &[&str]) -> Vec { + dims.iter().map(|d| d.to_string()).collect() + } + + #[test] + fn the_array_with_the_most_dimensions_decides() { + let ds = dataset(vec![ + ("time", arr(&["time"], &[4], &[4])), + ("data", arr(&["time", "lat", "lon"], &[4, 6, 8], &[1, 3, 4])), + ("time_lat", arr(&["time", "lat"], &[4, 6], &[4, 6])), + ]) + .unwrap(); + assert_eq!(ds.dimensions, names(&["time", "lat", "lon"])); + assert_eq!(ds.shape, vec![4, 6, 8]); + assert_eq!(ds.chunk_shape, vec![1, 3, 4]); + } + + #[test] + fn the_first_array_wins_a_rank_tie() { + let ds = dataset(vec![ + ("first", arr(&["x", "y"], &[4, 6], &[2, 3])), + ("second", arr(&["x", "y"], &[4, 6], &[4, 6])), + ]) + .unwrap(); + assert_eq!(ds.chunk_shape, vec![2, 3]); + } + + #[test] + fn prefix_arrays_and_scalars_pass() { + let ds = dataset(vec![ + ("data", arr(&["x", "y"], &[4, 6], &[2, 3])), + ("scalar", arr(&[], &[], &[])), + ("x_only", arr(&["x"], &[4], &[4])), + ]) + .unwrap(); + assert_eq!(ds.dimensions, names(&["x", "y"])); + assert_eq!(ds.arrays.len(), 3); + } + + #[test] + fn a_dimension_outside_the_prefix_is_an_error() { + let err = dataset(vec![ + ("data", arr(&["x", "y"], &[4, 6], &[2, 3])), + ("bounds", arr(&["y", "nv"], &[6, 2], &[6, 2])), + ]) + .unwrap_err(); + assert!(err.to_string().contains("\"nv\""), "{err}"); + } + + #[test] + fn an_anchor_that_does_not_extend_the_prefix_is_an_error() { + let err = dataset(vec![ + ("x_only", arr(&["x"], &[4], &[4])), + ("data", arr(&["y", "x"], &[6, 4], &[6, 4])), + ]) + .unwrap_err(); + assert!(err.to_string().contains("\"x\""), "{err}"); + } + + #[test] + fn chunks_cover_the_grid_in_c_order() { + let ds = dataset(vec![("data", arr(&["x", "y"], &[5, 4], &[2, 4]))]).unwrap(); + let chunks: Vec = ds + .chunks() + .into_iter() + .map(|chunk| chunk.downcast_ref::().unwrap().clone()) + .collect(); + let starts: Vec> = chunks.iter().map(|c| c.start.clone()).collect(); + let shapes: Vec> = chunks.iter().map(|c| c.shape.clone()).collect(); + assert_eq!(starts, vec![vec![0, 0], vec![2, 0], vec![4, 0]]); + assert_eq!(shapes, vec![vec![2, 4], vec![2, 4], vec![1, 4]]); + } + + #[test] + fn a_scalar_dataset_has_one_empty_chunk() { + let ds = dataset(vec![("scalar", arr(&[], &[], &[]))]).unwrap(); + let chunks = ds.chunks(); + assert_eq!(chunks.len(), 1); + let subset = chunks[0].downcast_ref::().unwrap(); + assert!(subset.start.is_empty()); + assert!(subset.shape.is_empty()); + } + + #[test] + fn an_empty_axis_has_no_chunks() { + let ds = dataset(vec![("data", arr(&["x", "y"], &[0, 4], &[2, 4]))]).unwrap(); + assert!(ds.chunks().is_empty()); + } + + /// A dataset with a coordinate per axis and one 2-D variable, cut on the + /// variable's chunk shape. + fn gridded(chunk: &[usize]) -> Dataset { + let time = NdArray::::try_new_from_vec_in_mem( + (0..4).map(|v| v * 100).collect(), + vec![4], + names(&["time"]), + None, + ) + .unwrap(); + let lat = NdArray::::try_new_from_vec_in_mem( + vec![-30.0, 0.0, 30.0], + vec![3], + names(&["lat"]), + None, + ) + .unwrap(); + let sst = { + let values = + ArrayD::from_shape_vec(vec![4, 3], (0..12).map(|v| v as f64).collect()).unwrap(); + let inner = + InMemoryArrayBackend::new(values, vec![4, 3], names(&["time", "lat"]), None); + NdArray::new_with_backend(Chunked { + inner, + chunk_shape: chunk.to_vec(), + }) + .unwrap() + }; + // Built directly: `new` applies the prefix rule, and `lat(lat)` is not + // a prefix of `sst(time, lat)`. The read itself maps every array onto + // the grid by dimension name. + let mut arrays: IndexMap> = IndexMap::new(); + arrays.insert("time".to_string(), Arc::new(time)); + arrays.insert("lat".to_string(), Arc::new(lat)); + arrays.insert("sst".to_string(), Arc::new(sst)); + Dataset { + name: "gridded".to_string(), + dimensions: names(&["time", "lat"]), + shape: vec![4, 3], + chunk_shape: chunk.to_vec(), + arrays, + } + } + + /// Poll every chunk, broadcast each, and stitch the rows back together. + async fn read_all(ds: &Dataset) -> arrow::record_batch::RecordBatch { + let mut batches = Vec::new(); + for chunk in ds.chunks() { + let nd = ds.poll_next(chunk).await.unwrap().unwrap(); + batches.push(nd.materialize().unwrap()); + } + let schema = build_dataset_schema(&ds.arrays); + arrow::compute::concat_batches(&schema, &batches).unwrap() + } + + /// The chunked read equals the whole-array read, row for row. + #[tokio::test] + async fn polling_every_chunk_reads_the_whole_grid() { + use arrow::array::{Float64Array, Int64Array}; + + let whole = read_all(&gridded(&[4, 3])).await; + assert_eq!(whole.num_rows(), 12); + + let sst = whole.column_by_name("sst").unwrap(); + let sst = sst.as_any().downcast_ref::().unwrap(); + assert_eq!( + sst.values(), + &(0..12).map(|v| v as f64).collect::>()[..] + ); + + let time = whole.column_by_name("time").unwrap(); + let time = time.as_any().downcast_ref::().unwrap(); + assert_eq!( + time.values(), + &[0, 0, 0, 100, 100, 100, 200, 200, 200, 300, 300, 300] + ); + + // A cut on the outer axis alone keeps the row order. + assert_eq!(read_all(&gridded(&[2, 3])).await, whole); + + // A cut on an inner axis reorders the rows, and loses none. + for chunk in [[3, 2], [1, 1], [2, 2]] { + let ds = gridded(&chunk); + assert!(ds.chunks().len() > 1, "chunk {chunk:?} must cut the grid"); + let mut chunked = rows(&read_all(&ds).await); + chunked.sort(); + let mut expected = rows(&whole); + expected.sort(); + assert_eq!(chunked, expected, "chunk {chunk:?}"); + } + } + + /// The `(time, lat, sst)` of every row, as integers so they sort. + fn rows(batch: &arrow::record_batch::RecordBatch) -> Vec<(i64, i64, i64)> { + use arrow::array::{Float64Array, Int64Array}; + let time = batch.column_by_name("time").unwrap(); + let time = time.as_any().downcast_ref::().unwrap(); + let lat = batch.column_by_name("lat").unwrap(); + let lat = lat.as_any().downcast_ref::().unwrap(); + let sst = batch.column_by_name("sst").unwrap(); + let sst = sst.as_any().downcast_ref::().unwrap(); + (0..batch.num_rows()) + .map(|i| (time.value(i), lat.value(i) as i64, sst.value(i) as i64)) + .collect() + } + + /// A chunk leaves un-broadcast: a coordinate keeps its own rank. + #[tokio::test] + async fn a_polled_chunk_is_not_broadcast() { + let ds = gridded(&[2, 3]); + let chunk = ds.chunks().into_iter().next().unwrap(); + let nd = ds.poll_next(chunk).await.unwrap().unwrap(); + assert_eq!(nd.target().rank(), 2); + assert_eq!(nd.num_rows(), 6); + assert_eq!(nd.column(0).dims().rank(), 1, "time stays 1-D"); + assert_eq!(nd.column(2).dims().rank(), 2, "sst is 2-D"); + } + + #[tokio::test] + async fn a_chunk_of_another_type_is_an_error() { + let ds = gridded(&[4, 3]); + let err = ds.poll_next(Arc::new(42usize)).await.unwrap_err(); + assert!(err.to_string().contains("not an ArraySubset"), "{err}"); + } + + #[test] + fn no_arrays_gives_no_dimensions() { + let ds = dataset(vec![]).unwrap(); + assert!(ds.dimensions.is_empty()); + assert!(ds.shape.is_empty()); + assert!(ds.chunk_shape.is_empty()); + } +} diff --git a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/mod.rs b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/mod.rs index d2bc4af3..f508a410 100644 --- a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/mod.rs @@ -1,5 +1,7 @@ pub mod any; +pub mod default; pub mod ragged; +pub mod source; pub mod variant; pub use any::AnyDataset; @@ -104,7 +106,11 @@ impl Dataset { let dimensions = self .dimensions .iter() - .filter(|(dim, _)| arrays.values().any(|array| array.dimensions().contains(dim))) + .filter(|(dim, _)| { + arrays + .values() + .any(|array| array.dimensions().contains(dim)) + }) .map(|(dim, size)| (dim.clone(), *size)) .collect(); @@ -164,12 +170,10 @@ impl Dataset { .map(|array| array.dimensions())?; // Already broadcast-safe: every variable's dims fit inside `max_dims`. - let needs_narrowing = self.arrays.values().any(|array| { - !array - .dimensions() - .iter() - .all(|dim| max_dims.contains(dim)) - }); + let needs_narrowing = self + .arrays + .values() + .any(|array| !array.dimensions().iter().all(|dim| max_dims.contains(dim))); if !needs_narrowing { return None; } @@ -1229,8 +1233,13 @@ mod tests { let len: usize = shape.iter().product(); let dim_names: Vec = dims.iter().map(|d| d.to_string()).collect(); Arc::new( - NdArray::::try_new_from_vec_in_mem(vec![0.0; len], shape.to_vec(), dim_names, None) - .unwrap(), + NdArray::::try_new_from_vec_in_mem( + vec![0.0; len], + shape.to_vec(), + dim_names, + None, + ) + .unwrap(), ) } @@ -1239,7 +1248,10 @@ mod tests { // 2D var plus a 1D subset — already broadcast-safe. let ds = make_dataset( "safe", - vec![("grid", arr(&["x", "y"]).await), ("scale", arr(&["y"]).await)], + vec![ + ("grid", arr(&["x", "y"]).await), + ("scale", arr(&["y"]).await), + ], ) .await; assert_eq!(ds.default_broadcast_dimensions(), None); @@ -1295,7 +1307,10 @@ mod tests { #[tokio::test] async fn test_default_dims_of_an_unnamed_file_follow_the_volume() { let arrays = vec![ - ("data", arr_shaped(&["len_1250", "len_23250"], &[1250, 23250]).await), + ( + "data", + arr_shaped(&["len_1250", "len_23250"], &[1250, 23250]).await, + ), ("meta_a", arr_shaped(&["len_3"], &[3]).await), ("meta_b", arr_shaped(&["len_3"], &[3]).await), ("meta_c", arr_shaped(&["len_3"], &[3]).await), @@ -1311,9 +1326,7 @@ mod tests { // Invented: the volume wins, and the payload survives. let unnamed = make_dataset("unnamed", arrays) .await - .with_invented_dimensions( - ["len_1250", "len_23250", "len_3"].map(String::from), - ); + .with_invented_dimensions(["len_1250", "len_23250", "len_3"].map(String::from)); assert_eq!( unnamed.default_broadcast_dimensions(), Some(vec!["len_1250".to_string(), "len_23250".to_string()]) @@ -1328,7 +1341,10 @@ mod tests { let ds = make_dataset( "mixed", vec![ - ("data", arr_shaped(&["time", "len_23250"], &[1250, 23250]).await), + ( + "data", + arr_shaped(&["time", "len_23250"], &[1250, 23250]).await, + ), ("meta_a", arr_shaped(&["len_3"], &[3]).await), ("meta_b", arr_shaped(&["len_3"], &[3]).await), ("meta_c", arr_shaped(&["len_3"], &[3]).await), @@ -1454,7 +1470,10 @@ mod tests { // Equal variable count and equal dimensionality → first-encountered wins. let ds = make_dataset( "fulltie", - vec![("first", arr(&["a", "b"]).await), ("second", arr(&["c", "d"]).await)], + vec![ + ("first", arr(&["a", "b"]).await), + ("second", arr(&["c", "d"]).await), + ], ) .await; assert_eq!( @@ -1598,10 +1617,7 @@ mod tests { let resolved = resolve_read_dimensions(&any, Some(vec!["y".to_string(), "nv".to_string()]), None); - assert_eq!( - resolved, - Some(vec!["y".to_string(), "nv".to_string()]) - ); + assert_eq!(resolved, Some(vec!["y".to_string(), "nv".to_string()])); } #[tokio::test] @@ -1618,7 +1634,10 @@ mod tests { .await; let any = AnyDataset::try_from_dataset(ds).await.unwrap(); - assert_eq!(resolve_read_dimensions(&any, Some(vec![]), None), Some(vec![])); + assert_eq!( + resolve_read_dimensions(&any, Some(vec![]), None), + Some(vec![]) + ); } #[tokio::test] @@ -1643,7 +1662,10 @@ mod tests { async fn test_resolve_read_dimensions_none_when_already_broadcast_safe() { let ds = make_dataset( "safe", - vec![("grid", arr(&["x", "y"]).await), ("scale", arr(&["y"]).await)], + vec![ + ("grid", arr(&["x", "y"]).await), + ("scale", arr(&["y"]).await), + ], ) .await; let any = AnyDataset::try_from_dataset(ds).await.unwrap(); diff --git a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/source.rs b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/source.rs new file mode 100644 index 00000000..99012612 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/source.rs @@ -0,0 +1,72 @@ +use beacon_datafusion_ext::nd::NdRecordBatch; +use crossbeam::queue::ArrayQueue; +use std::{any::Any, sync::Arc}; + +#[async_trait::async_trait] +pub trait DatasetSource: Send + Sync + std::fmt::Debug { + fn chunks(&self) -> Vec>; + + async fn poll_next( + &self, + chunk: Arc, + ) -> anyhow::Result>; +} + +/// A source shared by the partitions that read it, and the steps left to read. +/// +/// The queue is filled once, on construction. Every partition pops from the +/// same queue, so each step is read by one partition and by no other. +#[derive(Debug, Clone)] +pub struct SharedDatasetSource { + source: Arc, + chunks: Arc>>, +} + +impl SharedDatasetSource { + /// Cut `source` into steps and queue them. + /// + /// A step is one slab along the outermost axis: every chunk that shares + /// one chunk index on axis 0. The chunks of a step are in C order, and the + /// steps are queued in C order too, so a reader that drains the queue alone + /// sees the chunks in the order the source stores them. + /// + /// See [`chunk_steps`] for the cut. + pub fn new(source: Arc) -> anyhow::Result { + // Cut the source into steps, and queue them. The queue is shared by all + let steps = source.chunks(); + if steps.is_empty() { + return Err(anyhow::anyhow!("dataset source has no chunks")); + } + if steps.len() == 1 { + return Err(anyhow::anyhow!( + "dataset source has only one chunk, so no steps" + )); + } + let queue = ArrayQueue::new(steps.len()); + for step in steps { + queue + .push(step) + .map_err(|_| anyhow::anyhow!("dataset source queue is full, cannot push step"))?; + } + Ok(Self { + source, + chunks: Arc::new(queue), + }) + } + + /// How many steps are left. For tests and diagnostics. + pub fn remaining_steps(&self) -> usize { + self.chunks.len() + } + + pub fn next_step(&self) -> Option> { + self.chunks.pop() + } + + pub async fn poll_next( + &self, + chunk: Arc, + ) -> anyhow::Result> { + self.source.poll_next(chunk).await + } +} From ae86a923e44ed0ce8df20df7d5a321f47d453187 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 9 Sep 2026 01:26:54 +0200 Subject: [PATCH 10/16] wip --- beacon-db/beacon-core/src/runtime_builder.rs | 8 +++- .../beacon-core/tests/wide_scan_planning.rs | 48 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) create mode 100644 beacon-db/beacon-core/tests/wide_scan_planning.rs diff --git a/beacon-db/beacon-core/src/runtime_builder.rs b/beacon-db/beacon-core/src/runtime_builder.rs index 7a33a004..9279245b 100644 --- a/beacon-db/beacon-core/src/runtime_builder.rs +++ b/beacon-db/beacon-core/src/runtime_builder.rs @@ -43,7 +43,7 @@ use datafusion::{ runtime_env::{RuntimeEnv, RuntimeEnvBuilder}, SessionStateBuilder, }, - optimizer::OptimizerRule, + optimizer::{optimize_projections::OptimizeProjections, OptimizerRule}, prelude::{SessionConfig, SessionContext}, }; use object_store::ObjectStore; @@ -853,7 +853,11 @@ fn build_session_state( runtime_env: Arc, session_cell: SessionCell, ) -> anyhow::Result { - let mut optimizer_rules: Vec> = vec![]; + // Narrow every scan to the columns the query reads before any other rule. + // CSE copies every input column into an intermediate projection, one linear + // schema lookup per column, which is quadratic on a scan of 100k+ columns. + let mut optimizer_rules: Vec> = + vec![Arc::new(OptimizeProjections::new())]; // This is DataFusion's default logical rule set with `FederationOptimizerRule` // inserted, so replacing the defaults with it is intentional: sub-plans rooted // at remote tables get pushed down. The matching `FederatedPlanner` lives in diff --git a/beacon-db/beacon-core/tests/wide_scan_planning.rs b/beacon-db/beacon-core/tests/wide_scan_planning.rs new file mode 100644 index 00000000..20978e8d --- /dev/null +++ b/beacon-db/beacon-core/tests/wide_scan_planning.rs @@ -0,0 +1,48 @@ +//! Planning over a very wide table. +//! +//! An Atlas table can declare 100k+ columns. DataFusion's common-subexpression +//! rule copies every input column into an intermediate projection, one linear +//! schema lookup per column, so the scan must be narrowed before that rule +//! runs. The runtime puts `OptimizeProjections` first in its rule list for +//! that reason, and this test pins the order: the query below repeats `abs(c0)` +//! after simplification and has to plan in well under a minute. + +mod common; + +use std::sync::Arc; +use std::time::Duration; + +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use common::{runtime, total_rows}; +use datafusion::parquet::arrow::ArrowWriter; + +/// Wide enough that a quadratic pass takes minutes, narrow enough to write fast. +const COLUMNS: usize = 80_000; + +#[tokio::test(flavor = "multi_thread")] +async fn a_repeated_sub_expression_over_a_wide_scan_plans_quickly() { + let rt = runtime("wide-scan").await; + + let fields: Vec = (0..COLUMNS) + .map(|i| Field::new(format!("c{i}"), DataType::Float64, true)) + .collect(); + let schema = Arc::new(Schema::new(fields)); + let path = rt.datasets_dir().join("wide/w.parquet"); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + let file = std::fs::File::create(&path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema.clone(), None).unwrap(); + writer.write(&RecordBatch::new_empty(schema)).unwrap(); + writer.close().unwrap(); + + rt.sql("CREATE EXTERNAL TABLE wide STORED AS PARQUET LOCATION 'wide/'") + .await; + + // `coalesce(f(x), y)` simplifies to a CASE that names `f(x)` twice. + let query = rt.try_sql("SELECT coalesce(abs(c0), c1) AS d FROM wide LIMIT 5"); + let batches = tokio::time::timeout(Duration::from_secs(60), query) + .await + .expect("planning must not be quadratic in the column count") + .expect("the query runs"); + assert_eq!(total_rows(&batches), 0); +} From b5faf2462641820b861acae93e9752a94e61ea80 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 9 Sep 2026 13:02:56 +0200 Subject: [PATCH 11/16] wip --- Cargo.lock | 2 + .../beacon-arrow-atlas/Cargo.toml | 2 + .../beacon-arrow-atlas/src/datafusion/mod.rs | 2 + .../src/datafusion/opener.rs | 258 ++++------- .../beacon-arrow-atlas/src/datafusion/pool.rs | 399 ++++++++++++++++++ .../src/datafusion/source.rs | 22 +- .../beacon-arrow-atlas/src/datafusion/view.rs | 104 +++++ .../beacon-nd-array/src/arrow/batch.rs | 4 +- .../beacon-nd-array/src/dataset/default.rs | 16 +- .../beacon-nd-array/src/dataset/ragged.rs | 269 +++++++++++- 10 files changed, 876 insertions(+), 202 deletions(-) create mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs create mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs diff --git a/Cargo.lock b/Cargo.lock index 39a12c1b..9cbdd5c4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1452,12 +1452,14 @@ dependencies = [ "beacon-datafusion-ext", "beacon-nd-array", "chrono", + "crossbeam", "datafusion 53.1.0", "futures", "indexmap 2.14.0", "moka", "ndarray 0.17.2", "object_store 0.13.2", + "parking_lot", "tempfile", "tokio", "tracing", diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/Cargo.toml b/beacon-db/beacon-file-formats/beacon-arrow-atlas/Cargo.toml index b514f673..fad3dbde 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/Cargo.toml +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/Cargo.toml @@ -22,6 +22,8 @@ tokio = { workspace = true } ndarray = { workspace = true } chrono = { workspace = true } moka = { workspace = true } +crossbeam = { workspace = true } +parking_lot = { workspace = true } beacon-datafusion-ext = { path = "../../beacon-datafusion-ext" } beacon-nd-array = { path = "../beacon-nd-array" } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs index adc14497..54b2aaf1 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs @@ -40,9 +40,11 @@ use crate::store::{ATLAS_MARKER, AtlasReaderCache, get_or_open_atlas, top_level_ pub mod metrics; pub mod opener; pub mod options; +pub mod pool; pub mod pruning; pub mod source; pub mod table_function; +pub mod view; pub use options::AtlasOptions; pub use source::AtlasSource; diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs index 257c717b..02cec72b 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs @@ -1,26 +1,23 @@ //! One partition's opener: a collection in, nd batches out. //! -//! A column view says where one column of the scan comes from, for every -//! dataset at once: the variable's segment, or the attribute values keyed by -//! dataset. One dataset is then a lazy [`Dataset`] over those views. It reads -//! one stored chunk at a time as an [`NdRecordBatch`], each column on the axes -//! the dataset stores it on. +//! The opener reads through the [`AtlasReaderPool`]. The first partition to +//! reach a collection opens it and queues its datasets, and every partition +//! then streams the datasets it pops. What lives here is the column +//! resolution the read and the pruning share: a column view says where one +//! column of the scan comes from, for every dataset at once, and +//! [`under_fields`] puts one chunk of a dataset under the scan's fields. -use std::{any::Any, sync::Arc}; +use std::sync::Arc; use arrow::{ - array::{ArrayRef, RecordBatch, new_null_array}, + array::{ArrayRef, new_null_array}, compute::cast, datatypes::{Field, FieldRef, Schema, SchemaRef}, }; use atlas::{ArrayFile, Atlas, Attr}; -use beacon_datafusion_ext::nd::{Dimensions, NdArrowArray, NdRecordBatch, encode_nd_record_batch}; +use beacon_datafusion_ext::nd::{Dimensions, NdArrowArray, NdRecordBatch}; use beacon_datafusion_ext::type_widening::is_type_conflict; -use beacon_nd_array::{ - NdArrayD, - arrow::metrics::ReadMetrics, - dataset::{default::Dataset, source::DatasetSource}, -}; +use beacon_nd_array::arrow::metrics::ReadMetrics; use datafusion::{ datasource::{ listing::PartitionedFile, @@ -29,14 +26,13 @@ use datafusion::{ error::{DataFusionError, Result}, physical_plan::PhysicalExpr, }; -use futures::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream}; +use futures::{FutureExt, StreamExt, TryStreamExt}; use indexmap::IndexMap; use object_store::ObjectStore; use crate::{ - compat, - datafusion::{metrics::AtlasScanMetrics, pruning::prune_datasets}, - store::{AtlasReaderCache, get_or_open_atlas}, + datafusion::{metrics::AtlasScanMetrics, pool::AtlasReaderPool}, + store::AtlasReaderCache, }; /// One partition's opener: a collection in, its batches out. @@ -58,15 +54,20 @@ pub struct AtlasOpener { pub predicate: Option>, pub read_metrics: ReadMetrics, pub scan_metrics: AtlasScanMetrics, + /// The scan's pools, one per collection, shared by every partition. + pub reader_pool: Arc, } impl FileOpener for AtlasOpener { /// One collection in, one encoded batch per stored chunk of every dataset /// worth reading out. /// - /// The column views are built once per collection, and every dataset then - /// reads against them. A dataset the deletion mask hides is not read, and - /// neither is one the predicate rules out from the statistics in memory. + /// The collection is opened through the reader pool. The first partition + /// to reach it opens it, prunes its datasets in one pass over the footer's + /// statistics, and queues the survivors. A dataset the deletion mask hides + /// is not queued, and neither is one the predicate rules out. Every + /// partition then streams the datasets it pops off that queue, so the + /// partitions that share a collection share its work. fn open(&self, file: PartitionedFile) -> Result { let store = self.object_store.clone(); let cache = self.cache.clone(); @@ -74,41 +75,29 @@ impl FileOpener for AtlasOpener { let logical_schema = self.logical_schema.clone(); let predicate = self.predicate.clone(); let scan_metrics = self.scan_metrics.clone(); + let pool = Arc::clone(&self.reader_pool); let fut = async move { - let open_timer = scan_metrics.open_time.timer(); - let atlas = get_or_open_atlas(Some(&cache), store, &file.object_meta) + let location = file.object_meta.location.clone(); + let stream = pool + .try_open_into_pooled_stream( + Some(&cache), + store, + file.object_meta, + logical_schema, + projected_schema, + predicate, + scan_metrics, + ) .await .map_err(|e| { DataFusionError::Execution(format!( - "Failed to open atlas collection '{}': {e}", - file.object_meta.location + "Failed to open atlas collection '{location}': {e}" )) })?; - let views = Arc::new(column_views(&atlas, &logical_schema).await?); - drop(open_timer); - - let mut datasets = atlas.list_datasets(); - if let Some(predicate) = &predicate { - let prune_timer = scan_metrics.prune_time.timer(); - let listed = datasets.len(); - datasets = prune_datasets(&views, datasets, predicate, &logical_schema).await; - scan_metrics.datasets_pruned.add(listed - datasets.len()); - drop(prune_timer); - } - - let stream = futures::stream::iter(datasets) - .map(move |dataset| { - dataset_stream( - Arc::clone(&views), - dataset, - Arc::clone(&projected_schema), - scan_metrics.clone(), - ) - }) - .flatten() - .boxed(); - Ok(stream) + Ok(stream + .map_err(|e| DataFusionError::External(e.into())) + .boxed()) }; Ok(fut.boxed()) @@ -156,122 +145,29 @@ pub(crate) async fn column_views( Ok(views) } -/// One dataset's batches: one encoded nd batch per stored chunk, in C order. -/// -/// The dataset is built when the stream reaches it, and each chunk is read -/// when the stream reaches that, so the next dataset is not touched before -/// this one is drained. -fn dataset_stream( - views: Arc>>, - dataset: String, - schema: SchemaRef, - metrics: AtlasScanMetrics, -) -> BoxStream<'static, Result> { - futures::stream::once(async move { - let read = Arc::new(DatasetRead::build(&views, &dataset)?); - metrics.datasets_scanned.add(1); - let chunks = read.dataset.chunks(); - Ok::<_, DataFusionError>(futures::stream::iter(chunks).then(move |chunk| { - let read = Arc::clone(&read); - let schema = Arc::clone(&schema); - async move { - let nd = read.chunk(chunk).await?; - // The encoding names the columns and types the scan declared. - // The scan's schema carries each field's marks as well, so the - // batch takes that schema. - let batch = encode_nd_record_batch(&nd)?.with_schema(schema)?; - Ok::<_, DataFusionError>(batch) - } - })) - }) - .try_flatten() - .boxed() -} - -/// One dataset of the collection as a lazy nd dataset, under the scan's fields. +/// `nd` under `fields`: every field in order, on the same target grid. /// -/// The dataset holds an array for every field it has: the variable's entry in -/// its segment, read on demand through the atlas backend, or an attribute -/// value on no axis. A field it lacks has no array, and reads as a rank-0 -/// null. -struct DatasetRead { - /// The scan's fields, in order. - fields: Vec, - /// The lazy dataset, its arrays keyed by field name. - dataset: Dataset, -} - -impl DatasetRead { - /// No array data is read here. The dataset's layout comes from the - /// segments, and its chunk grid is the one the writer chose. - fn build(views: &IndexMap>, dataset: &str) -> Result { - let mut arrays: IndexMap> = IndexMap::new(); - for (field, view) in views { - let array = match view { - None => None, - Some(AtlasColumnView::Array { segment }) => match segment.array(dataset) { - Some(info) => Some( - compat::array_to_nd_array(Arc::clone(segment), dataset, &info.dtype) - .map_err(execution)?, - ), - None => None, - }, - Some(AtlasColumnView::GlobalAttribute { map }) - | Some(AtlasColumnView::VariableAttribute { map, .. }) => { - // A list has no rank-0 form, and the schema holds no list - // column. A dataset that stores a list under a scalar - // column's key reads as null. - map.get(dataset) - .and_then(|attr| compat::attribute_to_nd_array(attr).ok()) +/// A column comes out under the array's own type, and the table may declare a +/// wider one: that is a cast. A field the dataset lacks is a rank-0 null, +/// which broadcasts to an all-null column. The decoder makes the same of a +/// null struct row, so the scan sees one thing either way. +pub(crate) fn under_fields(nd: &NdRecordBatch, fields: &[FieldRef]) -> Result { + let mut columns = Vec::with_capacity(fields.len()); + for field in fields { + let column = match nd.schema().column_with_name(field.name()) { + Some((index, _)) => { + let column = nd.column(index); + match as_field_type(Arc::clone(column.values()), field)? { + Some(values) => NdArrowArray::try_new(values, column.dims().clone())?, + None => null_scalar(field), } - }; - if let Some(array) = array { - arrays.insert(field.name().clone(), array); } - } - let dataset = Dataset::new(dataset.to_string(), arrays).map_err(execution)?; - Ok(Self { - fields: views.keys().cloned().collect(), - dataset, - }) - } - - /// One chunk of the dataset, under the scan's fields. - /// - /// A column comes out under the array's own type, and the table may - /// declare a wider one: that is a cast. A field the dataset lacks is a - /// rank-0 null, which broadcasts to an all-null column. The decoder makes - /// the same of a null struct row, so the scan sees one thing either way. - async fn chunk(&self, chunk: Arc) -> Result { - let nd = self - .dataset - .poll_next(chunk) - .await - .map_err(execution)? - .ok_or_else(|| { - DataFusionError::Execution(format!( - "dataset '{}' read no batch for a chunk of its own grid", - self.dataset.name - )) - })?; - - let mut columns = Vec::with_capacity(self.fields.len()); - for field in &self.fields { - let column = match nd.schema().column_with_name(field.name()) { - Some((index, _)) => { - let column = nd.column(index); - match as_field_type(Arc::clone(column.values()), field)? { - Some(values) => NdArrowArray::try_new(values, column.dims().clone())?, - None => null_scalar(field), - } - } - None => null_scalar(field), - }; - columns.push(column); - } - let schema = Arc::new(Schema::new(self.fields.clone())); - NdRecordBatch::try_new(schema, columns, nd.target().clone()) + None => null_scalar(field), + }; + columns.push(column); } + let schema = Arc::new(Schema::new(fields.to_vec())); + NdRecordBatch::try_new(schema, columns, nd.target().clone()) } /// A rank-0 null. It broadcasts to an all-null column of the target grid. @@ -303,11 +199,6 @@ fn external(error: impl std::error::Error + Send + Sync + 'static) -> DataFusion DataFusionError::External(Box::new(error)) } -/// A read error, as the scan reports it. -fn execution(error: impl std::fmt::Display) -> DataFusionError { - DataFusionError::Execution(error.to_string()) -} - /// Where one column of the scan comes from, for every dataset of a collection. /// /// The scan reads through it, and pruning judges through it, so both see one @@ -332,6 +223,7 @@ mod tests { use beacon_datafusion_ext::type_widening::ArrowTypeWidening; use super::*; + use crate::datafusion::view::AtlasView; use crate::{compat, test_support}; use std::path::Path; @@ -342,28 +234,33 @@ mod tests { use datafusion::scalar::ScalarValue; use futures::TryStreamExt; - /// The column views of a fixture, over the schema `infer_schema` derives. - async fn views(dir: &std::path::Path) -> IndexMap> { + /// The schema `infer_schema` derives for a fixture. + async fn schema(dir: &Path) -> SchemaRef { let atlas = test_support::open(dir).await; - let schema = compat::collection_arrow_schema( - &atlas.footer().collection_schema(), - &ArrowTypeWidening::default_extension(), + Arc::new( + compat::collection_arrow_schema( + &atlas.footer().collection_schema(), + &ArrowTypeWidening::default_extension(), + ) + .unwrap(), ) - .unwrap(); - column_views(&atlas, &schema).await.unwrap() } - /// Every chunk of `dataset`, read as the opener reads it, and the rows of - /// all of them in chunk order. - async fn read(dir: &std::path::Path, dataset: &str) -> (Vec, RecordBatch) { - let views = views(dir).await; - let read = DatasetRead::build(&views, dataset).unwrap(); + /// Every chunk of `dataset`, read as the scan reads it: through the view, + /// under the scan's fields. And the rows of all of them in chunk order. + async fn read(dir: &Path, dataset: &str) -> (Vec, RecordBatch) { + let schema = schema(dir).await; + let (store, marker) = test_support::store_and_marker(dir); + let view = AtlasView::new(None, store, marker, Arc::clone(&schema)) + .await + .unwrap(); + let source = view.dataset(dataset).await.unwrap().unwrap(); let mut chunks = Vec::new(); - for chunk in read.dataset.chunks() { - chunks.push(read.chunk(chunk).await.unwrap()); + for chunk in source.chunks() { + let nd = source.poll_next(chunk).await.unwrap().unwrap(); + chunks.push(under_fields(&nd, schema.fields()).unwrap()); } let batches: Vec = chunks.iter().map(|nd| nd.materialize().unwrap()).collect(); - let schema = Arc::new(Schema::new(views.keys().cloned().collect::>())); let batch = arrow::compute::concat_batches(&schema, &batches).unwrap(); (chunks, batch) } @@ -537,6 +434,7 @@ mod tests { predicate: None, read_metrics: ReadMetrics::new(&metrics, 0), scan_metrics: AtlasScanMetrics::new(&metrics, 0), + reader_pool: Arc::new(AtlasReaderPool::new()), }; (opener, PartitionedFile::from(marker)) } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs new file mode 100644 index 00000000..f1c85a30 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs @@ -0,0 +1,399 @@ +//! A shared reader over one collection: the datasets left to read, and the +//! view every partition reads them through. +//! +//! The pool for a collection is built once, on the first open, and every +//! later open of the same collection gets a handle to it. A handle is a +//! stream. It pops the next dataset off the shared queue, reads it one stored +//! chunk at a time, and yields each chunk as one nd-encoded batch. Two handles +//! that poll at once drain different datasets, so a dataset is read by one +//! partition and by no other. + +use std::collections::HashMap; +use std::fmt::Debug; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::array::RecordBatch; +use arrow::datatypes::SchemaRef; +use beacon_datafusion_ext::nd::encode_nd_record_batch; +use crossbeam::queue::ArrayQueue; +use datafusion::physical_plan::PhysicalExpr; +use futures::stream::BoxStream; +use futures::{Stream, StreamExt, TryStreamExt}; +use object_store::{ObjectMeta, ObjectStore}; +use parking_lot::RwLock; +use tokio::sync::OnceCell; + +use crate::datafusion::metrics::AtlasScanMetrics; +use crate::datafusion::opener::under_fields; +use crate::datafusion::view::AtlasView; +use crate::store::AtlasReaderCache; + +/// The pools of one scan, one per collection, keyed by the container's path. +#[derive(Default, Clone)] +pub struct AtlasReaderPool { + pools: Arc>>>>>, +} + +impl Debug for AtlasReaderPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AtlasReaderPool").finish() + } +} + +impl AtlasReaderPool { + pub fn new() -> Self { + Self { + pools: Arc::new(RwLock::new(HashMap::new())), + } + } + + /// One consumer of the collection at `object_meta`. + /// + /// The first call for a collection opens it, through `cache` when given, + /// prunes its datasets with `pruning_predicate`, and queues the survivors. + /// Every call gets a stream over that queue. The open and the prune are + /// timed on the first caller's `scan_metrics`, and each consumer counts + /// the datasets it reads on its own. + /// + /// `logical_schema` is what the columns and the predicate are written + /// against. `projected_schema` is the same schema nd-encoded, and every + /// batch goes out under it. + #[allow(clippy::too_many_arguments)] + pub async fn try_open_into_pooled_stream( + &self, + cache: Option<&AtlasReaderCache>, + store: Arc, + object_meta: ObjectMeta, + logical_schema: SchemaRef, + projected_schema: SchemaRef, + pruning_predicate: Option>, + scan_metrics: AtlasScanMetrics, + ) -> anyhow::Result>> { + let cell = if self.pools.read().contains_key(&object_meta.location) { + self.pools + .read() + .get(&object_meta.location) + .unwrap() + .clone() + } else { + let cell = Arc::new(OnceCell::new()); + self.pools + .write() + .insert(object_meta.location.clone(), cell.clone()); + cell + }; + let pool = cell + .get_or_try_init(|| async { + let open_timer = scan_metrics.open_time.timer(); + let atlas_view = AtlasView::new( + cache, + store.clone(), + object_meta.clone(), + logical_schema.clone(), + ) + .await?; + drop(open_timer); + let datasets = atlas_view + .list_datasets(pruning_predicate, scan_metrics.clone()) + .await?; + let queue = ArrayQueue::new(datasets.len()); + for dataset in datasets { + queue.push(dataset).unwrap(); + } + + Ok::, anyhow::Error>(Arc::new(Level1Pool { + inner: Arc::new(InnerLevel1Pool { + atlas_view, + queue, + projected_schema: projected_schema.clone(), + }), + })) + }) + .await + .cloned()?; + + Ok(pool.as_ref().clone().into_stream(scan_metrics).boxed()) + } +} + +/// A handle on a collection's shared queue. +/// +/// The handle holds no read state, so it is shared between threads freely. +/// [`Level1Pool::into_stream`] makes one consumer of the queue. The +/// partitions that share a collection each hold a consumer, and the queue +/// shares the work between them. +#[derive(Clone)] +pub struct Level1Pool { + inner: Arc, +} + +impl Level1Pool { + /// One consumer of the queue. It starts on no dataset, and counts the + /// datasets it reads on `scan_metrics`. + pub fn into_stream(self, scan_metrics: AtlasScanMetrics) -> Level1PoolStream { + Level1PoolStream { + inner: self.inner, + scan_metrics, + current: None, + } + } +} + +struct InnerLevel1Pool { + atlas_view: AtlasView, + /// The datasets left to read, in listing order. + queue: ArrayQueue, + /// The scan's output schema, nd-encoded. Every batch goes out under it. + projected_schema: SchemaRef, +} + +/// One consumer of a collection's shared queue. +/// +/// It streams the datasets it pops, one nd-encoded batch per stored chunk. +/// The dataset in hand is a boxed stream, which is `Send` and not `Sync`, so +/// the consumer lives with the partition that polls it and never in the +/// shared handle. +pub struct Level1PoolStream { + inner: Arc, + /// The partition's metrics. The datasets this consumer reads count here. + scan_metrics: AtlasScanMetrics, + /// The dataset this consumer is draining, if any. + current: Option>>, +} + +impl Stream for Level1PoolStream { + type Item = anyhow::Result; // Record Batches using nd encoding + + /// The next batch of the dataset in hand, or the first batch of the next + /// dataset on the queue. The stream ends when the queue is empty and the + /// dataset in hand is drained. + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + loop { + if let Some(current) = this.current.as_mut() { + match current.as_mut().poll_next(cx) { + Poll::Ready(Some(item)) => return Poll::Ready(Some(item)), + Poll::Ready(None) => this.current = None, + Poll::Pending => return Poll::Pending, + } + } + match this.inner.queue.pop() { + Some(dataset) => { + this.current = Some(dataset_stream( + Arc::clone(&this.inner), + this.scan_metrics.clone(), + dataset, + )); + } + None => return Poll::Ready(None), + } + } + } +} + +/// One dataset's batches: one encoded nd batch per stored chunk, in C order. +/// +/// The dataset is built when the stream is first polled, and each chunk is +/// read when the stream reaches it. A chunk goes out under the table's +/// fields, cast to the types the table declares, and under the encoded table +/// schema, so every batch of the pool has the one schema the scan expects. +fn dataset_stream( + pool: Arc, + scan_metrics: AtlasScanMetrics, + dataset: String, +) -> BoxStream<'static, anyhow::Result> { + futures::stream::once(async move { + let source = pool + .atlas_view + .dataset(&dataset) + .await? + .ok_or_else(|| anyhow::anyhow!("dataset '{dataset}' is not in the collection"))?; + scan_metrics.datasets_scanned.add(1); + let chunks = source.chunks(); + let dataset = Arc::::from(dataset); + Ok::<_, anyhow::Error>(futures::stream::iter(chunks).then(move |chunk| { + let pool = Arc::clone(&pool); + let source = Arc::clone(&source); + let dataset = Arc::clone(&dataset); + async move { + let nd = source.poll_next(chunk).await?.ok_or_else(|| { + anyhow::anyhow!("dataset '{dataset}' read no batch for a chunk of its own grid") + })?; + let nd = under_fields(&nd, pool.atlas_view.table_schema().fields())?; + let batch = + encode_nd_record_batch(&nd)?.with_schema(Arc::clone(&pool.projected_schema))?; + Ok(batch) + } + })) + }) + .try_flatten() + .boxed() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{compat, test_support}; + use beacon_datafusion_ext::nd::{decode_nd_record_batch, encoded_schema}; + use beacon_datafusion_ext::type_widening::ArrowTypeWidening; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Column as ColumnExpr, Literal}; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; + use datafusion::scalar::ScalarValue; + use std::path::Path; + + /// The logical schema of a fixture, the one a scan hands the pool. + async fn logical_schema(dir: &Path) -> SchemaRef { + let atlas = test_support::open(dir).await; + Arc::new( + compat::collection_arrow_schema( + &atlas.footer().collection_schema(), + &ArrowTypeWidening::default_extension(), + ) + .unwrap(), + ) + } + + /// One handle on the pool of a fixture. + async fn handle( + pool: &AtlasReaderPool, + dir: &Path, + predicate: Option>, + metrics: AtlasScanMetrics, + ) -> BoxStream<'static, anyhow::Result> { + let (store, marker) = test_support::store_and_marker(dir); + let logical = logical_schema(dir).await; + let projected = Arc::new(encoded_schema(&logical)); + pool.try_open_into_pooled_stream( + None, store, marker, logical, projected, predicate, metrics, + ) + .await + .unwrap() + } + + /// The scan holds the pool in a `FileSource`, which must be `Send + Sync`. + /// A consumer is not `Sync`, so it must never sit in the shared handle. + #[test] + fn the_pool_is_send_and_sync() { + fn assert_send_sync() {} + assert_send_sync::(); + assert_send_sync::(); + } + + fn rows(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .map(|batch| decode_nd_record_batch(batch).unwrap().num_rows()) + .collect() + } + + /// One encoded batch per dataset, in listing order, each on the encoded + /// table schema. + #[tokio::test] + async fn the_pool_streams_one_encoded_batch_per_dataset() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let set = ExecutionPlanMetricsSet::new(); + let metrics = AtlasScanMetrics::new(&set, 0); + let pool = AtlasReaderPool::new(); + + let batches: Vec = handle(&pool, tmp.path(), None, metrics.clone()) + .await + .try_collect() + .await + .unwrap(); + + let logical = logical_schema(tmp.path()).await; + let encoded = Arc::new(encoded_schema(&logical)); + assert_eq!(batches.len(), 2); + for batch in &batches { + assert_eq!(batch.schema(), encoded, "the table schema, encoded"); + } + assert_eq!(rows(&batches), vec![4, 3], "winter, then summer"); + assert_eq!(metrics.datasets_scanned.value(), 2); + } + + /// A column the dataset lacks reads as a null column, so every batch fits + /// the one schema. `summer` declares no `cycle`. + #[tokio::test] + async fn a_column_the_dataset_lacks_is_all_null() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let set = ExecutionPlanMetricsSet::new(); + let pool = AtlasReaderPool::new(); + + let batches: Vec = + handle(&pool, tmp.path(), None, AtlasScanMetrics::new(&set, 0)) + .await + .try_collect() + .await + .unwrap(); + + let summer = decode_nd_record_batch(&batches[1]) + .unwrap() + .materialize() + .unwrap(); + assert_eq!(summer.num_rows(), 3); + assert_eq!(summer.column_by_name("cycle").unwrap().null_count(), 3); + } + + /// Two handles on one collection share one queue. A dataset one handle + /// reads is not read again by the other. + #[tokio::test] + async fn two_handles_share_one_queue() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let set = ExecutionPlanMetricsSet::new(); + let metrics = AtlasScanMetrics::new(&set, 0); + let pool = AtlasReaderPool::new(); + + let mut first = handle(&pool, tmp.path(), None, metrics.clone()).await; + let mut second = handle(&pool, tmp.path(), None, metrics.clone()).await; + + let winter = first.try_next().await.unwrap().unwrap(); + assert_eq!(rows(&[winter]), vec![4]); + let summer = second.try_next().await.unwrap().unwrap(); + assert_eq!( + rows(&[summer]), + vec![3], + "the second handle got the next dataset" + ); + assert!( + first.try_next().await.unwrap().is_none(), + "nothing left for the first" + ); + assert!( + second.try_next().await.unwrap().is_none(), + "nor for the second" + ); + assert_eq!(metrics.datasets_scanned.value(), 2); + } + + /// A predicate the statistics can judge skips the datasets it rules out + /// before any of them is read. + #[tokio::test] + async fn a_predicate_prunes_datasets_before_the_read() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 10).await; + let set = ExecutionPlanMetricsSet::new(); + let metrics = AtlasScanMetrics::new(&set, 0); + let pool = AtlasReaderPool::new(); + let predicate: Arc = Arc::new(BinaryExpr::new( + Arc::new(ColumnExpr::new("temperature", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Float32(Some(45.0)))), + )); + + let batches: Vec = handle(&pool, tmp.path(), Some(predicate), metrics.clone()) + .await + .try_collect() + .await + .unwrap(); + + assert_eq!(batches.len(), 5, "d5 to d9 reach past 45"); + assert_eq!(metrics.datasets_pruned.value(), 5); + assert_eq!(metrics.datasets_scanned.value(), 5); + } +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs index 633a564a..d6903df3 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs @@ -10,14 +10,17 @@ //! //! # What an open does //! -//! Opening a collection costs one footer read through the reader cache. The -//! opener then lists the datasets, prunes them in one vectorised pass over the -//! footer's statistics, and is left with the names it has to read. Those names -//! feed one stream: each dataset is built in turn, planned as a [`FileRead`], -//! and its batches are yielded before the next dataset is touched. +//! An open goes through the [`AtlasReaderPool`] the source holds. The first +//! partition to reach a collection opens it, at the cost of one footer read +//! through the reader cache, prunes its datasets in one vectorised pass over +//! the footer's statistics, and queues the names it has to read. Every +//! partition that opens the collection then streams the datasets it pops off +//! that queue: each dataset is built in turn, and its batches are yielded +//! before the next dataset is touched. //! -//! So a pruned dataset costs nothing at all, and a kept one costs its build and -//! its read. Nothing is listed at plan time, and nothing is queued. +//! So a pruned dataset costs nothing at all, a kept one costs its build and +//! its read, and a dataset is read by one partition and by no other. Nothing +//! is listed at plan time. //! //! [`AtlasFormat`]: super::AtlasFormat @@ -53,7 +56,7 @@ use object_store::ObjectStore; use beacon_datafusion_ext::nd::logical_schema; -use crate::datafusion::metrics::AtlasScanMetrics; +use crate::datafusion::{metrics::AtlasScanMetrics, pool::AtlasReaderPool}; use crate::store::{AtlasReaderCache, get_or_open_atlas}; use crate::{compat, datafusion::opener::AtlasOpener}; @@ -68,6 +71,7 @@ pub struct AtlasSource { projection: Option, /// The reader cache to consult, or `None` to open every collection afresh. cache: AtlasReaderCache, + reader_pool: Arc, } impl AtlasSource { @@ -84,6 +88,7 @@ impl AtlasSource { read_dimensions, projection: None, cache, + reader_pool: Arc::new(AtlasReaderPool::new()), } } @@ -123,6 +128,7 @@ impl FileSource for AtlasSource { predicate: self.predicate.clone(), read_metrics: ReadMetrics::new(&self.execution_plan_metrics, partition), scan_metrics: AtlasScanMetrics::new(&self.execution_plan_metrics, partition), + reader_pool: Arc::clone(&self.reader_pool), })) } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs new file mode 100644 index 00000000..25f0df85 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs @@ -0,0 +1,104 @@ +use arrow::datatypes::{FieldRef, SchemaRef}; +use atlas::Atlas; +use beacon_nd_array::{ + NdArrayD, + dataset::{default::DefaultDataset, source::DatasetSource}, +}; +use datafusion::physical_plan::PhysicalExpr; +use indexmap::IndexMap; +use object_store::{ObjectMeta, ObjectStore}; +use std::sync::Arc; + +use crate::{ + compat, + datafusion::{ + metrics::AtlasScanMetrics, + opener::{AtlasColumnView, column_views}, + pruning::prune_datasets, + }, + store::{AtlasReaderCache, get_or_open_atlas}, +}; + +#[derive(Clone)] +pub struct AtlasView { + atlas: Arc, + table_schema: SchemaRef, + column_views: Arc>>, +} + +impl AtlasView { + /// Open the collection at `object_meta`, through `cache` when given, and + /// resolve every column of `table_schema` against it. + pub async fn new( + cache: Option<&AtlasReaderCache>, + store: Arc, + object_meta: ObjectMeta, + table_schema: SchemaRef, + ) -> anyhow::Result { + let atlas = get_or_open_atlas(cache, store, &object_meta).await?; + let views = column_views(&atlas, &table_schema).await?; + + Ok(Self { + atlas, + table_schema, + column_views: Arc::new(views), + }) + } + + /// The table schema the view resolves columns for, in field order. + pub fn table_schema(&self) -> &SchemaRef { + &self.table_schema + } + + pub async fn list_datasets( + &self, + pruning_predicate: Option>, + scan_metrics: AtlasScanMetrics, + ) -> anyhow::Result> { + let mut datasets = self.atlas.list_datasets(); + + if let Some(predicate) = &pruning_predicate { + let prune_timer = scan_metrics.prune_time.timer(); + let listed = datasets.len(); + datasets = + prune_datasets(&self.column_views, datasets, predicate, &self.table_schema).await; + scan_metrics.datasets_pruned.add(listed - datasets.len()); + drop(prune_timer); + } + + Ok(datasets) + } + + pub async fn dataset( + &self, + dataset_name: &str, + ) -> anyhow::Result>> { + let mut arrays: IndexMap> = IndexMap::new(); + for (field, view) in &*self.column_views { + let array = match view { + None => None, + Some(AtlasColumnView::Array { segment }) => match segment.array(dataset_name) { + Some(info) => Some(compat::array_to_nd_array( + Arc::clone(segment), + dataset_name, + &info.dtype, + )?), + None => None, + }, + Some(AtlasColumnView::GlobalAttribute { map }) + | Some(AtlasColumnView::VariableAttribute { map, .. }) => { + // A list has no rank-0 form, and the schema holds no list + // column. A dataset that stores a list under a scalar + // column's key reads as null. + map.get(dataset_name) + .and_then(|attr| compat::attribute_to_nd_array(attr).ok()) + } + }; + if let Some(array) = array { + arrays.insert(field.name().clone(), array); + } + } + let dataset = DefaultDataset::new(dataset_name.to_string(), arrays)?; + Ok(Some(Arc::new(dataset))) + } +} diff --git a/beacon-db/beacon-file-formats/beacon-nd-array/src/arrow/batch.rs b/beacon-db/beacon-file-formats/beacon-nd-array/src/arrow/batch.rs index c08c2405..9501249b 100644 --- a/beacon-db/beacon-file-formats/beacon-nd-array/src/arrow/batch.rs +++ b/beacon-db/beacon-file-formats/beacon-nd-array/src/arrow/batch.rs @@ -181,7 +181,7 @@ pub(crate) async fn read_ragged_range( /// observation variables (all obs-dim groups), variable attributes, /// and global attributes. Row-size variables are excluded. Fields are /// sorted alphabetically. -fn ragged_record_batch_schema(ragged: &RaggedDataset) -> Arc { +pub(crate) fn ragged_record_batch_schema(ragged: &RaggedDataset) -> Arc { let mut fields: Vec = Vec::new(); for (name, var) in &ragged.variables { @@ -244,7 +244,7 @@ fn plan_ragged_batches( /// observation row belonging to cast `i`. Observation variables are /// already contiguous across all casts. Attributes are repeated to /// fill all rows. -async fn ragged_batch_to_record_batch( +pub(crate) async fn ragged_batch_to_record_batch( cast_data: &Dataset, schema: &Arc, obs_dims: &std::collections::HashSet, diff --git a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs index fccb4ea6..d9d44f30 100644 --- a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs +++ b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs @@ -14,7 +14,7 @@ use crate::{ }; #[derive(Debug, Clone)] -pub struct Dataset { +pub struct DefaultDataset { pub name: String, pub dimensions: Vec, pub shape: Vec, @@ -22,7 +22,7 @@ pub struct Dataset { pub arrays: IndexMap>, } -impl Dataset { +impl DefaultDataset { pub fn new(name: String, arrays: IndexMap>) -> anyhow::Result { let mut dimensions = Vec::new(); let mut shape = Vec::new(); @@ -75,7 +75,7 @@ impl Dataset { } #[async_trait::async_trait] -impl DatasetSource for Dataset { +impl DatasetSource for DefaultDataset { /// Every chunk of the dataset grid, in C order, as an [`ArraySubset`]. /// /// The cut comes from [`generate_chunk_subsets`], so a boundary chunk @@ -158,8 +158,8 @@ mod tests { Arc::new(NdArray::new_with_backend(backend).unwrap()) } - fn dataset(arrays: Vec<(&str, Arc)>) -> anyhow::Result { - Dataset::new( + fn dataset(arrays: Vec<(&str, Arc)>) -> anyhow::Result { + DefaultDataset::new( "test".to_string(), arrays .into_iter() @@ -259,7 +259,7 @@ mod tests { /// A dataset with a coordinate per axis and one 2-D variable, cut on the /// variable's chunk shape. - fn gridded(chunk: &[usize]) -> Dataset { + fn gridded(chunk: &[usize]) -> DefaultDataset { let time = NdArray::::try_new_from_vec_in_mem( (0..4).map(|v| v * 100).collect(), vec![4], @@ -292,7 +292,7 @@ mod tests { arrays.insert("time".to_string(), Arc::new(time)); arrays.insert("lat".to_string(), Arc::new(lat)); arrays.insert("sst".to_string(), Arc::new(sst)); - Dataset { + DefaultDataset { name: "gridded".to_string(), dimensions: names(&["time", "lat"]), shape: vec![4, 3], @@ -302,7 +302,7 @@ mod tests { } /// Poll every chunk, broadcast each, and stitch the rows back together. - async fn read_all(ds: &Dataset) -> arrow::record_batch::RecordBatch { + async fn read_all(ds: &DefaultDataset) -> arrow::record_batch::RecordBatch { let mut batches = Vec::new(); for chunk in ds.chunks() { let nd = ds.poll_next(chunk).await.unwrap().unwrap(); diff --git a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/ragged.rs b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/ragged.rs index 78b4c781..5b560500 100644 --- a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/ragged.rs +++ b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/ragged.rs @@ -1,5 +1,11 @@ -use std::{collections::HashMap, sync::Arc}; +use std::{ + any::Any, + collections::{HashMap, HashSet}, + sync::Arc, +}; +use arrow::record_batch::RecordBatch; +use beacon_datafusion_ext::nd::{Dimension, Dimensions, NdArrowArray, NdRecordBatch}; use indexmap::IndexMap; use tokio::sync::OnceCell; @@ -8,7 +14,10 @@ use num_traits::ToPrimitive; use crate::{ NdArray, NdArrayD, array::subset::ArraySubset, - dataset::Dataset, + arrow::batch::{ + generate_chunk_subsets, ragged_batch_to_record_batch, ragged_record_batch_schema, + }, + dataset::{Dataset, source::DatasetSource}, datatypes::{NdArrayDataType, NdArrayType}, }; @@ -364,7 +373,9 @@ impl RaggedDataset { } RaggedArray::ObservationVariable(array) => { let obs_dim = array.dimensions().first().cloned().ok_or_else(|| { - anyhow::anyhow!("observation variable {name} must have at least one dimension") + anyhow::anyhow!( + "observation variable {name} must have at least one dimension" + ) })?; let cum = &offsets[&obs_dim]; let obs_start = cum[index]; @@ -445,7 +456,9 @@ impl RaggedDataset { } RaggedArray::ObservationVariable(array) => { let obs_dim = array.dimensions().first().cloned().ok_or_else(|| { - anyhow::anyhow!("observation variable {name} must have at least one dimension") + anyhow::anyhow!( + "observation variable {name} must have at least one dimension" + ) })?; let cum = &offsets[&obs_dim]; let obs_start = cum[start]; @@ -499,6 +512,89 @@ impl RaggedDataset { } } +impl RaggedDataset { + /// The chunk size the file stores the instance dimension in. + /// + /// The smallest chunk any instance variable or row-size variable reports + /// on that axis, so a chunk of casts lies inside one stored chunk of each + /// of them. An array with no chunk layout reports its whole axis, so a + /// file with no chunking at all is one chunk of every cast. + fn instance_chunk(&self) -> usize { + let instance_arrays = self.variables.values().filter_map(|var| match var { + RaggedArray::InstanceVariable(array) => Some(array), + _ => None, + }); + instance_arrays + .chain(self.row_size_arrays.values()) + .filter_map(|array| array.chunk_shape().first().copied()) + .filter(|&chunk| chunk > 0) + .min() + .unwrap_or(self.n_instances) + .max(1) + } +} + +#[async_trait::async_trait] +impl DatasetSource for RaggedDataset { + /// Every chunk of the instance dimension, in order, as an [`ArraySubset`] + /// over that one axis. + /// + /// The cut follows the stored chunking of the instance dimension, see + /// [`RaggedDataset::instance_chunk`]. A boundary chunk shrinks to fit. + /// The observation rows of a chunk follow from the offsets when it is + /// read. A dataset with no casts has no chunks. + fn chunks(&self) -> Vec> { + generate_chunk_subsets(&[self.n_instances], &[self.instance_chunk()]) + .into_iter() + .map(|subset| Arc::new(subset) as Arc) + .collect() + } + + /// Read one chunk of casts into an [`NdRecordBatch`]. + /// + /// `chunk` is one of [`DatasetSource::chunks`]. The casts are read in one + /// pass per array, and their rows come out flat: an instance value repeats + /// for every observation row of its cast, and an attribute for every row. + /// The batch then sits on one synthetic `row` axis, on which every column + /// is full rank, the layout the flat nd encoding gives a ragged batch. + async fn poll_next( + &self, + chunk: Arc, + ) -> anyhow::Result> { + let subset = chunk + .downcast_ref::() + .ok_or_else(|| anyhow::anyhow!("chunk is not an ArraySubset"))?; + let (Some(&start), Some(&len)) = (subset.start.first(), subset.shape.first()) else { + anyhow::bail!("a ragged chunk spans the instance dimension; this one spans no axis"); + }; + let range = start..start + len; + + let casts = self.get_casts_range(range.start, range.end).await?; + let schema = ragged_record_batch_schema(self); + let obs_dims: HashSet = self.observation_dimensions().map(String::from).collect(); + let offsets = self.cumulative_offsets().await?; + let flat = + ragged_batch_to_record_batch(&casts, &schema, &obs_dims, offsets, &range).await?; + Ok(Some(flat_batch_to_nd(&flat)?)) + } +} + +/// A flat batch as an nd batch on one synthetic `row` axis. +/// +/// Every column is full rank on that axis, so the broadcast above the scan is +/// the identity. This is the layout `encode_flat_batch_as_nd` gives a ragged +/// batch, built here without the encoding. +fn flat_batch_to_nd(batch: &RecordBatch) -> anyhow::Result { + let rows = batch.num_rows(); + let row_dim = || Dimensions::try_new(vec![Dimension::new("row", rows)]); + let columns = batch + .columns() + .iter() + .map(|column| NdArrowArray::try_new(column.clone(), row_dim()?)) + .collect::, _>>()?; + Ok(NdRecordBatch::try_new(batch.schema(), columns, row_dim()?)?) +} + /// Iterator over cast indices, yielding `(index, &RaggedDataset)`. pub struct RaggedIter<'a> { ragged: &'a RaggedDataset, @@ -553,3 +649,168 @@ impl Clone for RaggedDataset { } } } + +#[cfg(test)] +mod source_tests { + use super::*; + use crate::array::backend::{ArrayBackend, mem::InMemoryArrayBackend}; + use arrow::array::{Array, Float64Array, StringArray}; + use ndarray::ArrayD; + + /// An in-memory array that reports the chunk shape a test hands it. + #[derive(Debug)] + struct Chunked { + inner: InMemoryArrayBackend, + chunk_shape: Vec, + } + + #[async_trait::async_trait] + impl ArrayBackend for Chunked { + fn len(&self) -> usize { + self.inner.len() + } + fn shape(&self) -> Vec { + self.inner.shape() + } + fn chunk_shape(&self) -> Vec { + self.chunk_shape.clone() + } + fn dimensions(&self) -> Vec { + self.inner.dimensions() + } + async fn read_subset(&self, subset: ArraySubset) -> anyhow::Result> { + self.inner.read_subset(subset).await + } + } + + fn f64s(values: Vec, dim: &str) -> Arc { + let len = values.len(); + Arc::new( + NdArray::::try_new_from_vec_in_mem(values, vec![len], vec![dim.to_string()], None) + .unwrap(), + ) + } + + fn text(value: &str) -> Arc { + Arc::new( + NdArray::::try_new_from_vec_in_mem( + vec![value.to_string()], + vec![], + vec![] as Vec, + None, + ) + .unwrap(), + ) + } + + /// Three casts of two, one and three observations, the row-size variable + /// stored in chunks of `chunk` casts. + async fn ragged(chunk: usize) -> RaggedDataset { + let sizes = vec![2, 1, 3]; + let inner = InMemoryArrayBackend::new( + ArrayD::from_shape_vec(vec![3], sizes).unwrap(), + vec![3], + vec!["casts".to_string()], + None, + ); + let row_size: Arc = Arc::new( + NdArray::new_with_backend(Chunked { + inner, + chunk_shape: vec![chunk], + }) + .unwrap(), + ); + + let mut arrays: IndexMap> = IndexMap::new(); + arrays.insert("row_size".to_string(), row_size); + arrays.insert("row_size.sample_dimension".to_string(), text("obs")); + arrays.insert( + "station".to_string(), + f64s(vec![100.0, 200.0, 300.0], "casts"), + ); + arrays.insert( + "depth".to_string(), + f64s(vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0], "obs"), + ); + arrays.insert(".title".to_string(), text("casts")); + + let dataset = Dataset::new("ragged".to_string(), arrays).await; + RaggedDataset::try_new(&dataset).await.unwrap() + } + + fn subsets(ragged: &RaggedDataset) -> Vec<(usize, usize)> { + ragged + .chunks() + .iter() + .map(|chunk| { + let subset = chunk.downcast_ref::().unwrap(); + (subset.start[0], subset.shape[0]) + }) + .collect() + } + + #[tokio::test] + async fn chunks_follow_the_stored_chunking_of_the_instance_dimension() { + assert_eq!(subsets(&ragged(2).await), vec![(0, 2), (2, 1)]); + assert_eq!(subsets(&ragged(1).await), vec![(0, 1), (1, 1), (2, 1)]); + assert_eq!( + subsets(&ragged(3).await), + vec![(0, 3)], + "no chunking is one chunk of every cast" + ); + } + + /// Each chunk comes out flat on a `row` axis, and the chunks together are + /// every observation row once, in cast order. + #[tokio::test] + async fn polling_every_chunk_reads_every_row_once() { + let ragged = ragged(2).await; + let mut station = Vec::new(); + let mut depth = Vec::new(); + let mut rows = Vec::new(); + for chunk in ragged.chunks() { + let nd = ragged.poll_next(chunk).await.unwrap().unwrap(); + assert_eq!(nd.target().rank(), 1, "one synthetic row axis"); + assert!( + nd.columns().iter().all(|column| column.dims().rank() == 1), + "every column is full rank on it" + ); + rows.push(nd.num_rows()); + + let batch = nd.materialize().unwrap(); + let get = |name: &str| batch.column_by_name(name).unwrap().clone(); + station.extend( + get("station") + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied(), + ); + depth.extend( + get("depth") + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied(), + ); + let title = get(".title"); + let title = title.as_any().downcast_ref::().unwrap(); + assert!((0..title.len()).all(|row| title.value(row) == "casts")); + } + + assert_eq!(rows, vec![3, 3], "casts 0 and 1, then cast 2"); + assert_eq!(station, vec![100.0, 100.0, 200.0, 300.0, 300.0, 300.0]); + assert_eq!(depth, vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0]); + } + + #[tokio::test] + async fn a_chunk_of_another_type_is_an_error() { + let ragged = ragged(3).await; + let err = ragged.poll_next(Arc::new(42usize)).await.unwrap_err(); + assert!(err.to_string().contains("not an ArraySubset"), "{err}"); + } +} From 969c6600e5dd3a2d54c4a6517eb9ac65ab3b5bfe Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 9 Sep 2026 13:19:55 +0200 Subject: [PATCH 12/16] wip --- .../beacon-arrow-atlas/src/datafusion/pool.rs | 31 ++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs index f1c85a30..30eb0991 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs @@ -98,7 +98,10 @@ impl AtlasReaderPool { let datasets = atlas_view .list_datasets(pruning_predicate, scan_metrics.clone()) .await?; - let queue = ArrayQueue::new(datasets.len()); + // A queue has at least one slot: `ArrayQueue::new(0)` panics. + // With no dataset to read the slot stays empty, the first + // `pop` finds nothing, and the stream ends at once. + let queue = ArrayQueue::new(datasets.len().max(1)); for dataset in datasets { queue.push(dataset).unwrap(); } @@ -371,6 +374,32 @@ mod tests { assert_eq!(metrics.datasets_scanned.value(), 2); } + /// A predicate that rules out every dataset leaves nothing to read. The + /// stream ends at once, and nothing panics on the empty queue. + #[tokio::test] + async fn a_collection_with_nothing_to_read_streams_nothing() { + let tmp = tempfile::tempdir().unwrap(); + test_support::ranged(tmp.path(), 10).await; + let set = ExecutionPlanMetricsSet::new(); + let metrics = AtlasScanMetrics::new(&set, 0); + let pool = AtlasReaderPool::new(); + let predicate: Arc = Arc::new(BinaryExpr::new( + Arc::new(ColumnExpr::new("temperature", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Float32(Some(1.0e9)))), + )); + + let batches: Vec = handle(&pool, tmp.path(), Some(predicate), metrics.clone()) + .await + .try_collect() + .await + .unwrap(); + + assert!(batches.is_empty()); + assert_eq!(metrics.datasets_pruned.value(), 10); + assert_eq!(metrics.datasets_scanned.value(), 0); + } + /// A predicate the statistics can judge skips the datasets it rules out /// before any of them is read. #[tokio::test] From 4332daa5b72157aed4a8209174f4e2345d220352 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 9 Sep 2026 13:36:42 +0200 Subject: [PATCH 13/16] wip --- .../beacon-arrow-atlas/src/datafusion/mod.rs | 204 ++++++++++++++++-- .../src/datafusion/source.rs | 13 +- .../beacon-arrow-atlas/src/lib.rs | 12 +- 3 files changed, 200 insertions(+), 29 deletions(-) diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs index 54b2aaf1..d2571ca7 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs @@ -293,12 +293,14 @@ impl FileFormat for AtlasFormat { Ok(Statistics::new_unknown(&table_schema)) } - /// Plan one entry per collection, then wrap the scan in the nd spine. + /// Plan every collection into every partition, then wrap the scan in the + /// nd spine. /// /// Nothing is opened here. The markers the listing found are deduped to the - /// outermost collections and dealt round-robin over the target partitions, - /// and each partition's opener lists, prunes and reads the collections it - /// holds. Parallelism is therefore bounded by the collection count. + /// outermost collections, and every target partition gets all of them in + /// its own rotation, see [`deal_rotated`]. The partitions that open one + /// collection share its datasets through the reader pool, so parallelism + /// is bounded by the dataset count, not the collection count. async fn create_physical_plan( &self, state: &dyn Session, @@ -314,21 +316,10 @@ impl FileFormat for AtlasFormat { .collect(); let markers = top_level_atlas_markers(&listed); - // One collection is one unit of work, and a container is never split, - // so the deal here is the whole distribution. - let partitions = state - .config() - .target_partitions() - .clamp(1, markers.len().max(1)); - let mut dealt: Vec> = vec![Vec::new(); partitions]; - for (index, marker) in markers.iter().enumerate() { - dealt[index % partitions].push(PartitionedFile::from(marker.clone())); - } - let file_groups: Vec = dealt - .into_iter() - .filter(|group| !group.is_empty()) - .map(FileGroup::new) - .collect(); + // A container is never split, and the reader pool shares one between + // the partitions that open it, so the deal here is the whole + // distribution. + let file_groups = deal_rotated(&markers, state.config().target_partitions()); tracing::debug!( collections = markers.len(), partitions = file_groups.len(), @@ -380,3 +371,178 @@ impl FileFormat for AtlasFormat { )) } } + +/// Deal `markers` over `partitions` groups, every collection to every group. +/// +/// The reader pool shares a collection between the partitions that open it, +/// so a partition may hold every collection and still read nothing twice. +/// Every partition then reads until every collection is drained, whatever +/// the collections' sizes, and parallelism is bounded by the dataset count +/// rather than the collection count. +/// +/// Each group is the whole list, rotated. Group `p` starts `p * n / partitions` +/// collections round the ring, so the start points spread evenly: with as many +/// groups as collections each starts on its own, with fewer they start as far +/// apart as they can, and with more they double up as evenly as they can. A +/// partition therefore works alone on its collection until the partitions +/// meet, and the shared queue takes over from there. +pub(crate) fn deal_rotated(markers: &[ObjectMeta], partitions: usize) -> Vec { + let n = markers.len(); + if n == 0 { + return Vec::new(); + } + let partitions = partitions.max(1); + (0..partitions) + .map(|p| { + let offset = p * n / partitions; + (0..n) + .map(|i| PartitionedFile::from(markers[(offset + i) % n].clone())) + .collect() + }) + .collect() +} + +#[cfg(test)] +mod deal_tests { + use super::*; + + fn markers(n: usize) -> Vec { + (0..n) + .map(|i| ObjectMeta { + location: object_store::path::Path::from(format!("c{i}/{ATLAS_MARKER}")), + last_modified: chrono::Utc::now(), + size: 0, + e_tag: None, + version: None, + }) + .collect() + } + + /// The collections of a group, by their directory. + fn dealt(group: &FileGroup) -> Vec { + group + .files() + .iter() + .map(|file| { + file.object_meta + .location + .parts() + .next() + .unwrap() + .as_ref() + .to_string() + }) + .collect() + } + + /// The collection each group starts on. + fn starts(groups: &[FileGroup]) -> Vec { + groups.iter().map(|group| dealt(group)[0].clone()).collect() + } + + #[test] + fn every_partition_holds_every_collection_in_its_own_rotation() { + let groups = deal_rotated(&markers(4), 2); + + assert_eq!(groups.len(), 2); + assert_eq!(dealt(&groups[0]), ["c0", "c1", "c2", "c3"]); + assert_eq!(dealt(&groups[1]), ["c2", "c3", "c0", "c1"]); + } + + #[test] + fn the_starts_spread_evenly_over_the_collections() { + assert_eq!( + starts(&deal_rotated(&markers(3), 3)), + ["c0", "c1", "c2"], + "one start per collection" + ); + assert_eq!( + starts(&deal_rotated(&markers(4), 3)), + ["c0", "c1", "c2"], + "fewer partitions start as far apart as they can" + ); + let groups = deal_rotated(&markers(2), 4); + assert_eq!( + starts(&groups), + ["c0", "c0", "c1", "c1"], + "more partitions double up evenly" + ); + assert!(groups.iter().all(|group| group.len() == 2)); + } + + #[test] + fn no_collection_makes_no_group() { + assert!(deal_rotated(&markers(0), 4).is_empty()); + assert_eq!( + deal_rotated(&markers(2), 0).len(), + 1, + "no partition reads as one" + ); + } +} + +#[cfg(test)] +mod scan_tests { + use super::*; + use crate::test_support; + use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig}; + use datafusion::prelude::{SessionConfig, SessionContext}; + use std::path::Path; + + /// The rows of every collection under `dir`, read in a session of + /// `partitions` target partitions, and the partition count of the plan. + async fn rows(dir: &Path, partitions: usize) -> (usize, usize) { + // The collections sit in subdirectories, which a listing skips by + // default. + let config = SessionConfig::new() + .with_target_partitions(partitions) + .set_bool( + "datafusion.execution.listing_table_ignore_subdirectory", + false, + ); + let ctx = SessionContext::new_with_config(config); + let url = ListingTableUrl::parse(format!("{}/", dir.display())).unwrap(); + let options = ListingOptions::new(Arc::new(AtlasFormat::default())) + .with_file_extension(ATLAS_MARKER) + .with_collect_stat(false); + let config = ListingTableConfig::new(url) + .with_listing_options(options) + .infer_schema(&ctx.state()) + .await + .unwrap(); + let table = Arc::new(ListingTable::try_new(config).unwrap()); + + let df = ctx.read_table(table).unwrap(); + let plan = df.clone().create_physical_plan().await.unwrap(); + let partition_count = plan.properties().output_partitioning().partition_count(); + let batches = df.collect().await.unwrap(); + (batches.iter().map(|b| b.num_rows()).sum(), partition_count) + } + + /// Three collections, read by three partitions that each hold all of + /// them. Every row comes out once, and no row comes out twice. + #[tokio::test] + async fn every_partition_reads_through_the_pool_and_no_dataset_twice() { + let tmp = tempfile::tempdir().unwrap(); + for name in ["a", "b", "c"] { + let dir = tmp.path().join(name); + std::fs::create_dir_all(&dir).unwrap(); + test_support::ranged(&dir, 5).await; + } + + let (alone, one) = rows(tmp.path(), 1).await; + let (shared, three) = rows(tmp.path(), 3).await; + + assert_eq!(one, 1); + assert_eq!(three, 3, "every partition holds a group"); + assert_eq!( + alone, + 3 * 5 * 4, + "five datasets of four rows per collection" + ); + assert_eq!( + shared, alone, + "a dataset is read by one partition and no other" + ); + } +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs index d6903df3..8b56cb71 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs @@ -3,10 +3,11 @@ //! # One collection is one unit of work //! //! A plan entry is a collection: its `data.atlas` container, as the listing -//! found it. [`AtlasFormat`] dedupes the markers and deals them round-robin -//! over the target partitions, and a partition reads each collection it holds -//! from end to end. A container is never split by byte range, because a byte -//! range of one means nothing. +//! found it. [`AtlasFormat`] dedupes the markers and deals every one of them +//! to every target partition, each partition in its own rotation. A container +//! is never split by byte range, because a byte range of one means nothing. +//! The partitions that open one collection share its datasets through the +//! reader pool instead. //! //! # What an open does //! @@ -148,7 +149,9 @@ impl FileSource for AtlasSource { } /// A container is one unit. A byte range of it names nothing a reader can - /// open, so the plan's groups stand as the format dealt them. + /// open. The format deals every collection to every partition itself, and + /// the reader pool shares the datasets, so the plan's groups stand as the + /// format dealt them. fn supports_repartitioning(&self) -> bool { false } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs index e8e6611b..5cba4844 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs @@ -35,11 +35,13 @@ //! //! # One collection is one unit of work //! -//! The scan plans one entry per collection and deals them over the partitions. -//! A partition opens each collection it holds once, prunes every dataset in one -//! pass over the footer, and streams the survivors one after another. A pruned -//! dataset therefore costs nothing, and parallelism is bounded by the -//! collection count. See [`datafusion::source`]. +//! The scan plans one entry per collection and deals every entry to every +//! partition, each in its own rotation. The first partition to open a +//! collection prunes every dataset in one pass over the footer and queues the +//! survivors in a reader pool. Every partition that opens the collection then +//! streams the datasets it pops off that queue. A pruned dataset therefore +//! costs nothing, a dataset is read once, and parallelism is bounded by the +//! dataset count. See [`datafusion::source`]. //! //! # Columns //! From 65700168087bb7010c1045fc08ed06a4939c7b4b Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 9 Sep 2026 14:13:35 +0200 Subject: [PATCH 14/16] wip --- .../src/datafusion/error.rs | 14 + .../beacon-arrow-atlas/src/datafusion/mod.rs | 20 +- .../src/datafusion/opener.rs | 351 ++--------------- .../src/datafusion/options.rs | 5 +- .../beacon-arrow-atlas/src/datafusion/pool.rs | 126 +++--- .../src/datafusion/pruning.rs | 12 +- .../src/datafusion/source.rs | 44 +-- .../beacon-arrow-atlas/src/datafusion/view.rs | 371 +++++++++++++++++- .../beacon-arrow-atlas/src/store.rs | 7 +- .../beacon-arrow-atlas/src/test_support.rs | 197 ---------- 10 files changed, 487 insertions(+), 660 deletions(-) create mode 100644 beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/error.rs diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/error.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/error.rs new file mode 100644 index 00000000..28f5885d --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/error.rs @@ -0,0 +1,14 @@ +//! The seam between the crate's own errors and DataFusion's. +//! +//! Inside the crate an error is an `anyhow::Error`, with context added where +//! the context says something the cause does not: the collection, the +//! dataset, the column. At the DataFusion boundary it becomes one +//! `DataFusionError::External`, chain and all, so a query reports what went +//! wrong and where without a second layer of formatting. + +use datafusion::error::DataFusionError; + +/// A crate error, as DataFusion reports it. +pub(crate) fn external(error: anyhow::Error) -> DataFusionError { + DataFusionError::External(error.into()) +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs index d2571ca7..f414b280 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs @@ -10,6 +10,7 @@ use std::any::Any; use std::collections::HashMap; use std::sync::Arc; +use anyhow::Context as _; use arrow::datatypes::{Schema, SchemaRef}; use beacon_datafusion_ext::format_ext::{ DatasetMetadata, FileFormatFactoryExt, SchemaOptions, SchemaUnit, units_over_stores, @@ -35,8 +36,10 @@ use datafusion::{ use object_store::{ObjectMeta, ObjectStore}; use crate::compat; +use crate::datafusion::error::external; use crate::store::{ATLAS_MARKER, AtlasReaderCache, get_or_open_atlas, top_level_atlas_markers}; +pub(crate) mod error; pub mod metrics; pub mod opener; pub mod options; @@ -65,8 +68,8 @@ impl AtlasFormatFactory { Self { options } } - /// A format with this table's effective settings, wired to the shared cache - /// when caching is on. + /// A format with this table's effective settings. Each format owns a + /// reader cache of its own. pub(crate) fn build(&self, options: AtlasOptions) -> AtlasFormat { AtlasFormat::new(options) } @@ -248,16 +251,17 @@ impl FileFormat for AtlasFormat { for marker in &markers { let atlas = get_or_open_atlas(Some(&self.cache), Arc::clone(store), marker) .await - .map_err(|e| exec_datafusion_err!("{e}"))?; + .map_err(external)?; let schema = compat::collection_arrow_schema(&atlas.footer().collection_schema(), &widening) - .map_err(|e| { - exec_datafusion_err!( - "Failed to read the schema of atlas collection '{}': {e}", + .with_context(|| { + format!( + "reading the schema of atlas collection '{}'", marker.location ) - })?; + }) + .map_err(external)?; schemas.push(LabeledSchema::new( Arc::new(schema), marker.location.as_ref(), @@ -298,7 +302,7 @@ impl FileFormat for AtlasFormat { /// /// Nothing is opened here. The markers the listing found are deduped to the /// outermost collections, and every target partition gets all of them in - /// its own rotation, see [`deal_rotated`]. The partitions that open one + /// its own rotation, see `deal_rotated`. The partitions that open one /// collection share its datasets through the reader pool, so parallelism /// is bounded by the dataset count, not the collection count. async fn create_physical_plan( diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs index 02cec72b..b4183d54 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs @@ -2,36 +2,29 @@ //! //! The opener reads through the [`AtlasReaderPool`]. The first partition to //! reach a collection opens it and queues its datasets, and every partition -//! then streams the datasets it pops. What lives here is the column -//! resolution the read and the pruning share: a column view says where one -//! column of the scan comes from, for every dataset at once, and -//! [`under_fields`] puts one chunk of a dataset under the scan's fields. +//! then streams the datasets it pops. How a column resolves, and how a chunk +//! comes out under the scan's fields, lives in [`view`](super::view). use std::sync::Arc; -use arrow::{ - array::{ArrayRef, new_null_array}, - compute::cast, - datatypes::{Field, FieldRef, Schema, SchemaRef}, -}; -use atlas::{ArrayFile, Atlas, Attr}; -use beacon_datafusion_ext::nd::{Dimensions, NdArrowArray, NdRecordBatch}; -use beacon_datafusion_ext::type_widening::is_type_conflict; -use beacon_nd_array::arrow::metrics::ReadMetrics; +use arrow::datatypes::SchemaRef; use datafusion::{ datasource::{ listing::PartitionedFile, physical_plan::{FileOpenFuture, FileOpener}, }, - error::{DataFusionError, Result}, + error::Result, physical_plan::PhysicalExpr, }; use futures::{FutureExt, StreamExt, TryStreamExt}; -use indexmap::IndexMap; use object_store::ObjectStore; use crate::{ - datafusion::{metrics::AtlasScanMetrics, pool::AtlasReaderPool}, + datafusion::{ + error::external, + metrics::AtlasScanMetrics, + pool::{AtlasReaderPool, PoolOpen}, + }, store::AtlasReaderCache, }; @@ -50,9 +43,7 @@ pub struct AtlasOpener { /// and the pruning engine are written against. pub logical_schema: SchemaRef, pub read_dimensions: Option>, - pub batch_size: usize, pub predicate: Option>, - pub read_metrics: ReadMetrics, pub scan_metrics: AtlasScanMetrics, /// The scan's pools, one per collection, shared by every partition. pub reader_pool: Arc, @@ -78,192 +69,38 @@ impl FileOpener for AtlasOpener { let pool = Arc::clone(&self.reader_pool); let fut = async move { - let location = file.object_meta.location.clone(); + let open = PoolOpen { + cache: Some(&cache), + logical_schema, + projected_schema, + predicate, + scan_metrics, + }; let stream = pool - .try_open_into_pooled_stream( - Some(&cache), - store, - file.object_meta, - logical_schema, - projected_schema, - predicate, - scan_metrics, - ) - .await - .map_err(|e| { - DataFusionError::Execution(format!( - "Failed to open atlas collection '{location}': {e}" - )) - })?; - Ok(stream - .map_err(|e| DataFusionError::External(e.into())) - .boxed()) - }; - - Ok(fut.boxed()) - } -} - -/// Where each column of the scan comes from, for every dataset at once. -/// -/// One segment open per array, and one attribute sweep per key. Each costs the -/// same however many datasets the collection holds, so a partition pays them -/// once and reads every dataset against the result. A column no dataset -/// declares gets `None`. -pub(crate) async fn column_views( - atlas: &Atlas, - logical_schema: &Schema, -) -> Result>> { - let mut views = IndexMap::with_capacity(logical_schema.fields().len()); - for field in logical_schema.fields() { - let view = if let Some(key) = field.name().strip_prefix('.') { - let map = atlas - .attributes_by_dataset(None, key) + .open(store, file.object_meta, open) .await .map_err(external)?; - Some(AtlasColumnView::GlobalAttribute { map }) - } else if let Some((array, key)) = field.name().split_once('.') { - let map = atlas - .attributes_by_dataset(Some(array), key) - .await - .map_err(external)?; - Some(AtlasColumnView::VariableAttribute { - variable: array.to_string(), - map, - }) - } else { - atlas - .try_segment(field.name()) - .await - .map_err(external)? - .map(|segment| AtlasColumnView::Array { - segment: Arc::clone(segment), - }) + Ok(stream.map_err(external).boxed()) }; - views.insert(Arc::clone(field), view); - } - Ok(views) -} - -/// `nd` under `fields`: every field in order, on the same target grid. -/// -/// A column comes out under the array's own type, and the table may declare a -/// wider one: that is a cast. A field the dataset lacks is a rank-0 null, -/// which broadcasts to an all-null column. The decoder makes the same of a -/// null struct row, so the scan sees one thing either way. -pub(crate) fn under_fields(nd: &NdRecordBatch, fields: &[FieldRef]) -> Result { - let mut columns = Vec::with_capacity(fields.len()); - for field in fields { - let column = match nd.schema().column_with_name(field.name()) { - Some((index, _)) => { - let column = nd.column(index); - match as_field_type(Arc::clone(column.values()), field)? { - Some(values) => NdArrowArray::try_new(values, column.dims().clone())?, - None => null_scalar(field), - } - } - None => null_scalar(field), - }; - columns.push(column); - } - let schema = Arc::new(Schema::new(fields.to_vec())); - NdRecordBatch::try_new(schema, columns, nd.target().clone()) -} -/// A rank-0 null. It broadcasts to an all-null column of the target grid. -fn null_scalar(field: &Field) -> NdArrowArray { - NdArrowArray::try_new(new_null_array(field.data_type(), 1), Dimensions::scalar()) - .expect("one element on no axis") -} - -/// `values` in the type the table declares for `field`, or `None` for values -/// the table cannot hold. -/// -/// A dataset may store a column narrower than the merged type, and the merge -/// widened it: that is a cast. A column the merge could not join is marked, -/// and a dataset of the other family then reads as null. That is what the mark -/// promises the scan. -fn as_field_type(values: ArrayRef, field: &Field) -> Result> { - if values.data_type() == field.data_type() { - return Ok(Some(values)); - } - match cast(&values, field.data_type()) { - Ok(values) => Ok(Some(values)), - Err(_) if is_type_conflict(field) => Ok(None), - Err(error) => Err(error.into()), + Ok(fut.boxed()) } } -/// An atlas error, as the scan reports it. -fn external(error: impl std::error::Error + Send + Sync + 'static) -> DataFusionError { - DataFusionError::External(Box::new(error)) -} - -/// Where one column of the scan comes from, for every dataset of a collection. -/// -/// The scan reads through it, and pruning judges through it, so both see one -/// resolution of a column name. -pub(crate) enum AtlasColumnView { - /// The variable's segment. It holds the array for every dataset that - /// declares it, keyed by dataset name. - Array { segment: Arc }, - /// A dataset-level attribute, its value per dataset. - GlobalAttribute { map: IndexMap }, - /// An attribute of one array, its value per dataset. - VariableAttribute { - variable: String, - map: IndexMap, - }, -} - #[cfg(test)] mod tests { - use arrow::array::{Array, AsArray, RecordBatch}; - use arrow::datatypes::{Float32Type, Float64Type, Int32Type, Int64Type}; - use beacon_datafusion_ext::type_widening::ArrowTypeWidening; - - use super::*; - use crate::datafusion::view::AtlasView; - use crate::{compat, test_support}; - use std::path::Path; - + use arrow::array::{Array, ArrayRef, RecordBatch}; use beacon_datafusion_ext::nd::{decode_nd_record_batch, encoded_schema}; + use beacon_datafusion_ext::type_widening::ArrowTypeWidening; use datafusion::logical_expr::Operator; use datafusion::physical_expr::expressions::{BinaryExpr, Column as ColumnExpr, Literal}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::scalar::ScalarValue; use futures::TryStreamExt; + use std::path::Path; - /// The schema `infer_schema` derives for a fixture. - async fn schema(dir: &Path) -> SchemaRef { - let atlas = test_support::open(dir).await; - Arc::new( - compat::collection_arrow_schema( - &atlas.footer().collection_schema(), - &ArrowTypeWidening::default_extension(), - ) - .unwrap(), - ) - } - - /// Every chunk of `dataset`, read as the scan reads it: through the view, - /// under the scan's fields. And the rows of all of them in chunk order. - async fn read(dir: &Path, dataset: &str) -> (Vec, RecordBatch) { - let schema = schema(dir).await; - let (store, marker) = test_support::store_and_marker(dir); - let view = AtlasView::new(None, store, marker, Arc::clone(&schema)) - .await - .unwrap(); - let source = view.dataset(dataset).await.unwrap().unwrap(); - let mut chunks = Vec::new(); - for chunk in source.chunks() { - let nd = source.poll_next(chunk).await.unwrap().unwrap(); - chunks.push(under_fields(&nd, schema.fields()).unwrap()); - } - let batches: Vec = chunks.iter().map(|nd| nd.materialize().unwrap()).collect(); - let batch = arrow::compute::concat_batches(&schema, &batches).unwrap(); - (chunks, batch) - } + use super::*; + use crate::{compat, test_support}; fn column<'a>(batch: &'a RecordBatch, name: &str) -> &'a ArrayRef { batch @@ -271,144 +108,6 @@ mod tests { .unwrap_or_else(|| panic!("no column {name}")) } - /// Every column comes out on the dataset's grid. An attribute has no axis - /// of its own, so it repeats on every row. - #[tokio::test] - async fn a_dataset_reads_every_column_on_its_own_grid() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let (nd, batch) = read(tmp.path(), "winter").await; - - assert_eq!(nd.len(), 1, "an unchunked array is one chunk"); - assert_eq!(nd[0].target().shape(), vec![4]); - assert_eq!(batch.num_rows(), 4); - assert_eq!( - column(&batch, "temperature") - .as_primitive::() - .values() - .to_vec(), - vec![1.0, 2.0, 3.0, 4.0] - ); - assert_eq!( - column(&batch, "cycle") - .as_primitive::() - .values() - .to_vec(), - vec![10, 20, 30, 40] - ); - let season = column(&batch, ".season").as_string::(); - assert!( - (0..4).all(|row| season.value(row) == "winter"), - "a rank-0 attribute repeats on every row" - ); - assert_eq!( - column(&batch, ".year") - .as_primitive::() - .values() - .to_vec(), - vec![2024; 4] - ); - assert_eq!( - column(&batch, "temperature.units") - .as_string::() - .value(3), - "celsius" - ); - } - - /// `summer` declares neither `cycle` nor `time`, sets no `year`, and has no - /// `units` on `temperature`. Each is a column of nulls on summer's grid. - #[tokio::test] - async fn a_column_the_dataset_lacks_is_all_null() { - let tmp = tempfile::tempdir().unwrap(); - test_support::two_datasets(tmp.path()).await; - - let (_, batch) = read(tmp.path(), "summer").await; - - assert_eq!(batch.num_rows(), 3); - for missing in ["cycle", "time", ".year", "temperature.units"] { - assert_eq!(column(&batch, missing).null_count(), 3, "{missing}"); - } - assert_eq!( - column(&batch, "temperature") - .as_primitive::() - .values() - .to_vec(), - vec![20.0, 21.0, 22.0] - ); - assert_eq!( - column(&batch, ".season").as_string::().value(2), - "summer" - ); - } - - /// A 2-D array reads one stored chunk at a time, keeps both axes, and a - /// cell nobody wrote reads as null. - #[tokio::test] - async fn a_fill_value_reads_as_null_on_a_two_dimensional_grid() { - let tmp = tempfile::tempdir().unwrap(); - test_support::chunked_grid(tmp.path()).await; - - let (nd, batch) = read(tmp.path(), "grid").await; - - assert_eq!(nd.len(), 4, "a [4, 6] grid chunked [2, 3]"); - for chunk in &nd { - assert_eq!(chunk.target().shape(), vec![2, 3]); - } - assert_eq!(batch.num_rows(), 24); - let temperature = column(&batch, "temperature").as_primitive::(); - assert_eq!( - temperature.value(4), - 7.0, - "row 1, column 1 of the grid: the fifth cell of the first chunk" - ); - let mut cells = temperature.values().to_vec(); - cells.sort_by(|a, b| a.partial_cmp(b).unwrap()); - assert_eq!(cells, (0..24).map(f64::from).collect::>()); - let sparse = column(&batch, "sparse"); - assert_eq!(sparse.null_count(), 12, "two of four rows were written"); - assert!( - sparse.is_valid(0), - "the first chunk lies in the written rows" - ); - assert!(sparse.is_null(23), "the last chunk lies outside them"); - } - - /// `a` stores `value` as `Int16` and `b` as `Float32`. The table declares - /// `Float64`, so each dataset casts up to it. `flag` is `a`'s alone. - #[tokio::test] - async fn a_narrower_dataset_casts_to_the_merged_type() { - let tmp = tempfile::tempdir().unwrap(); - test_support::widening(tmp.path()).await; - - let (_, a) = read(tmp.path(), "a").await; - let (_, b) = read(tmp.path(), "b").await; - - assert_eq!( - column(&a, "value") - .as_primitive::() - .values() - .to_vec(), - vec![1.0, 2.0] - ); - assert_eq!( - column(&b, "value") - .as_primitive::() - .values() - .to_vec(), - vec![3.5, 4.5] - ); - assert_eq!( - column(&a, "flag") - .as_primitive::() - .values() - .to_vec(), - vec![7, 8] - ); - assert_eq!(column(&b, "flag").null_count(), 2); - } - // ── the opener ────────────────────────────────────────────────────── /// An opener over a fixture, built the way `AtlasSource` builds one. @@ -430,9 +129,7 @@ mod tests { projected_schema, logical_schema, read_dimensions: None, - batch_size: 8192, predicate: None, - read_metrics: ReadMetrics::new(&metrics, 0), scan_metrics: AtlasScanMetrics::new(&metrics, 0), reader_pool: Arc::new(AtlasReaderPool::new()), }; diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/options.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/options.rs index d6b488de..95f98654 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/options.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/options.rs @@ -2,9 +2,8 @@ /// Settings that change *what* a scan reads, as opposed to how fast it does so. /// -/// The runtime settings live in [`AtlasConfig`](crate::AtlasConfig). These come -/// from the query: `read_atlas(paths, dimensions)` sets them, and so does -/// `CREATE EXTERNAL TABLE ... OPTIONS ('read_dimensions' '…')`. +/// These come from the query: `read_atlas(paths, dimensions)` sets them, and +/// so does `CREATE EXTERNAL TABLE ... OPTIONS ('read_dimensions' '…')`. #[derive(Debug, Default, Clone, PartialEq, Eq)] pub struct AtlasOptions { /// The dimensions the table reads, or `None` to pick a broadcast-compatible diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs index 30eb0991..a9c59798 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs @@ -14,6 +14,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use anyhow::Context as _; use arrow::array::RecordBatch; use arrow::datatypes::SchemaRef; use beacon_datafusion_ext::nd::encode_nd_record_batch; @@ -21,19 +22,37 @@ use crossbeam::queue::ArrayQueue; use datafusion::physical_plan::PhysicalExpr; use futures::stream::BoxStream; use futures::{Stream, StreamExt, TryStreamExt}; -use object_store::{ObjectMeta, ObjectStore}; +use object_store::{ObjectMeta, ObjectStore, path::Path}; use parking_lot::RwLock; use tokio::sync::OnceCell; use crate::datafusion::metrics::AtlasScanMetrics; -use crate::datafusion::opener::under_fields; -use crate::datafusion::view::AtlasView; +use crate::datafusion::view::{AtlasView, under_fields}; use crate::store::AtlasReaderCache; +/// One cell per collection. The cell fills on the first open and never again. +type Pools = HashMap>>>; + /// The pools of one scan, one per collection, keyed by the container's path. #[derive(Default, Clone)] pub struct AtlasReaderPool { - pools: Arc>>>>>, + pools: Arc>, +} + +/// What one open of a collection needs, beyond the collection itself. +/// +/// `logical_schema` is what the columns and the predicate are written +/// against. `projected_schema` is the same schema nd-encoded, and every batch +/// goes out under it. +pub(crate) struct PoolOpen<'a> { + /// The reader cache to open through, or `None` to open afresh. + pub cache: Option<&'a AtlasReaderCache>, + pub logical_schema: SchemaRef, + pub projected_schema: SchemaRef, + /// The predicate to prune datasets with, if any. + pub predicate: Option>, + /// The partition's metrics. The datasets this consumer reads count here. + pub scan_metrics: AtlasScanMetrics, } impl Debug for AtlasReaderPool { @@ -51,73 +70,53 @@ impl AtlasReaderPool { /// One consumer of the collection at `object_meta`. /// - /// The first call for a collection opens it, through `cache` when given, - /// prunes its datasets with `pruning_predicate`, and queues the survivors. - /// Every call gets a stream over that queue. The open and the prune are - /// timed on the first caller's `scan_metrics`, and each consumer counts - /// the datasets it reads on its own. - /// - /// `logical_schema` is what the columns and the predicate are written - /// against. `projected_schema` is the same schema nd-encoded, and every - /// batch goes out under it. - #[allow(clippy::too_many_arguments)] - pub async fn try_open_into_pooled_stream( + /// The first call for a collection opens it, prunes its datasets, and + /// queues the survivors. Every call gets a stream over that queue. The + /// open and the prune are timed on the first caller's metrics, and each + /// consumer counts the datasets it reads on its own. + pub(crate) async fn open( &self, - cache: Option<&AtlasReaderCache>, store: Arc, object_meta: ObjectMeta, - logical_schema: SchemaRef, - projected_schema: SchemaRef, - pruning_predicate: Option>, - scan_metrics: AtlasScanMetrics, - ) -> anyhow::Result>> { - let cell = if self.pools.read().contains_key(&object_meta.location) { - self.pools - .read() - .get(&object_meta.location) - .unwrap() - .clone() - } else { - let cell = Arc::new(OnceCell::new()); + open: PoolOpen<'_>, + ) -> anyhow::Result { + let cell = Arc::clone( self.pools .write() - .insert(object_meta.location.clone(), cell.clone()); - cell - }; + .entry(object_meta.location.clone()) + .or_default(), + ); let pool = cell .get_or_try_init(|| async { - let open_timer = scan_metrics.open_time.timer(); - let atlas_view = AtlasView::new( - cache, - store.clone(), - object_meta.clone(), - logical_schema.clone(), - ) - .await?; + let open_timer = open.scan_metrics.open_time.timer(); + let atlas_view = + AtlasView::new(open.cache, store, object_meta, open.logical_schema).await?; drop(open_timer); let datasets = atlas_view - .list_datasets(pruning_predicate, scan_metrics.clone()) + .list_datasets(open.predicate, open.scan_metrics.clone()) .await?; // A queue has at least one slot: `ArrayQueue::new(0)` panics. // With no dataset to read the slot stays empty, the first // `pop` finds nothing, and the stream ends at once. let queue = ArrayQueue::new(datasets.len().max(1)); for dataset in datasets { - queue.push(dataset).unwrap(); + queue + .push(dataset) + .map_err(|dataset| anyhow::anyhow!("no slot for dataset '{dataset}'"))?; } Ok::, anyhow::Error>(Arc::new(Level1Pool { inner: Arc::new(InnerLevel1Pool { atlas_view, queue, - projected_schema: projected_schema.clone(), + projected_schema: open.projected_schema, }), })) }) .await .cloned()?; - Ok(pool.as_ref().clone().into_stream(scan_metrics).boxed()) + Ok(pool.as_ref().clone().into_stream(open.scan_metrics)) } } @@ -128,14 +127,14 @@ impl AtlasReaderPool { /// partitions that share a collection each hold a consumer, and the queue /// shares the work between them. #[derive(Clone)] -pub struct Level1Pool { +pub(crate) struct Level1Pool { inner: Arc, } impl Level1Pool { /// One consumer of the queue. It starts on no dataset, and counts the /// datasets it reads on `scan_metrics`. - pub fn into_stream(self, scan_metrics: AtlasScanMetrics) -> Level1PoolStream { + pub(crate) fn into_stream(self, scan_metrics: AtlasScanMetrics) -> Level1PoolStream { Level1PoolStream { inner: self.inner, scan_metrics, @@ -158,7 +157,7 @@ struct InnerLevel1Pool { /// The dataset in hand is a boxed stream, which is `Send` and not `Sync`, so /// the consumer lives with the partition that polls it and never in the /// shared handle. -pub struct Level1PoolStream { +pub(crate) struct Level1PoolStream { inner: Arc, /// The partition's metrics. The datasets this consumer reads count here. scan_metrics: AtlasScanMetrics, @@ -208,11 +207,7 @@ fn dataset_stream( dataset: String, ) -> BoxStream<'static, anyhow::Result> { futures::stream::once(async move { - let source = pool - .atlas_view - .dataset(&dataset) - .await? - .ok_or_else(|| anyhow::anyhow!("dataset '{dataset}' is not in the collection"))?; + let source = pool.atlas_view.dataset(&dataset).await?; scan_metrics.datasets_scanned.add(1); let chunks = source.chunks(); let dataset = Arc::::from(dataset); @@ -221,9 +216,15 @@ fn dataset_stream( let source = Arc::clone(&source); let dataset = Arc::clone(&dataset); async move { - let nd = source.poll_next(chunk).await?.ok_or_else(|| { - anyhow::anyhow!("dataset '{dataset}' read no batch for a chunk of its own grid") - })?; + let nd = source + .poll_next(chunk) + .await + .with_context(|| format!("reading a chunk of dataset '{dataset}'"))? + .ok_or_else(|| { + anyhow::anyhow!( + "dataset '{dataset}' read no batch for a chunk of its own grid" + ) + })?; let nd = under_fields(&nd, pool.atlas_view.table_schema().fields())?; let batch = encode_nd_record_batch(&nd)?.with_schema(Arc::clone(&pool.projected_schema))?; @@ -265,15 +266,18 @@ mod tests { dir: &Path, predicate: Option>, metrics: AtlasScanMetrics, - ) -> BoxStream<'static, anyhow::Result> { + ) -> Level1PoolStream { let (store, marker) = test_support::store_and_marker(dir); let logical = logical_schema(dir).await; let projected = Arc::new(encoded_schema(&logical)); - pool.try_open_into_pooled_stream( - None, store, marker, logical, projected, predicate, metrics, - ) - .await - .unwrap() + let open = PoolOpen { + cache: None, + logical_schema: logical, + projected_schema: projected, + predicate, + scan_metrics: metrics, + }; + pool.open(store, marker, open).await.unwrap() } /// The scan holds the pool in a `FileSource`, which must be `Send + Sync`. diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs index 0dd3d158..e87ca30e 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs @@ -4,17 +4,17 @@ //! //! A collection can hold millions of datasets. Evaluating a predicate against //! each one in turn would cost millions of evaluations. Instead the opener -//! builds one [`PruningIndex`] over the collection: one row per live dataset, +//! builds one `PruningIndex` over the collection: one row per live dataset, //! and one column of typed Arrow statistics per column the predicate names. //! DataFusion's [`PruningPredicate`] then judges the whole collection in one //! vectorised pass, and the result is one bit per dataset. //! -//! The inputs are the opener's own column views. A variable's segment records +//! The inputs are the view's column views. A variable's segment records //! the statistics of every dataset that wrote it, and an attribute view holds //! every dataset's value. Both are in memory once the views exist, so the //! index costs no I/O. Reading the views rather than asking the collection //! again also keeps pruning on the columns the scan reads: a column resolves -//! one way, in [`column_views`](super::opener::column_views). +//! one way, in `column_views`. //! //! # A column the dataset lacks //! @@ -48,7 +48,7 @@ use datafusion::physical_optimizer::pruning::PruningPredicate; use datafusion::scalar::ScalarValue; use indexmap::IndexMap; -use super::opener::AtlasColumnView; +use super::view::AtlasColumnView; /// The datasets of `names` that `predicate` could still match, in order. /// @@ -184,7 +184,7 @@ fn build_index( pack_array_column(segment, names, target) } Some(AtlasColumnView::GlobalAttribute { map }) - | Some(AtlasColumnView::VariableAttribute { map, .. }) => { + | Some(AtlasColumnView::VariableAttribute { map }) => { pack_attribute_column(map, names, target) } }; @@ -356,7 +356,7 @@ mod tests { }; use super::*; - use crate::datafusion::opener::column_views; + use crate::datafusion::view::column_views; use crate::test_support; fn schema(name: &str, data_type: DataType) -> SchemaRef { diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs index 8b56cb71..35dd5b10 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/source.rs @@ -27,51 +27,38 @@ use std::any::Any; use std::sync::Arc; -use std::time::Instant; -use arrow::datatypes::SchemaRef; -use atlas::{Atlas, DatasetView}; -use beacon_nd_array::arrow::{ - file_read::FileRead, metrics::ReadMetrics, partition::FilePartitions, -}; use datafusion::{ config::ConfigOptions, datasource::{ - listing::PartitionedFile, - physical_plan::{FileOpenFuture, FileOpener, FileScanConfig, FileSource}, + physical_plan::{FileOpener, FileScanConfig, FileSource}, table_schema::TableSchema, }, - error::{DataFusionError, Result}, - physical_expr::{ - PhysicalExpr, conjunction, - projection::ProjectionExprs, - utils::{collect_columns, reassign_expr_columns}, - }, + error::Result, + physical_expr::{PhysicalExpr, conjunction, projection::ProjectionExprs}, physical_plan::{ filter_pushdown::{FilterPushdownPropagation, PushedDown}, metrics::ExecutionPlanMetricsSet, }, }; -use futures::{FutureExt, StreamExt, TryStreamExt}; use object_store::ObjectStore; use beacon_datafusion_ext::nd::logical_schema; -use crate::datafusion::{metrics::AtlasScanMetrics, pool::AtlasReaderPool}; -use crate::store::{AtlasReaderCache, get_or_open_atlas}; -use crate::{compat, datafusion::opener::AtlasOpener}; +use crate::datafusion::{metrics::AtlasScanMetrics, opener::AtlasOpener, pool::AtlasReaderPool}; +use crate::store::AtlasReaderCache; /// DataFusion [`FileSource`] for Atlas collections. #[derive(Debug, Clone)] pub struct AtlasSource { table_schema: TableSchema, execution_plan_metrics: ExecutionPlanMetricsSet, - batch_size: usize, predicate: Option>, read_dimensions: Option>, projection: Option, - /// The reader cache to consult, or `None` to open every collection afresh. + /// The reader cache every open goes through. cache: AtlasReaderCache, + /// The scan's pools, one per collection, shared by every partition. reader_pool: Arc, } @@ -84,7 +71,6 @@ impl AtlasSource { Self { table_schema, execution_plan_metrics: ExecutionPlanMetricsSet::new(), - batch_size: usize::MAX, predicate: None, read_dimensions, projection: None, @@ -93,12 +79,6 @@ impl AtlasSource { } } - /// Consult `cache` for opened collections, or open them afresh. - pub fn with_cache(mut self, cache: AtlasReaderCache) -> Self { - self.cache = cache; - self - } - /// Carry a projection the scan pushed down. /// /// The format rebuilds the source in `create_physical_plan`, and without @@ -125,9 +105,7 @@ impl FileSource for AtlasSource { logical_schema: logical_schema(&projected_schema)?, projected_schema, read_dimensions: self.read_dimensions.clone(), - batch_size: self.batch_size, predicate: self.predicate.clone(), - read_metrics: ReadMetrics::new(&self.execution_plan_metrics, partition), scan_metrics: AtlasScanMetrics::new(&self.execution_plan_metrics, partition), reader_pool: Arc::clone(&self.reader_pool), })) @@ -141,11 +119,9 @@ impl FileSource for AtlasSource { &self.table_schema } - fn with_batch_size(&self, batch_size: usize) -> Arc { - Arc::new(Self { - batch_size, - ..self.clone() - }) + /// A batch is one stored chunk, sized by the writer. The scan has no say. + fn with_batch_size(&self, _batch_size: usize) -> Arc { + Arc::new(self.clone()) } /// A container is one unit. A byte range of it names nothing a reader can diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs index 25f0df85..f0cd7717 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs @@ -1,5 +1,28 @@ -use arrow::datatypes::{FieldRef, SchemaRef}; -use atlas::Atlas; +//! One collection, resolved against the table's schema. +//! +//! A column view says where one column of the scan comes from, for every +//! dataset at once: the variable's segment, or the attribute values keyed by +//! dataset. One segment open per array and one attribute sweep per key cost +//! the same however many datasets the collection holds, so a view pays them +//! once, and every dataset reads against the result. Pruning judges through +//! the same views, so the scan and the pruning see one resolution of a name. +//! +//! One dataset is then a lazy [`DatasetSource`] over those views. It reads one +//! stored chunk at a time as an [`NdRecordBatch`], each column on the axes the +//! dataset stores it on, and `under_fields` puts that chunk under the scan's +//! fields. + +use std::sync::Arc; + +use anyhow::Context as _; +use arrow::{ + array::{ArrayRef, new_null_array}, + compute::cast, + datatypes::{Field, FieldRef, Schema, SchemaRef}, +}; +use atlas::{ArrayFile, Atlas, Attr}; +use beacon_datafusion_ext::nd::{Dimensions, NdArrowArray, NdRecordBatch}; +use beacon_datafusion_ext::type_widening::is_type_conflict; use beacon_nd_array::{ NdArrayD, dataset::{default::DefaultDataset, source::DatasetSource}, @@ -7,18 +30,14 @@ use beacon_nd_array::{ use datafusion::physical_plan::PhysicalExpr; use indexmap::IndexMap; use object_store::{ObjectMeta, ObjectStore}; -use std::sync::Arc; use crate::{ compat, - datafusion::{ - metrics::AtlasScanMetrics, - opener::{AtlasColumnView, column_views}, - pruning::prune_datasets, - }, + datafusion::{metrics::AtlasScanMetrics, pruning::prune_datasets}, store::{AtlasReaderCache, get_or_open_atlas}, }; +/// An open collection and the resolution of every column of the table. #[derive(Clone)] pub struct AtlasView { atlas: Arc, @@ -36,7 +55,9 @@ impl AtlasView { table_schema: SchemaRef, ) -> anyhow::Result { let atlas = get_or_open_atlas(cache, store, &object_meta).await?; - let views = column_views(&atlas, &table_schema).await?; + let views = column_views(&atlas, &table_schema) + .await + .with_context(|| format!("resolving the columns of '{}'", object_meta.location))?; Ok(Self { atlas, @@ -50,6 +71,10 @@ impl AtlasView { &self.table_schema } + /// The datasets worth reading, in the collection's order. + /// + /// A dataset the deletion mask hides is not listed. With a predicate, one + /// the statistics rule out is dropped too, and counted on `scan_metrics`. pub async fn list_datasets( &self, pruning_predicate: Option>, @@ -69,24 +94,34 @@ impl AtlasView { Ok(datasets) } - pub async fn dataset( - &self, - dataset_name: &str, - ) -> anyhow::Result>> { + /// One dataset of the collection as a lazy nd dataset, under the table's + /// fields. + /// + /// The dataset holds an array for every field it has: the variable's + /// entry in its segment, read on demand through the atlas backend, or an + /// attribute value on no axis. A field it lacks has no array, and reads + /// as a rank-0 null. No array data is read here. The dataset's layout + /// comes from the segments, and its chunk grid is the one the writer + /// chose. + pub async fn dataset(&self, dataset_name: &str) -> anyhow::Result> { let mut arrays: IndexMap> = IndexMap::new(); for (field, view) in &*self.column_views { let array = match view { None => None, Some(AtlasColumnView::Array { segment }) => match segment.array(dataset_name) { - Some(info) => Some(compat::array_to_nd_array( - Arc::clone(segment), - dataset_name, - &info.dtype, - )?), + Some(info) => Some( + compat::array_to_nd_array(Arc::clone(segment), dataset_name, &info.dtype) + .with_context(|| { + format!( + "reading array '{}' of dataset '{dataset_name}'", + field.name() + ) + })?, + ), None => None, }, Some(AtlasColumnView::GlobalAttribute { map }) - | Some(AtlasColumnView::VariableAttribute { map, .. }) => { + | Some(AtlasColumnView::VariableAttribute { map }) => { // A list has no rank-0 form, and the schema holds no list // column. A dataset that stores a list under a scalar // column's key reads as null. @@ -98,7 +133,301 @@ impl AtlasView { arrays.insert(field.name().clone(), array); } } - let dataset = DefaultDataset::new(dataset_name.to_string(), arrays)?; - Ok(Some(Arc::new(dataset))) + let dataset = DefaultDataset::new(dataset_name.to_string(), arrays) + .with_context(|| format!("laying out dataset '{dataset_name}'"))?; + Ok(Arc::new(dataset)) + } +} + +/// Where one column of the scan comes from, for every dataset of a collection. +/// +/// The scan reads through it, and pruning judges through it, so both see one +/// resolution of a column name. +pub(crate) enum AtlasColumnView { + /// The variable's segment. It holds the array for every dataset that + /// declares it, keyed by dataset name. + Array { segment: Arc }, + /// A dataset-level attribute, its value per dataset. + GlobalAttribute { map: IndexMap }, + /// An attribute of one array, its value per dataset. + VariableAttribute { map: IndexMap }, +} + +/// Where each column of the scan comes from, for every dataset at once. +/// +/// A column no dataset declares gets `None`. +pub(crate) async fn column_views( + atlas: &Atlas, + logical_schema: &Schema, +) -> anyhow::Result>> { + let mut views = IndexMap::with_capacity(logical_schema.fields().len()); + for field in logical_schema.fields() { + let name = field.name(); + let view = if let Some(key) = name.strip_prefix('.') { + let map = atlas + .attributes_by_dataset(None, key) + .await + .with_context(|| format!("sweeping attribute '{key}' for column '{name}'"))?; + Some(AtlasColumnView::GlobalAttribute { map }) + } else if let Some((array, key)) = name.split_once('.') { + let map = atlas + .attributes_by_dataset(Some(array), key) + .await + .with_context(|| format!("sweeping attribute '{key}' for column '{name}'"))?; + Some(AtlasColumnView::VariableAttribute { map }) + } else { + atlas + .try_segment(name) + .await + .with_context(|| format!("opening the segment of column '{name}'"))? + .map(|segment| AtlasColumnView::Array { + segment: Arc::clone(segment), + }) + }; + views.insert(Arc::clone(field), view); + } + Ok(views) +} + +/// `nd` under `fields`: every field in order, on the same target grid. +/// +/// A column comes out under the array's own type, and the table may declare a +/// wider one: that is a cast. A field the dataset lacks is a rank-0 null, +/// which broadcasts to an all-null column. The decoder makes the same of a +/// null struct row, so the scan sees one thing either way. +pub(crate) fn under_fields( + nd: &NdRecordBatch, + fields: &[FieldRef], +) -> anyhow::Result { + let mut columns = Vec::with_capacity(fields.len()); + for field in fields { + let column = match nd.schema().column_with_name(field.name()) { + Some((index, _)) => { + let column = nd.column(index); + match as_field_type(Arc::clone(column.values()), field)? { + Some(values) => NdArrowArray::try_new(values, column.dims().clone())?, + None => null_scalar(field), + } + } + None => null_scalar(field), + }; + columns.push(column); + } + let schema = Arc::new(Schema::new(fields.to_vec())); + Ok(NdRecordBatch::try_new( + schema, + columns, + nd.target().clone(), + )?) +} + +/// A rank-0 null. It broadcasts to an all-null column of the target grid. +fn null_scalar(field: &Field) -> NdArrowArray { + NdArrowArray::try_new(new_null_array(field.data_type(), 1), Dimensions::scalar()) + .expect("one element on no axis") +} + +/// `values` in the type the table declares for `field`, or `None` for values +/// the table cannot hold. +/// +/// A dataset may store a column narrower than the merged type, and the merge +/// widened it: that is a cast. A column the merge could not join is marked, +/// and a dataset of the other family then reads as null. That is what the mark +/// promises the scan. +fn as_field_type(values: ArrayRef, field: &Field) -> anyhow::Result> { + if values.data_type() == field.data_type() { + return Ok(Some(values)); + } + match cast(&values, field.data_type()) { + Ok(values) => Ok(Some(values)), + Err(_) if is_type_conflict(field) => Ok(None), + Err(error) => Err(error) + .with_context(|| format!("casting column '{}' to {}", field.name(), field.data_type())), + } +} + +#[cfg(test)] +mod tests { + use arrow::array::{Array, AsArray, RecordBatch}; + use arrow::datatypes::{Float32Type, Float64Type, Int32Type, Int64Type}; + use beacon_datafusion_ext::type_widening::ArrowTypeWidening; + use std::path::Path; + + use super::*; + use crate::{compat, test_support}; + + /// The schema `infer_schema` derives for a fixture. + async fn schema(dir: &Path) -> SchemaRef { + let atlas = test_support::open(dir).await; + Arc::new( + compat::collection_arrow_schema( + &atlas.footer().collection_schema(), + &ArrowTypeWidening::default_extension(), + ) + .unwrap(), + ) + } + + /// Every chunk of `dataset`, read as the scan reads it: through the view, + /// under the scan's fields. And the rows of all of them in chunk order. + async fn read(dir: &Path, dataset: &str) -> (Vec, RecordBatch) { + let schema = schema(dir).await; + let (store, marker) = test_support::store_and_marker(dir); + let view = AtlasView::new(None, store, marker, Arc::clone(&schema)) + .await + .unwrap(); + let source = view.dataset(dataset).await.unwrap(); + let mut chunks = Vec::new(); + for chunk in source.chunks() { + let nd = source.poll_next(chunk).await.unwrap().unwrap(); + chunks.push(under_fields(&nd, schema.fields()).unwrap()); + } + let batches: Vec = chunks.iter().map(|nd| nd.materialize().unwrap()).collect(); + let batch = arrow::compute::concat_batches(&schema, &batches).unwrap(); + (chunks, batch) + } + + fn column<'a>(batch: &'a RecordBatch, name: &str) -> &'a ArrayRef { + batch + .column_by_name(name) + .unwrap_or_else(|| panic!("no column {name}")) + } + + /// Every column comes out on the dataset's grid. An attribute has no axis + /// of its own, so it repeats on every row. + #[tokio::test] + async fn a_dataset_reads_every_column_on_its_own_grid() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let (nd, batch) = read(tmp.path(), "winter").await; + + assert_eq!(nd.len(), 1, "an unchunked array is one chunk"); + assert_eq!(nd[0].target().shape(), vec![4]); + assert_eq!(batch.num_rows(), 4); + assert_eq!( + column(&batch, "temperature") + .as_primitive::() + .values() + .to_vec(), + vec![1.0, 2.0, 3.0, 4.0] + ); + assert_eq!( + column(&batch, "cycle") + .as_primitive::() + .values() + .to_vec(), + vec![10, 20, 30, 40] + ); + let season = column(&batch, ".season").as_string::(); + assert!( + (0..4).all(|row| season.value(row) == "winter"), + "a rank-0 attribute repeats on every row" + ); + assert_eq!( + column(&batch, ".year") + .as_primitive::() + .values() + .to_vec(), + vec![2024; 4] + ); + assert_eq!( + column(&batch, "temperature.units") + .as_string::() + .value(3), + "celsius" + ); + } + + /// `summer` declares neither `cycle` nor `time`, sets no `year`, and has no + /// `units` on `temperature`. Each is a column of nulls on summer's grid. + #[tokio::test] + async fn a_column_the_dataset_lacks_is_all_null() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + let (_, batch) = read(tmp.path(), "summer").await; + + assert_eq!(batch.num_rows(), 3); + for missing in ["cycle", "time", ".year", "temperature.units"] { + assert_eq!(column(&batch, missing).null_count(), 3, "{missing}"); + } + assert_eq!( + column(&batch, "temperature") + .as_primitive::() + .values() + .to_vec(), + vec![20.0, 21.0, 22.0] + ); + assert_eq!( + column(&batch, ".season").as_string::().value(2), + "summer" + ); + } + + /// A 2-D array reads one stored chunk at a time, keeps both axes, and a + /// cell nobody wrote reads as null. + #[tokio::test] + async fn a_fill_value_reads_as_null_on_a_two_dimensional_grid() { + let tmp = tempfile::tempdir().unwrap(); + test_support::chunked_grid(tmp.path()).await; + + let (nd, batch) = read(tmp.path(), "grid").await; + + assert_eq!(nd.len(), 4, "a [4, 6] grid chunked [2, 3]"); + for chunk in &nd { + assert_eq!(chunk.target().shape(), vec![2, 3]); + } + assert_eq!(batch.num_rows(), 24); + let temperature = column(&batch, "temperature").as_primitive::(); + assert_eq!( + temperature.value(4), + 7.0, + "row 1, column 1 of the grid: the fifth cell of the first chunk" + ); + let mut cells = temperature.values().to_vec(); + cells.sort_by(|a, b| a.partial_cmp(b).unwrap()); + assert_eq!(cells, (0..24).map(f64::from).collect::>()); + let sparse = column(&batch, "sparse"); + assert_eq!(sparse.null_count(), 12, "two of four rows were written"); + assert!( + sparse.is_valid(0), + "the first chunk lies in the written rows" + ); + assert!(sparse.is_null(23), "the last chunk lies outside them"); + } + + /// `a` stores `value` as `Int16` and `b` as `Float32`. The table declares + /// `Float64`, so each dataset casts up to it. `flag` is `a`'s alone. + #[tokio::test] + async fn a_narrower_dataset_casts_to_the_merged_type() { + let tmp = tempfile::tempdir().unwrap(); + test_support::widening(tmp.path()).await; + + let (_, a) = read(tmp.path(), "a").await; + let (_, b) = read(tmp.path(), "b").await; + + assert_eq!( + column(&a, "value") + .as_primitive::() + .values() + .to_vec(), + vec![1.0, 2.0] + ); + assert_eq!( + column(&b, "value") + .as_primitive::() + .values() + .to_vec(), + vec![3.5, 4.5] + ); + assert_eq!( + column(&a, "flag") + .as_primitive::() + .values() + .to_vec(), + vec![7, 8] + ); + assert_eq!(column(&b, "flag").null_count(), 2); } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/store.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/store.rs index 307800ee..f0a27a25 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/store.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/store.rs @@ -138,8 +138,7 @@ struct CacheKey { /// per-runtime state; there is no process-global cache. /// /// Each entry owns a 256 MiB block cache and a 64 MiB I/O cache of its own, so -/// the capacity is a memory bound as much as a handle count. See -/// [`AtlasConfig::reader_cache_size`](crate::AtlasConfig::reader_cache_size). +/// the capacity is a memory bound as much as a handle count. #[derive(Clone)] pub struct AtlasReaderCache { cache: Cache>, @@ -217,7 +216,9 @@ pub async fn get_or_open_atlas( .cache .try_get_with(key, async move { open_collection(store, &path).await }) .await - .map_err(|e: Arc| anyhow::anyhow!("{e}")) + // The cache shares one error between the readers that waited on the + // open, so it cannot be moved out. Its chain survives as text. + .map_err(|e: Arc| anyhow::anyhow!("{e:#}")) } #[cfg(test)] diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs index d6c52fe9..3e181a90 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs @@ -385,200 +385,3 @@ pub async fn declared_unwritten(dir: &Path) { writer.finish().await.expect("finish the collection"); } - -/// A wide collection on two dimensions, `profile` and `level`. -/// -/// Dataset `i` is named `set{i}` and holds `shapes[i]` as -/// `(profiles, levels)`. The pairs differ per dataset, so a row count is a sum -/// and not a product. -/// -/// Four arrays live on `profile` alone: `latitude`, `longitude`, `time` and -/// `platform`. Four live on `profile` and `level`: `pressure`, `temperature`, -/// `salinity` and `quality`. So a read narrowed to `profile` keeps the first -/// four and drops the rest. -/// -/// Every array carries four attributes, so the collection is wide. A predicate -/// column then sits far past the length of a small projection, which is what a -/// projection pushed down after a filter has to be re-indexed against. -/// -/// `temperature` holds `100 * i + level`, so each dataset owns a disjoint -/// range and a predicate over it has a known answer. -pub async fn wide_profiles(dir: &Path, shapes: &[(usize, usize)]) { - let writer = AtlasWriter::create_path(dir, WriterConfig::default()) - .await - .expect("create the collection"); - - for (index, &(profiles, levels)) in shapes.iter().enumerate() { - let mut set = writer - .add_dataset(&format!("set{index}")) - .await - .expect("add a dataset"); - let flat = vec!["profile".to_string()]; - let grid = vec!["profile".to_string(), "level".to_string()]; - - set.define_array::("latitude", flat.clone(), vec![profiles], None, None) - .await - .expect("define latitude"); - set.define_array::("longitude", flat.clone(), vec![profiles], None, None) - .await - .expect("define longitude"); - set.define_array::("time", flat.clone(), vec![profiles], None, None) - .await - .expect("define time"); - set.define_array::("platform", flat, vec![profiles], None, None) - .await - .expect("define platform"); - for name in ["pressure", "temperature", "salinity"] { - set.define_array::(name, grid.clone(), vec![profiles, levels], None, None) - .await - .unwrap_or_else(|e| panic!("define {name}: {e}")); - } - set.define_array::("quality", grid, vec![profiles, levels], None, None) - .await - .expect("define quality"); - - let latitudes: Vec = (0..profiles).map(|p| 10.0 + p as f64).collect(); - let longitudes: Vec = (0..profiles).map(|p| 100.0 + p as f64).collect(); - let times: Vec = (0..profiles) - .map(|p| TimestampNs(EPOCH_NANOS + p as i64 * DAY_NANOS)) - .collect(); - let platforms: Vec = (0..profiles).map(|_| format!("set{index}")).collect(); - set.write_array("latitude", vec![0], arr1(&latitudes).into_dyn().view()) - .await - .expect("write latitude"); - set.write_array("longitude", vec![0], arr1(&longitudes).into_dyn().view()) - .await - .expect("write longitude"); - set.write_array("time", vec![0], arr1(×).into_dyn().view()) - .await - .expect("write time"); - set.write_array("platform", vec![0], arr1(&platforms).into_dyn().view()) - .await - .expect("write platform"); - - let shape = IxDyn(&[profiles, levels]); - let base = 100.0 * index as f32; - let pressure = ArrayD::from_shape_fn(shape.clone(), |i| i[1] as f32); - let temperature = ArrayD::from_shape_fn(shape.clone(), |i| base + i[1] as f32); - let salinity = ArrayD::from_shape_fn(shape.clone(), |i| 30.0 + i[1] as f32); - let quality = ArrayD::from_shape_fn(shape, |_| "1".to_string()); - set.write_array("pressure", vec![0, 0], pressure.view()) - .await - .expect("write pressure"); - set.write_array("temperature", vec![0, 0], temperature.view()) - .await - .expect("write temperature"); - set.write_array("salinity", vec![0, 0], salinity.view()) - .await - .expect("write salinity"); - set.write_array("quality", vec![0, 0], quality.view()) - .await - .expect("write quality"); - - // The attributes make the collection wide. Four per array, so the - // column count is well past the arrays alone. - for name in [ - "latitude", - "longitude", - "time", - "platform", - "pressure", - "temperature", - "salinity", - "quality", - ] { - set.set_array_attribute(name, "long_name", Attr::String(format!("the {name}"))) - .expect("set long_name"); - set.set_array_attribute(name, "units", Attr::String("1".into())) - .expect("set units"); - set.set_array_attribute(name, "valid_min", Attr::Float64(-1000.0)) - .expect("set valid_min"); - set.set_array_attribute(name, "valid_max", Attr::Float64(1000.0)) - .expect("set valid_max"); - } - set.set_attribute("title", Attr::String("wide profiles".into())); - set.set_attribute("institution", Attr::String("test".into())); - set.finish().await.expect("finish a dataset"); - } - - writer.finish().await.expect("finish the collection"); -} - -/// A number column that one dataset stores as text a cast cannot read. -/// -/// - `a`: `value: Float64[2] = [1.5, 2.5]`. -/// - `b`: `value: String[2] = ["0.-90", "3.5"]`. -/// -/// `Float64` and `String` share no type, so `KeepFirst` keeps `Float64` and -/// marks the column. `'0.-90'` is not a number, so its cast has to read as -/// null. `'3.5'` casts cleanly, which separates "the cast ran" from "the cast -/// gave up on the column". -pub async fn conflicting_numbers(dir: &Path) { - let writer = AtlasWriter::create_path(dir, WriterConfig::default()) - .await - .expect("create the collection"); - - { - let mut a = writer.add_dataset("a").await.expect("add a"); - a.define_array::("value", vec!["obs".into()], vec![2], None, None) - .await - .expect("define value"); - a.write_array("value", vec![0], arr1(&[1.5f64, 2.5]).into_dyn().view()) - .await - .expect("write value"); - a.finish().await.expect("finish a"); - } - - { - let mut b = writer.add_dataset("b").await.expect("add b"); - b.define_array::("value", vec!["obs".into()], vec![2], None, None) - .await - .expect("define value"); - b.write_array( - "value", - vec![0], - arr1(&["0.-90".to_string(), "3.5".to_string()]) - .into_dyn() - .view(), - ) - .await - .expect("write value"); - b.finish().await.expect("finish b"); - } - - writer.finish().await.expect("finish the collection"); -} - -/// One dataset, `d`, holding `value` as a number or as text. -/// -/// Two collections built this way put the conflict across the *collection* -/// merge rather than inside one collection's own merge. -pub async fn value_typed(dir: &Path, as_text: bool) { - let writer = AtlasWriter::create_path(dir, WriterConfig::default()) - .await - .expect("create the collection"); - let mut d = writer.add_dataset("d").await.expect("add d"); - if as_text { - d.define_array::("value", vec!["obs".into()], vec![2], None, None) - .await - .expect("define value"); - d.write_array( - "value", - vec![0], - arr1(&["0.-30".to_string(), "3.5".to_string()]) - .into_dyn() - .view(), - ) - .await - .expect("write value"); - } else { - d.define_array::("value", vec!["obs".into()], vec![2], None, None) - .await - .expect("define value"); - d.write_array("value", vec![0], arr1(&[1.5f64, 2.5]).into_dyn().view()) - .await - .expect("write value"); - } - d.finish().await.expect("finish d"); - writer.finish().await.expect("finish the collection"); -} From ad88fe0b6433b970dbe3cc623d1d2fb2c32d04d1 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 9 Sep 2026 14:34:19 +0200 Subject: [PATCH 15/16] atlas: count(*) drives on the widest array, docs match the code A read that projects no column loaded no array, so every dataset sat on a rank-0 grid and `count(*)` returned the dataset count. The view now drives such a read with each dataset's widest array, and a chunk states its row count through `DatasetSource::chunk_rows`, so a count reads no cell. The changelog and the format, tuning, configuration and external-table pages named four `BEACON_ATLAS_*` settings and two metrics that no longer exist. They now describe the reader cache, pruning and the metrics as the code has them. The beacon-core atlas tests referenced the removed `AtlasConfig`. The pruning test now compares a predicate against a full read filtered in memory. --- CHANGELOG.md | 17 ++-- beacon-db/beacon-core/tests/atlas.rs | 50 +++++----- .../tests/common/iceberg_fixture.rs | 2 +- .../beacon-arrow-atlas/src/datafusion/mod.rs | 51 +++++++++-- .../beacon-arrow-atlas/src/datafusion/pool.rs | 51 ++++++++--- .../beacon-arrow-atlas/src/datafusion/view.rs | 91 +++++++++++++++++++ .../beacon-nd-array/src/dataset/default.rs | 19 ++++ .../beacon-nd-array/src/dataset/source.rs | 9 ++ docs/docs/2.0.0-rc5/formats/atlas.md | 12 +-- docs/docs/2.0.0-rc5/server/configuration.md | 9 -- .../2.0.0-rc5/server/performance-tuning.md | 35 ++----- .../2.0.0-rc5/sql/create-external-table.md | 2 +- 12 files changed, 246 insertions(+), 102 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bc7edf2a..68ef984c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,15 +22,11 @@ tag. Releases before 2.0.0 are recorded in the before 0.17 is not read at all — rewrite it with `atlas create`. Three behaviour changes come with the rebuild. A dataset-level attribute is a column under a leading dot, `".platform"`, matching NetCDF and Zarr instead of the bare key. A scan reads through the shared nd pipeline, - so one dataset is one unit of work and a collection divides over every core — and a dataset - stored in several chunks divides further, so one large dataset still uses all of them. And a - column two - datasets type in two families now refuses the merge by name rather than silently becoming text; - `BEACON_TYPE_WIDENING_ON_CONFLICT=keep_first` settles it the other way. Atlas collections are - also crawlable now, because a collection is one file whose extension is its format. - `BEACON_ATLAS_USE_READER_CACHE`, `BEACON_ATLAS_READER_CACHE_SIZE`, `BEACON_ATLAS_USE_PRUNING` - and `BEACON_ATLAS_ENABLE_STATISTICS` configure it, and the same keys work per table through - `OPTIONS`. + so one dataset is one unit of work: every partition of a query helps drain every collection, + and a dataset is read once. And a column two datasets type in two families now refuses the + merge by name rather than silently becoming text; `BEACON_TYPE_WIDENING_ON_CONFLICT=keep_first` + settles it the other way. Atlas collections are also crawlable now, because a collection is + one file whose extension is its format. - **A predicate over an Atlas collection skips whole datasets.** Atlas records the minimum, the maximum and the null count of every array, so a collection can be judged before it is read. @@ -43,7 +39,8 @@ tag. Releases before 2.0.0 are recorded in the too, and at the same cost. Pruning only ever removes datasets that hold no matching row: every path falls back to reading everything, and the filter above the scan still decides each row. `EXPLAIN ANALYZE` reports it as - `atlas_datasets_pruned`, `atlas_index_builds` and `atlas_index_rows`. + `atlas_datasets_pruned` and `atlas_datasets_scanned`, with the time spent as `atlas_open_time` + and `atlas_prune_time`. - **`BEACON_TYPE_WIDENING_ON_CONFLICT` settles a column that no type holds.** A collection can type one column as a number in one file and as a string in another. No type holds both, so the diff --git a/beacon-db/beacon-core/tests/atlas.rs b/beacon-db/beacon-core/tests/atlas.rs index 11b95505..5e9c071e 100644 --- a/beacon-db/beacon-core/tests/atlas.rs +++ b/beacon-db/beacon-core/tests/atlas.rs @@ -13,7 +13,7 @@ mod common; use std::path::Path; use beacon_arrow_atlas::atlas::{AtlasWriter, Attr, WriterConfig}; -use common::{scalar_i64, total_rows, TestRuntime}; +use common::{TestRuntime, scalar_i64, total_rows}; use ndarray::arr1; /// Write a collection of `n` datasets at `dir`, named `d0..d{n-1}`. @@ -163,42 +163,36 @@ async fn read_atlas_takes_a_dimension_list() { // ── pruning, through the assembled runtime ────────────────────────────── -/// A predicate returns the same rows whether or not whole datasets were -/// skipped to find them. +/// A predicate returns the same rows as a full read filtered afterwards, so +/// skipping whole datasets to find them changes no answer. #[tokio::test(flavor = "multi_thread")] async fn pruning_does_not_change_the_answer() { - let rt = common::runtime_with("atlas-pruning-on", |builder| { - builder.with_atlas_config(beacon_arrow_atlas::AtlasConfig { - use_pruning: true, - ..Default::default() - }) - }) - .await; - let unpruned = common::runtime_with("atlas-pruning-off", |builder| { - builder.with_atlas_config(beacon_arrow_atlas::AtlasConfig { - use_pruning: false, - ..Default::default() - }) - }) + let rt = common::runtime("atlas-pruning").await; + write_collection(&rt.datasets_dir().join("obs"), 10).await; + let all = temperatures( + &rt, + "SELECT temperature FROM read_atlas('obs/data.atlas') ORDER BY temperature", + ) .await; - - for rt in [&rt, &unpruned] { - write_collection(&rt.datasets_dir().join("obs"), 10).await; - } - - for predicate in [ - "temperature > 45", - "temperature < 25", - "temperature > 1000", - "temperature > 45 AND temperature < 75", - ] { + assert_eq!(all.len(), 40, "ten datasets of four rows"); + + let cases: [(&str, fn(f32) -> bool); 4] = [ + ("temperature > 45", |t| t > 45.0), + ("temperature < 25", |t| t < 25.0), + ("temperature > 1000", |t| t > 1000.0), + ("temperature > 45 AND temperature < 75", |t| { + t > 45.0 && t < 75.0 + }), + ]; + for (predicate, keep) in cases { let sql = format!( "SELECT temperature FROM read_atlas('obs/data.atlas') \ WHERE {predicate} ORDER BY temperature" ); + let expected: Vec = all.iter().copied().filter(|&t| keep(t)).collect(); assert_eq!( temperatures(&rt, &sql).await, - temperatures(&unpruned, &sql).await, + expected, "pruning changed the answer for `{predicate}`" ); } diff --git a/beacon-db/beacon-core/tests/common/iceberg_fixture.rs b/beacon-db/beacon-core/tests/common/iceberg_fixture.rs index f1345ff3..5697b871 100644 --- a/beacon-db/beacon-core/tests/common/iceberg_fixture.rs +++ b/beacon-db/beacon-core/tests/common/iceberg_fixture.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use datafusion::prelude::SessionContext; use iceberg::io::LocalFsStorageFactory; -use iceberg::memory::{MemoryCatalog, MemoryCatalogBuilder, MEMORY_CATALOG_WAREHOUSE}; +use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalog, MemoryCatalogBuilder}; use iceberg::spec::{NestedField, PrimitiveType, Schema, Type}; use iceberg::table::Table; use iceberg::transaction::{AddColumn, ApplyTransactionAction, Transaction}; diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs index f414b280..0574736d 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs @@ -489,13 +489,15 @@ mod deal_tests { mod scan_tests { use super::*; use crate::test_support; + use arrow::array::AsArray; + use arrow::datatypes::Int64Type; use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig}; use datafusion::prelude::{SessionConfig, SessionContext}; use std::path::Path; - /// The rows of every collection under `dir`, read in a session of - /// `partitions` target partitions, and the partition count of the plan. - async fn rows(dir: &Path, partitions: usize) -> (usize, usize) { + /// A table over every collection under `dir`, in a session of + /// `partitions` target partitions. + async fn table(dir: &Path, partitions: usize) -> (SessionContext, Arc) { // The collections sit in subdirectories, which a listing skips by // default. let config = SessionConfig::new() @@ -514,8 +516,22 @@ mod scan_tests { .infer_schema(&ctx.state()) .await .unwrap(); - let table = Arc::new(ListingTable::try_new(config).unwrap()); + (ctx, Arc::new(ListingTable::try_new(config).unwrap())) + } + + /// Three collections of five datasets, four rows each, under `dir`. + async fn three_collections(dir: &Path) { + for name in ["a", "b", "c"] { + let dir = dir.join(name); + std::fs::create_dir_all(&dir).unwrap(); + test_support::ranged(&dir, 5).await; + } + } + /// The rows of every collection under `dir`, read in a session of + /// `partitions` target partitions, and the partition count of the plan. + async fn rows(dir: &Path, partitions: usize) -> (usize, usize) { + let (ctx, table) = table(dir, partitions).await; let df = ctx.read_table(table).unwrap(); let plan = df.clone().create_physical_plan().await.unwrap(); let partition_count = plan.properties().output_partitioning().partition_count(); @@ -528,11 +544,7 @@ mod scan_tests { #[tokio::test] async fn every_partition_reads_through_the_pool_and_no_dataset_twice() { let tmp = tempfile::tempdir().unwrap(); - for name in ["a", "b", "c"] { - let dir = tmp.path().join(name); - std::fs::create_dir_all(&dir).unwrap(); - test_support::ranged(&dir, 5).await; - } + three_collections(tmp.path()).await; let (alone, one) = rows(tmp.path(), 1).await; let (shared, three) = rows(tmp.path(), 3).await; @@ -549,4 +561,25 @@ mod scan_tests { "a dataset is read by one partition and no other" ); } + + /// A count projects no column. It still counts every row of every + /// dataset, once, across the partitions, and reads no cell to do so. + #[tokio::test] + async fn a_count_reads_no_column_and_counts_every_row() { + let tmp = tempfile::tempdir().unwrap(); + three_collections(tmp.path()).await; + let (ctx, table) = table(tmp.path(), 3).await; + ctx.register_table("obs", table).unwrap(); + + let batches = ctx + .sql("SELECT count(*) FROM obs") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let count = batches[0].column(0).as_primitive::().value(0); + assert_eq!(count, 60, "three collections of five datasets of four rows"); + } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs index a9c59798..806caf8d 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs @@ -8,6 +8,7 @@ //! that poll at once drain different datasets, so a dataset is read by one //! partition and by no other. +use std::any::Any; use std::collections::HashMap; use std::fmt::Debug; use std::pin::Pin; @@ -15,9 +16,10 @@ use std::sync::Arc; use std::task::{Context, Poll}; use anyhow::Context as _; -use arrow::array::RecordBatch; +use arrow::array::{RecordBatch, RecordBatchOptions}; use arrow::datatypes::SchemaRef; -use beacon_datafusion_ext::nd::encode_nd_record_batch; +use beacon_datafusion_ext::nd::{NdRecordBatch, encode_nd_record_batch}; +use beacon_nd_array::dataset::source::DatasetSource; use crossbeam::queue::ArrayQueue; use datafusion::physical_plan::PhysicalExpr; use futures::stream::BoxStream; @@ -216,15 +218,16 @@ fn dataset_stream( let source = Arc::clone(&source); let dataset = Arc::clone(&dataset); async move { - let nd = source - .poll_next(chunk) - .await - .with_context(|| format!("reading a chunk of dataset '{dataset}'"))? - .ok_or_else(|| { - anyhow::anyhow!( - "dataset '{dataset}' read no batch for a chunk of its own grid" - ) - })?; + // A read of no column counts rows. A chunk that states its + // rows is not read at all. + if pool.projected_schema.fields().is_empty() { + let rows = match source.chunk_rows(&chunk) { + Some(rows) => rows, + None => read_chunk(&source, chunk, &dataset).await?.num_rows(), + }; + return count_batch(Arc::clone(&pool.projected_schema), rows); + } + let nd = read_chunk(&source, chunk, &dataset).await?; let nd = under_fields(&nd, pool.atlas_view.table_schema().fields())?; let batch = encode_nd_record_batch(&nd)?.with_schema(Arc::clone(&pool.projected_schema))?; @@ -236,6 +239,32 @@ fn dataset_stream( .boxed() } +/// One chunk of `source`, read. +async fn read_chunk( + source: &Arc, + chunk: Arc, + dataset: &str, +) -> anyhow::Result { + source + .poll_next(chunk) + .await + .with_context(|| format!("reading a chunk of dataset '{dataset}'"))? + .ok_or_else(|| { + anyhow::anyhow!("dataset '{dataset}' read no batch for a chunk of its own grid") + }) +} + +/// The batch a read of no column emits for `rows` rows: the count, and +/// nothing else. The decoder reads a zero-column batch's row count as its +/// payload. +fn count_batch(schema: SchemaRef, rows: usize) -> anyhow::Result { + Ok(RecordBatch::try_new_with_options( + schema, + Vec::new(), + &RecordBatchOptions::new().with_row_count(Some(rows)), + )?) +} + #[cfg(test)] mod tests { use super::*; diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs index f0cd7717..21760bac 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs @@ -37,12 +37,23 @@ use crate::{ store::{AtlasReaderCache, get_or_open_atlas}, }; +/// The segment of every array of a collection, by name. +type DrivingSegments = Arc<[(String, Arc)]>; + /// An open collection and the resolution of every column of the table. #[derive(Clone)] pub struct AtlasView { atlas: Arc, table_schema: SchemaRef, column_views: Arc>>, + /// The segment of every array, by name, for a read that projects no + /// column. + /// + /// Such a read counts rows. With no array read, every dataset would sit + /// on a rank-0 grid of one row. The read is driven by the widest array of + /// each dataset instead, so the count is the dataset's full grid. `None` + /// when the table projects a column. + driving: Option, } impl AtlasView { @@ -58,11 +69,20 @@ impl AtlasView { let views = column_views(&atlas, &table_schema) .await .with_context(|| format!("resolving the columns of '{}'", object_meta.location))?; + let driving = if table_schema.fields().is_empty() { + let segments = driving_segments(&atlas) + .await + .with_context(|| format!("resolving the arrays of '{}'", object_meta.location))?; + Some(segments) + } else { + None + }; Ok(Self { atlas, table_schema, column_views: Arc::new(views), + driving, }) } @@ -133,6 +153,11 @@ impl AtlasView { arrays.insert(field.name().clone(), array); } } + if let Some(segments) = &self.driving + && let Some((name, array)) = widest_array(segments, dataset_name) + { + arrays.insert(name, array); + } let dataset = DefaultDataset::new(dataset_name.to_string(), arrays) .with_context(|| format!("laying out dataset '{dataset_name}'"))?; Ok(Arc::new(dataset)) @@ -189,6 +214,50 @@ pub(crate) async fn column_views( Ok(views) } +/// The segment of every array of the collection, by name. +async fn driving_segments(atlas: &Atlas) -> anyhow::Result { + let names: Vec = atlas + .footer() + .collection_schema() + .arrays + .keys() + .map(|name| name.to_string()) + .collect(); + let mut segments = Vec::with_capacity(names.len()); + for name in names { + let segment = atlas + .try_segment(&name) + .await + .with_context(|| format!("opening the segment of array '{name}'"))?; + if let Some(segment) = segment { + segments.push((name, Arc::clone(segment))); + } + } + Ok(segments.into()) +} + +/// The array of `dataset` with the most cells, out of `segments`, with its +/// name. An array Beacon cannot read is passed over: it drives nothing. +fn widest_array( + segments: &[(String, Arc)], + dataset: &str, +) -> Option<(String, Arc)> { + let mut widest: Option<(usize, String, Arc)> = None; + for (name, segment) in segments { + let Some(info) = segment.array(dataset) else { + continue; + }; + let Ok(array) = compat::array_to_nd_array(Arc::clone(segment), dataset, &info.dtype) else { + continue; + }; + let cells: usize = array.shape().iter().product(); + if widest.as_ref().is_none_or(|(most, _, _)| cells > *most) { + widest = Some((cells, name.clone(), array)); + } + } + widest.map(|(_, name, array)| (name, array)) +} + /// `nd` under `fields`: every field in order, on the same target grid. /// /// A column comes out under the array's own type, and the table may declare a @@ -339,6 +408,28 @@ mod tests { ); } + /// A read of no column is driven by each dataset's widest array, so a + /// count sees the dataset's full grid without reading a cell. + #[tokio::test] + async fn a_read_of_no_column_counts_the_full_grid() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let (store, marker) = test_support::store_and_marker(tmp.path()); + let view = AtlasView::new(None, store, marker, Arc::new(Schema::empty())) + .await + .unwrap(); + + for (dataset, rows) in [("winter", 4), ("summer", 3)] { + let source = view.dataset(dataset).await.unwrap(); + let counted: usize = source + .chunks() + .iter() + .map(|chunk| source.chunk_rows(chunk).unwrap()) + .sum(); + assert_eq!(counted, rows, "{dataset}"); + } + } + /// `summer` declares neither `cycle` nor `time`, sets no `year`, and has no /// `units` on `temperature`. Each is a column of nulls on summer's grid. #[tokio::test] diff --git a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs index d9d44f30..1f1ce303 100644 --- a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs +++ b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs @@ -88,6 +88,13 @@ impl DatasetSource for DefaultDataset { .collect() } + /// The cells of the chunk's grid. A scalar dataset has one. + fn chunk_rows(&self, chunk: &Arc) -> Option { + chunk + .downcast_ref::() + .map(|subset| subset.shape.iter().product()) + } + /// Read one chunk into an un-broadcast [`NdRecordBatch`]. /// /// `chunk` is one of [`DatasetSource::chunks`]. Each array is sliced on @@ -301,6 +308,18 @@ mod tests { } } + /// A chunk states its own row count, so a count reads nothing. + #[test] + fn a_chunk_states_its_rows() { + let ds = gridded(&[2, 2]); + let rows: Vec = ds + .chunks() + .iter() + .map(|chunk| ds.chunk_rows(chunk).unwrap()) + .collect(); + assert_eq!(rows, vec![4, 2, 4, 2], "a [4, 3] grid chunked [2, 2]"); + } + /// Poll every chunk, broadcast each, and stitch the rows back together. async fn read_all(ds: &DefaultDataset) -> arrow::record_batch::RecordBatch { let mut batches = Vec::new(); diff --git a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/source.rs b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/source.rs index 99012612..a08aedb0 100644 --- a/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/source.rs +++ b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/source.rs @@ -6,6 +6,15 @@ use std::{any::Any, sync::Arc}; pub trait DatasetSource: Send + Sync + std::fmt::Debug { fn chunks(&self) -> Vec>; + /// How many rows `chunk` holds, when that is known without a read. + /// + /// A read that projects no column counts rows and reads nothing else. + /// `None` says the count is not known ahead: read the chunk and count + /// what comes back. + fn chunk_rows(&self, _chunk: &Arc) -> Option { + None + } + async fn poll_next( &self, chunk: Arc, diff --git a/docs/docs/2.0.0-rc5/formats/atlas.md b/docs/docs/2.0.0-rc5/formats/atlas.md index 404a09f8..20d54d6a 100644 --- a/docs/docs/2.0.0-rc5/formats/atlas.md +++ b/docs/docs/2.0.0-rc5/formats/atlas.md @@ -69,10 +69,9 @@ What Beacon does with that: vectorised pass and never opens the ones that cannot match. Judging a column costs one request, whatever the dataset count. A dataset-level attribute is exact, so `WHERE ".platform" = 'p3'` prunes on it too. -- **One dataset is one unit of work.** A collection's datasets are spread across every core, and a - worker takes the next one when it is free, so a collection of a million small datasets and one of - four large ones both divide evenly. A dataset stored in several chunks divides further, so a - single large dataset still uses every core. +- **One dataset is one unit of work.** Every core reads from every collection of a query. A core + takes the next dataset when it is free, and a dataset is read once, so a collection of a million + small datasets and one of four large ones both divide evenly. - **Column projection.** Only the arrays a query names get read, and only their attributes are fetched. - **Object storage.** A collection reads from local disk, S3, GCS, Azure and HTTP alike. @@ -185,14 +184,11 @@ See [Create External Tables](/docs/2.0.0-rc5/data-sources/external-tables) for t ### `OPTIONS` -`STORED AS ATLAS` reads four keys: +`STORED AS ATLAS` reads one key: | Option | Type | Default | Description | | --- | --- | --- | --- | | `read_dimensions` | List of dimension names | The default grid of each dataset | The dimensions the table reads. An array survives only when the list holds every one of its own. | -| `use_pruning` | Boolean | `true` (`BEACON_ATLAS_USE_PRUNING`) | Drop the datasets a predicate rules out before reading them. Turning it off only costs speed: pruning never changes an answer. | -| `use_reader_cache` | Boolean | `true` (`BEACON_ATLAS_USE_READER_CACHE`) | Reuse an opened collection across queries. | -| `enable_statistics` | Boolean | `true` (`BEACON_ATLAS_ENABLE_STATISTICS`) | Whether `ANALYZE FILES` records this collection's column ranges. A query never measures a collection, so this affects the analyzer alone. | ```sql CREATE EXTERNAL TABLE sensor_atlas diff --git a/docs/docs/2.0.0-rc5/server/configuration.md b/docs/docs/2.0.0-rc5/server/configuration.md index cf86a7be..2bb00598 100644 --- a/docs/docs/2.0.0-rc5/server/configuration.md +++ b/docs/docs/2.0.0-rc5/server/configuration.md @@ -304,15 +304,6 @@ small stores, where even a rank-1 read per store adds up. `valid_min` and `valid_max` are never used as a range. They state which values are *valid*, not which values a store holds, so a store may hold values outside them. -### Atlas - -| Variable | Default | Description | -| --- | --- | --- | -| `BEACON_ATLAS_USE_READER_CACHE` | `true` | Keep opened Atlas collections in memory, so a repeated query does not read the footer of `data.atlas` again. | -| `BEACON_ATLAS_READER_CACHE_SIZE` | `32` | How many opened collections to keep. Each holds its own 256 MiB block cache and 64 MiB slab cache, so this bounds memory as well as handles. | -| `BEACON_ATLAS_USE_PRUNING` | `true` | Skip the datasets a predicate rules out from the collection's own statistics, before reading them. Off only costs speed; pruning never changes an answer. | -| `BEACON_ATLAS_ENABLE_STATISTICS` | `true` | Let `ANALYZE FILES` record a collection's column ranges. They come from the footer, so they cost no array read. | - ### Beacon Binary Format (BBF) | Variable | Default | Description | diff --git a/docs/docs/2.0.0-rc5/server/performance-tuning.md b/docs/docs/2.0.0-rc5/server/performance-tuning.md index cd1a8123..77639a1f 100644 --- a/docs/docs/2.0.0-rc5/server/performance-tuning.md +++ b/docs/docs/2.0.0-rc5/server/performance-tuning.md @@ -218,31 +218,16 @@ with statistics, next to the chunk pruning. It drops whole datasets before it re ## Atlas Tuning Beacon opens an [Atlas](/docs/2.0.0-rc5/formats/atlas) collection by reading the footer of its -`data.atlas` file. Beacon caches the open collections. It therefore does not read that footer again -for every query. +`data.atlas` file. Each table keeps its open collections in a cache of 512 entries, so a query +does not read that footer again. Each cached collection holds its own block cache, 256 MiB of +decompressed blocks and 64 MiB of raw slabs, so the cache is a memory bound as much as a handle +count. There is no setting for it. -### Reader cache (no repeated store open) +### Dataset pruning -#### `BEACON_ATLAS_USE_READER_CACHE` and `BEACON_ATLAS_READER_CACHE_SIZE` +Pruning is always on. A query with a predicate judges every dataset of a collection against the +statistics in its footer, in one pass, and never opens the ones that cannot match. Pruning never +changes an answer, and the filter above the scan still decides every row. -With the reader cache on, Beacon uses an open collection again. It therefore does not read the -footer for every query, and the decompressed blocks of a collection stay warm between them. - -Each cached collection holds its own block cache — 256 MiB of decompressed blocks and 64 MiB of raw -slabs — so the size is a memory bound as much as a handle count. - -Recommendations: - -- Keep `BEACON_ATLAS_USE_READER_CACHE=true`, the default, when several queries read the same Atlas - collections. -- Increase `BEACON_ATLAS_READER_CACHE_SIZE`, default `32`, if you query more Atlas collections than - the cache holds. Remember the memory bound above before raising it far. - -#### `BEACON_ATLAS_USE_PRUNING` - -On by default. A query with a predicate judges every dataset of a collection against the statistics -in its footer, in one pass, and never opens the ones that cannot match. Turning it off only costs -speed: pruning never changes an answer, and the filter above the scan still decides every row. - -`EXPLAIN ANALYZE` reports what it did as `atlas_datasets_pruned`, `atlas_datasets_scanned` and -`atlas_index_rows`. +`EXPLAIN ANALYZE` reports what it did as `atlas_datasets_pruned` and `atlas_datasets_scanned`, with +the time spent as `atlas_open_time` and `atlas_prune_time`. diff --git a/docs/docs/2.0.0-rc5/sql/create-external-table.md b/docs/docs/2.0.0-rc5/sql/create-external-table.md index 218ffdb5..bcc8099a 100644 --- a/docs/docs/2.0.0-rc5/sql/create-external-table.md +++ b/docs/docs/2.0.0-rc5/sql/create-external-table.md @@ -165,7 +165,7 @@ validates it and then reads the server setting alone. See the format page of eac | `NC` | `read_dimensions`, `use_rust_reader`, `enable_statistics` | [NetCDF](/docs/2.0.0-rc5/formats/netcdf#options) | | `HDF5`, `H5` | `read_dimensions`, `use_rust_reader`, `enable_statistics`, `unify_phony_dimensions`, `convention` | [HDF5](/docs/2.0.0-rc5/formats/hdf5#options) | | `ZARR` | `read_dimensions`, `enable_statistics` | [Zarr](/docs/2.0.0-rc5/formats/zarr#options) | -| `ATLAS` | `read_dimensions`, `use_pruning`, `use_reader_cache`, `enable_statistics` | [Atlas](/docs/2.0.0-rc5/formats/atlas#options) | +| `ATLAS` | `read_dimensions` | [Atlas](/docs/2.0.0-rc5/formats/atlas#options) | | `CSV` | `delimiter`, `infer_records` | [CSV](/docs/2.0.0-rc5/formats/csv#options) | | `BBF` | `split_streams_slice` | [BBF](/docs/2.0.0-rc5/formats/bbf#options) | | `DELTA` | `version`, `timestamp` | [Delta Lake](/docs/2.0.0-rc5/formats/delta-lake#options) | From 34ee818980866870f931f5aabf57e9d5c883a6a4 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 9 Sep 2026 14:34:51 +0200 Subject: [PATCH 16/16] Restore an import order rustfmt changed in the iceberg test fixture --- beacon-db/beacon-core/tests/common/iceberg_fixture.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/beacon-db/beacon-core/tests/common/iceberg_fixture.rs b/beacon-db/beacon-core/tests/common/iceberg_fixture.rs index 5697b871..f1345ff3 100644 --- a/beacon-db/beacon-core/tests/common/iceberg_fixture.rs +++ b/beacon-db/beacon-core/tests/common/iceberg_fixture.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use datafusion::prelude::SessionContext; use iceberg::io::LocalFsStorageFactory; -use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalog, MemoryCatalogBuilder}; +use iceberg::memory::{MemoryCatalog, MemoryCatalogBuilder, MEMORY_CATALOG_WAREHOUSE}; use iceberg::spec::{NestedField, PrimitiveType, Schema, Type}; use iceberg::table::Table; use iceberg::transaction::{AddColumn, ApplyTransactionAction, Transaction};