diff --git a/CHANGELOG.md b/CHANGELOG.md index 814ebe25..86063a25 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,36 @@ tag. Releases before 2.0.0 are recorded in the ### Added +- **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.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: 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. + 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. 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` and `atlas_datasets_scanned`, with the time spent as `atlas_open_time` + and `atlas_prune_time`. + - **The server root is a home page instead of a jump to Swagger.** `http://localhost:5001/` sent every visitor straight to the Swagger UI, which hid the admin panel, the API reference and the documentation from anyone who did not know their paths. The root now answers with a small page diff --git a/Cargo.lock b/Cargo.lock index b48d5802..9cbdd5c4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -287,6 +287,28 @@ dependencies = [ "password-hash", ] +[[package]] +name = "array-format" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "672b40e04f7c403e9322972c476414f8847e5e0aed0006b80128977d39138ab9" +dependencies = [ + "bytes", + "ecow", + "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 +1191,32 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "atlas-rust" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbb4b330af5674f3b1ee4e7ac7892d0c21c1719f9f9ec2ca3b38198f16c55466" +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", + "smallvec", + "smol_str 0.3.6", + "tempfile", + "thiserror 2.0.20", + "tokio", + "tracing", + "zstd", +] + [[package]] name = "atoi" version = "2.0.0" @@ -1392,6 +1440,31 @@ 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", + "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", +] + [[package]] name = "beacon-arrow-bbf" version = "2.0.0-rc.5" @@ -1723,6 +1796,7 @@ dependencies = [ "async-stream", "async-trait", "base64", + "beacon-arrow-atlas", "beacon-arrow-bbf", "beacon-arrow-csv", "beacon-arrow-geoparquet", @@ -1756,6 +1830,7 @@ dependencies = [ "glob", "iceberg", "iceberg-datafusion", + "ndarray 0.17.2", "num_cpus", "object_store 0.13.2", "parking_lot", @@ -1880,6 +1955,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 +2195,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", @@ -5064,6 +5141,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" @@ -11071,6 +11154,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" @@ -11086,6 +11172,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 f39e5465..307e4a7f 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.17.2" + 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 2b0adf7a..92841be9 100644 --- a/beacon-db/beacon-core/src/runtime_builder.rs +++ b/beacon-db/beacon-core/src/runtime_builder.rs @@ -1,11 +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::{AtlasFormatFactory, AtlasOptions}; use beacon_arrow_bbf::datafusion::BBFFormatFactory; use beacon_arrow_csv::datafusion::CsvFormatFactory; use beacon_arrow_geoparquet::datafusion::GeoParquetFormatFactory; @@ -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; @@ -797,6 +797,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())), Arc::new(BBFFormatFactory::new(Default::default())), Arc::new(GeoParquetFormatFactory::default()), Arc::new(NetCDFFormatFactory::new( @@ -858,7 +859,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/atlas.rs b/beacon-db/beacon-core/tests/atlas.rs new file mode 100644 index 00000000..5e9c071e --- /dev/null +++ b/beacon-db/beacon-core/tests/atlas.rs @@ -0,0 +1,215 @@ +//! 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::{TestRuntime, scalar_i64, total_rows}; +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 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("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; + 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, + expected, + "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-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); +} diff --git a/beacon-db/beacon-datafusion-ext/src/nd/encoding.rs b/beacon-db/beacon-datafusion-ext/src/nd/encoding.rs index dfab0335..931403c8 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()) } @@ -312,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 7cee86e4..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_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-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-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..fad3dbde 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,9 @@ tracing = { workspace = true } tokio = { workspace = true } ndarray = { workspace = true } chrono = { workspace = true } -crossbeam = { workspace = true } moka = { workspace = true } -atlas-rust = "0.14.0" +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/backend.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/backend.rs index f8d9efde..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,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 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::{ArrayFile, 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 `dataset`'s entry in `segment` from `start`. async fn read( - view: &atlas::DatasetView, - array_name: &str, + segment: &ArrayFile, + dataset: &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, + segment: &ArrayFile, + dataset: &str, start: Vec, shape: Vec, ) -> anyhow::Result> { - let arr = view - .read_array::<$ty>(array_name, start, shape) + let values = segment + .read_array::<$ty>(dataset, 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, - view.name() + "Failed to read dataset '{dataset}' from its atlas segment: {e}" ) })?; - 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, + segment: &ArrayFile, + dataset: &str, start: Vec, shape: Vec, ) -> anyhow::Result> { - let arr = view - .read_array::(array_name, start, shape) + let values = segment + .read_array::(dataset, 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, - view.name() + "Failed to read the timestamps of dataset '{dataset}' from its atlas segment: {e}" ) })?; - // 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 dataset's entry of an atlas segment lazily, one 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 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 { - atlas: Arc, - dataset_name: String, - array_name: String, + segment: Arc, + dataset: String, shape: Vec, dimensions: Vec, chunk_shape: Vec, @@ -135,39 +125,44 @@ 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.dataset) .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, - shape: Vec, - dimensions: Vec, - chunk_shape: Vec, - fill_value: Option, - ) -> Self { - Self { - atlas, - dataset_name, - array_name, +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, - } + }) } } #[async_trait::async_trait] -impl ArrayBackend for AtlasArrayBackend { +impl ArrayBackend for AtlasArrayBackend { fn len(&self) -> usize { self.shape.iter().product() } @@ -180,6 +175,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 +188,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.segment, &self.dataset, 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 +233,210 @@ impl ArrayBackend for AttributeBackend { #[cfg(test)] mod tests { use super::*; - use crate::reader::test_support::build_two_dataset_store; - use atlas::Atlas; + use crate::test_support; - // ── AttributeBackend ─────────────────────────────────────────────── + /// 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 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()); + async fn the_backend_reports_what_the_segment_holds() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + 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), + vec!["obs".to_string()] + ); + assert_eq!(ArrayBackend::::chunk_shape(&backend), vec![4]); + assert_eq!(ArrayBackend::::fill_value(&backend), Some(-1)); + 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 attribute_backend_read_subset_returns_value() { - let backend = AttributeBackend::new(42i32); - let arr = backend - .read_subset(ArraySubset { - start: vec![], - shape: vec![], - }) + async fn a_full_read_returns_every_value() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + + 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 - .expect("read"); - assert_eq!(arr.ndim(), 0); - let raw = arr.into_raw_vec_and_offset().0; - assert_eq!(raw, vec![42i32]); + .unwrap(); + assert_eq!(values.into_raw_vec_and_offset().0, vec![1.0, 2.0, 3.0, 4.0]); } - // ── AtlasReadable::fill_element ──────────────────────────────────── + #[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::::try_new( + segment(tmp.path(), "cycle").await, + "winter".to_string(), + ) + .unwrap(); + 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]); + } - #[test] - fn fill_element_passthrough_numeric() { - use atlas::FillValue; + /// 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::::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 + .read_subset(ArraySubset::new(vec![1, 2], vec![2, 2])) + .await + .unwrap(); + assert_eq!(values.shape(), &[2, 2]); assert_eq!( - ::fill_element(Some(&FillValue::Int(-7))), - -7i32 + values.into_raw_vec_and_offset().0, + vec![8.0, 9.0, 14.0, 15.0] ); - 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); } - // ── AtlasArrayBackend ────────────────────────────────────────────── + /// 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::::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 + .unwrap(); + assert_eq!(values.into_raw_vec_and_offset().0, vec![-999.0; 3]); + } #[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"); - - let backend = AtlasArrayBackend::::new( - Arc::new(atlas), - "winter".into(), - "temperature".into(), - vec![4], - vec!["obs".into()], - vec![4], - Some(-1.0f32), - ); + 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::::try_new( + segment(tmp.path(), "time").await, + "winter".to_string(), + ) + .unwrap(); + let values = backend + .read_subset(ArraySubset::new(vec![0], vec![2])) + .await + .unwrap(); assert_eq!( - as ArrayBackend>::shape(&backend), - vec![4] + 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 a_string_array_reads_its_values() { + let tmp = tempfile::tempdir().unwrap(); + test_support::incompatible(tmp.path()).await; + + 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 + .unwrap(); assert_eq!( - as ArrayBackend>::dimensions(&backend), - vec!["obs".to_string()] + 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!( - as ArrayBackend>::chunk_shape(&backend), - vec![4] + ::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!( - as ArrayBackend>::fill_value(&backend), - Some(-1.0f32) + ::fill_element(Some(&FillValue::TimestampNs( + i64::MIN + ))), + TimestampNanosecond(i64::MIN) ); - assert_eq!(backend.len(), 4); } - #[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"); - - let backend = AtlasArrayBackend::::new( - Arc::new(atlas), - "winter".into(), - "temperature".into(), - vec![4], - vec!["obs".into()], - vec![4], - None, - ); - let arr = backend - .read_subset(ArraySubset { - start: vec![0], - shape: vec![4], - }) - .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]); - } + // ── AttributeBackend ──────────────────────────────────────────────── #[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], - None, - ); - let arr = backend - .read_subset(ArraySubset { - start: vec![1], - shape: vec![2], - }) + 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..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 @@ -1,184 +1,264 @@ -//! 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::collections::BTreeMap; 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, Field, Schema}; +use arrow::error::ArrowError; +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, AtlasReadable, AttributeBackend}; +use crate::backend::{AtlasArrayBackend, 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) +} + +/// 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) } -/// Build the Arrow schema for a whole atlas store from its collection-wide -/// [`MergedSchema`] — **no per-dataset iteration and no disk I/O**. +/// 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) +} + +/// A stable tag for a dtype, for keys that group datasets by shape. +pub(crate) fn dtype_tag(dtype: &DType) -> String { + format!("{dtype:?}") +} + +// ─── Collection schema ─────────────────────────────────────────────────────── + +/// The Arrow schema of one collection, from its footer alone. /// -/// 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. +/// 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. /// -/// 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)); - } +/// 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); } } - 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)); + 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() +} - fields.sort_by(|a, b| a.name().cmp(b.name())); - Schema::new(fields) +/// 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()) } -/// Convert an atlas array (described by its [`ArraySchema`]) into a lazy -/// [`NdArrayD`] backed by [`AtlasArrayBackend`]. +// ─── Lazy arrays ───────────────────────────────────────────────────────────── + +/// Wrap one dataset's entry of an atlas segment as a lazy [`NdArrayD`]. /// -/// `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. +/// No array data is read here. `dtype` comes from the collection footer, and +/// 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. /// -/// `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, - array_name: &str, - schema: &ArraySchema, - fill_value: Option, + segment: Arc, + dataset: &str, + dtype: &DType, ) -> anyhow::Result> { - let shape = schema.shape.clone(); - let dimensions = schema.dimension_names.clone(); - let chunk_shape = schema.chunk_shape.clone(); - - macro_rules! mk { + macro_rules! lazy { ($ty:ty) => {{ - let fill: Option<$ty> = fill_value - .as_ref() - .map(|fv| <$ty as AtlasReadable>::fill_element(Some(fv))); - let backend = AtlasArrayBackend::<$ty>::new( - atlas.clone(), - dataset_name.to_string(), - array_name.to_string(), - shape.clone(), - dimensions.clone(), - chunk_shape.clone(), - fill, - ); - let nd = NdArray::new_with_backend(backend)?; - Ok::, anyhow::Error>(Arc::new(nd)) + let backend = AtlasArrayBackend::<$ty>::try_new(segment, dataset.to_string())?; + 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), + 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 + "dataset '{dataset}' holds a Bool array, 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 + "dataset '{dataset}' holds a FixedSizeList array, 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 + "dataset '{dataset}' holds a List array, 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, + ) }; } @@ -196,21 +276,9 @@ 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)), - 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,141 +286,291 @@ 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"); + } + + #[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"); + } + + // ── 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:?}"); } } - 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) + /// `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 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}"); + #[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:?}"); + } + } + + /// 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 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}"); + 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_bool_round_trips() { - let nd = attribute_to_nd_array(&Attr::Bool(true)).expect("convert"); + async fn a_bool_attribute_is_a_column() { + let nd = attribute_to_nd_array(&Attr::Bool(true)).unwrap(); 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]); } - #[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]); + #[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}"); } - #[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()]); + // ── 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 attribute_timestamp_round_trips() { - let nanos = 1_700_000_000_000_000_000i64; - let nd = attribute_to_nd_array(&Attr::TimestampNanoseconds(nanos)).expect("convert"); - assert_eq!(nd.datatype(), NdArrayDataType::Timestamp); - let typed = nd - .as_any() - .downcast_ref::>() - .expect("downcast"); + 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!( - typed.clone_into_raw_vec().await, - vec![TimestampNanosecond(nanos)] + names(&schema), + vec![ + ".season", + ".year", + "cycle", + "temperature", + "temperature.units", + "time" + ] ); } #[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 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 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*. + /// 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 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"); + 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(); - 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" + schema.field_with_name("value").unwrap().data_type(), + &DataType::Float64, + "Int16 and Float32 widen to Float64" ); - - let schema = atlas_merged_schema_to_arrow(&merged, None); assert_eq!( - schema.field_with_name("value").expect("value field").data_type(), - &DataType::Utf8 + 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}" ); - // The column only one dataset declares still appears, at its own type. + } + + /// 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").expect("only_a field").data_type(), + 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/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/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/metrics.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/metrics.rs index e0f25738..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 @@ -1,31 +1,20 @@ -//! Execution metrics for the atlas scan, surfaced through DataFusion's standard -//! metrics reporting (e.g. `EXPLAIN ANALYZE`). -//! -//! 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. - 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. pub prune_time: Time, - /// Wall time building lazy datasets — metadata, backends, projected - /// attribute values, and the per-dataset schema adapter. 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, } 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..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 @@ -1,140 +1,77 @@ -//! 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 anyhow::Context as _; +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::util::{ATLAS_MARKER, top_level_atlas_markers}; +use crate::compat; +use crate::datafusion::error::external; +use crate::store::{ATLAS_MARKER, AtlasReaderCache, get_or_open_atlas, top_level_atlas_markers}; -pub mod cache; +pub(crate) mod error; 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 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::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 -} - -/// Parse a boolean value supplied through a `CREATE EXTERNAL TABLE` option. -fn parse_bool_option(key: &str, value: &str) -> datafusion::error::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 ───────────────────────────────────────────────────────────────── +/// The name this format answers to: `STORED AS ATLAS`, `read_atlas`. +pub const ATLAS_FORMAT: &str = "atlas"; +/// 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`. - 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 } } - /// Build an [`AtlasFormat`] with the given per-table effective settings, - /// wiring in the shared reader cache when caching is enabled. - fn build_format( - &self, - options: AtlasOptions, - use_reader_cache: bool, - use_pruning: bool, - ) -> AtlasFormat { - let cache = use_reader_cache.then(|| self.cache.clone()); + /// 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) - .with_cache(cache) - .with_pruning(use_pruning) } } @@ -142,39 +79,24 @@ 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") { - use_reader_cache = parse_bool_option("use_reader_cache", value)?; - } - if let Some(value) = format_options.get("use_pruning") { - use_pruning = parse_bool_option("use_pruning", value)?; - } - - Ok(Arc::new(self.build_format(options, use_reader_cache, use_pruning))) + Ok(Arc::new(self.build(options))) } fn default(&self) -> Arc { - Arc::new(self.build_format( - 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 { @@ -184,85 +106,97 @@ 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. /// - /// 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. + /// 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. + 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 plain format. Atlas measures no column. + /// + /// 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, - 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> { + Ok(Arc::new(AtlasFormat::default())) } } -// ─── 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. - cache: Option, - /// Whether a predicate scan prunes non-matching datasets before reading. - use_pruning: bool, + cache: AtlasReaderCache, +} + +impl Default for AtlasFormat { + fn default() -> Self { + Self::new(AtlasOptions::default()) + } } impl AtlasFormat { pub fn new(options: AtlasOptions) -> Self { Self { options, - cache: None, - use_pruning: false, + cache: AtlasReaderCache::new(512), } } +} - /// Wire in a reader cache (`Some`) or disable caching (`None`). - pub fn with_cache(mut self, cache: Option) -> Self { - self.cache = cache; - self - } - - /// Enable or disable dataset pruning for predicate scans. - pub fn with_pruning(mut self, use_pruning: bool) -> Self { - self.use_pruning = use_pruning; - 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 +209,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,127 +217,149 @@ 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. + /// + /// 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, 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. - let read_dimensions = self.options.read_dimensions.clone(); - let mut schemas = Vec::new(); + // One rule for both merges: the datasets inside a collection, and the + // collections of this table. + let widening = session_widening(state); + + 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 atlas = get_or_open_atlas(Some(&self.cache), Arc::clone(store), marker) + .await + .map_err(external)?; + 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. + compat::collection_arrow_schema(&atlas.footer().collection_schema(), &widening) + .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(), )); } - // 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) } + /// Unknown for every column. + /// + /// 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, table_schema: SchemaRef, _object: &ObjectMeta, - ) -> datafusion::error::Result { + ) -> Result { Ok(Statistics::new_unknown(&table_schema)) } + /// 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 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, 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(); - let object_store = state - .runtime_env() - .object_store(conf.object_store_url.clone())?; - let target_partitions = state.config().target_partitions().max(1); + ) -> Result> { + beacon_nd_array::arrow::morsel::reject_partition_columns("Atlas", &conf)?; - 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(); - for marker in &markers { - let atlas = get_or_open_atlas(self.cache.as_ref(), object_store.clone(), marker).await?; - 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])); - } - } + // 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!( - elapsed_ms = plan_start.elapsed().as_millis() as u64, - stores = markers.len(), - datasets = total_datasets, + collections = markers.len(), partitions = file_groups.len(), - "atlas create_physical_plan partitioning", + "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) - .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)) .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, + self.cache.clone(), + )) } async fn create_writer_physical_plan( @@ -411,678 +368,218 @@ 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") - } - - /// 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 - }) - .await - .clone() - } - - /// 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()) - } - - /// `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, - } +/// 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 tests { - use super::test_support::{fixture_marker_object_meta, test_store}; +mod deal_tests { 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; - - // ── discovery ─────────────────────────────────────────────────────── - - #[test] - fn factory_get_ext_is_atlas() { - let factory = AtlasFormatFactory::new(Default::default(), Default::default()); - assert_eq!(factory.get_ext(), "atlas"); - assert_eq!(factory.file_format_name(), "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(), + 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, - }, - ObjectMeta { - location: OsPath::from("store_b/atlas.msgpack"), - 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 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:?}"); - } - } - - #[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()); - } - - #[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(); - - 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(); - } - - #[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; - - 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); - } - - /// 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. - #[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); - - 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 batches = ctx - .sql("SELECT temperature FROM atlas_fast") - .await - .unwrap() + }) .collect() - .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); } - #[tokio::test] - async fn projection_prunes_columns_through_datafusion() { - let ctx = SessionContext::new(); - register_example(&ctx).await; - - let df = ctx.sql("SELECT temperature FROM atlas_t").await.unwrap(); - let names: Vec = df - .schema() - .fields() + /// The collections of a group, by their directory. + fn dealt(group: &FileGroup) -> Vec { + group + .files() .iter() - .map(|f| f.name().clone()) - .collect(); - assert_eq!(names, vec!["temperature".to_string()]); - } - - #[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() + .map(|file| { + file.object_meta + .location + .parts() + .next() + .unwrap() + .as_ref() + .to_string() + }) .collect() - .await - .unwrap() - .iter() - .map(|b| b.num_rows()) - .sum(); - assert_eq!(rows, 0, "no temperature exceeds 1e6"); } - #[tokio::test] - async fn count_star_counts_every_dataset_row() { - let ctx = SessionContext::new(); - register_example(&ctx).await; - - use arrow::array::Int64Array; - let batches = ctx - .sql("SELECT COUNT(*) AS n FROM atlas_t") - .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"); + /// The collection each group starts on. + fn starts(groups: &[FileGroup]) -> Vec { + groups.iter().map(|group| dealt(group)[0].clone()).collect() } - // ── cross-dataset dtype widening: cast + null-fill ────────────────── + #[test] + fn every_partition_holds_every_collection_in_its_own_rotation() { + let groups = deal_rotated(&markers(4), 2); - /// 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 + assert_eq!(groups.len(), 2); + assert_eq!(dealt(&groups[0]), ["c0", "c1", "c2", "c3"]); + assert_eq!(dealt(&groups[1]), ["c2", "c3", "c0", "c1"]); } - #[tokio::test] - async fn widened_array_dtype_is_cast_from_each_dataset() { - use arrow::array::Float32Array; - use arrow::datatypes::DataType; - - let ctx = SessionContext::new(); - let _tmp = register_widening(&ctx).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(); + #[test] + fn the_starts_spread_evenly_over_the_collections() { assert_eq!( - df.schema().field_with_unqualified_name("value").unwrap().data_type(), - &DataType::Float32, - "merged value column must be the widened super-type" + starts(&deal_rotated(&markers(3), 3)), + ["c0", "c1", "c2"], + "one start per collection" ); - - 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)); - } - } - // 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"); - } - - #[tokio::test] - async fn missing_column_is_null_filled_per_dataset() { - let ctx = SessionContext::new(); - let _tmp = register_widening(&ctx).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") - .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"); - } - - /// 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; - - let ctx = SessionContext::new(); - 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()) - .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" + starts(&deal_rotated(&markers(4), 3)), + ["c0", "c1", "c2"], + "fewer partitions start as far apart as they can" ); - - 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 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)); } - // ── partition splitting ───────────────────────────────────────────── - #[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()); + 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" + ); } +} - // ── dataset pruning toggle ────────────────────────────────────────── - - /// 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) +#[cfg(test)] +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; + + /// 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() + .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(); - ctx.register_table("ranged", Arc::new(ListingTable::try_new(config).unwrap())) - .unwrap(); - tmp - } - - 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}")) - .await - .unwrap() - .collect() - .await - .unwrap() - .iter() - .map(|b| b.num_rows()) - .sum() + (ctx, Arc::new(ListingTable::try_new(config).unwrap())) } - #[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); - } - - #[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; - - 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(); - - 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 + /// 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; } - - // 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"); } - #[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 - .sql("SELECT temperature FROM ranged WHERE temperature > 45") - .await - .unwrap() - .collect() - .await - .unwrap() - .iter() - .map(|b| b.num_rows()) - .sum(); - assert_eq!(rows, 20, "d5..d9 × 4 rows, across 8 partitions"); + /// 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(); + 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 scan_metrics_report_pruned_and_scanned_counts() { - use datafusion::physical_plan::metrics::MetricsSet; - use datafusion::physical_plan::{ExecutionPlan, collect}; - - let ctx = SessionContext::new(); - let _tmp = register_ranged(&ctx, 10, true).await; - let plan = ctx - .sql("SELECT temperature FROM ranged WHERE temperature > 45") - .await - .unwrap() - .create_physical_plan() - .await - .unwrap(); - collect(plan.clone(), ctx.task_ctx()).await.unwrap(); + async fn every_partition_reads_through_the_pool_and_no_dataset_twice() { + let tmp = tempfile::tempdir().unwrap(); + three_collections(tmp.path()).await; - // 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()); + let (alone, one) = rows(tmp.path(), 1).await; + let (shared, three) = rows(tmp.path(), 3).await; - // `> 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()); + 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" + ); } + /// 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 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; + 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 temperature FROM ranged WHERE temperature > 45 ORDER BY temperature") + .sql("SELECT count(*) FROM obs") .await .unwrap() .collect() .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)); - } - } - assert!(vals.iter().all(|v| *v > 45.0)); - assert_eq!(vals.len(), 20); - } - - #[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() - .await - .unwrap() - .iter() - .map(|b| b.num_rows()) - .sum(); - assert_eq!(rows, 7, "partitioned scan must not drop or duplicate rows"); + 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/opener.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs new file mode 100644 index 00000000..b4183d54 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/opener.rs @@ -0,0 +1,223 @@ +//! One partition's opener: a collection in, nd batches out. +//! +//! 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. 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::datatypes::SchemaRef; +use datafusion::{ + datasource::{ + listing::PartitionedFile, + physical_plan::{FileOpenFuture, FileOpener}, + }, + error::Result, + physical_plan::PhysicalExpr, +}; +use futures::{FutureExt, StreamExt, TryStreamExt}; +use object_store::ObjectStore; + +use crate::{ + datafusion::{ + error::external, + metrics::AtlasScanMetrics, + pool::{AtlasReaderPool, PoolOpen}, + }, + 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 predicate: Option>, + 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 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(); + 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 pool = Arc::clone(&self.reader_pool); + + let fut = async move { + let open = PoolOpen { + cache: Some(&cache), + logical_schema, + projected_schema, + predicate, + scan_metrics, + }; + let stream = pool + .open(store, file.object_meta, open) + .await + .map_err(external)?; + Ok(stream.map_err(external).boxed()) + }; + + Ok(fut.boxed()) + } +} + +#[cfg(test)] +mod tests { + 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; + + use super::*; + use crate::{compat, test_support}; + + fn column<'a>(batch: &'a RecordBatch, name: &str) -> &'a ArrayRef { + batch + .column_by_name(name) + .unwrap_or_else(|| panic!("no column {name}")) + } + + // ── 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, + predicate: None, + scan_metrics: AtlasScanMetrics::new(&metrics, 0), + reader_pool: Arc::new(AtlasReaderPool::new()), + }; + (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/options.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/options.rs index 04782656..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 @@ -1,10 +1,16 @@ -/// 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. +/// +/// 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/pool.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs new file mode 100644 index 00000000..806caf8d --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pool.rs @@ -0,0 +1,461 @@ +//! 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::any::Any; +use std::collections::HashMap; +use std::fmt::Debug; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use anyhow::Context as _; +use arrow::array::{RecordBatch, RecordBatchOptions}; +use arrow::datatypes::SchemaRef; +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; +use futures::{Stream, StreamExt, TryStreamExt}; +use object_store::{ObjectMeta, ObjectStore, path::Path}; +use parking_lot::RwLock; +use tokio::sync::OnceCell; + +use crate::datafusion::metrics::AtlasScanMetrics; +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>, +} + +/// 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 { + 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, 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, + store: Arc, + object_meta: ObjectMeta, + open: PoolOpen<'_>, + ) -> anyhow::Result { + let cell = Arc::clone( + self.pools + .write() + .entry(object_meta.location.clone()) + .or_default(), + ); + let pool = cell + .get_or_try_init(|| async { + 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(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) + .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: open.projected_schema, + }), + })) + }) + .await + .cloned()?; + + Ok(pool.as_ref().clone().into_stream(open.scan_metrics)) + } +} + +/// 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(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(crate) 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(crate) 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?; + 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 { + // 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))?; + Ok(batch) + } + })) + }) + .try_flatten() + .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::*; + 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, + ) -> Level1PoolStream { + let (store, marker) = test_support::store_and_marker(dir); + let logical = logical_schema(dir).await; + let projected = Arc::new(encoded_schema(&logical)); + 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`. + /// 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 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] + 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/pruning.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/pruning.rs index 61fc168b..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 @@ -1,479 +1,737 @@ -//! Optional dataset-level predicate pruning. +//! Dropping the datasets a predicate cannot match, from what is in memory. //! -//! 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. 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. +//! +//! 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`. +//! +//! # 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, 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::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; use datafusion::physical_expr::utils::collect_columns; use datafusion::physical_optimizer::pruning::PruningPredicate; use datafusion::scalar::ScalarValue; +use indexmap::IndexMap; -/// 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)] -pub enum CandidateFilter { - /// Pruning didn't apply (fail-open) — keep every dataset. - KeepAll, - /// Only these dataset names could match; everything else is prunable. - Only(HashSet), -} +use super::view::AtlasColumnView; -impl CandidateFilter { - /// Restrict `names` to the datasets that survive pruning, preserving order. - pub fn retain(&self, names: Vec) -> Vec { - match self { - CandidateFilter::KeepAll => names, - CandidateFilter::Only(set) => { - names.into_iter().filter(|n| set.contains(n)).collect() - } +/// 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; + } + + // 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(); + }; + + 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() } } } -/// Per-store memo of the [`CandidateFilter`], so a store is pruned **once** per -/// query rather than once per scan partition. -/// -/// 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. -#[derive(Clone)] -pub struct PruneCache { - cache: moka::future::Cache>, +// ─── The index ─────────────────────────────────────────────────────────────── + +/// One column's statistics, one row per dataset. +struct StatColumn { + min: ArrayRef, + max: ArrayRef, + null_count: ArrayRef, + row_count: ArrayRef, } -impl PruneCache { - pub fn new() -> Self { - Self { - cache: moka::future::Cache::builder().max_capacity(256).build(), - } +/// 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)) } - /// Return the memoized filter for `key`, computing it via `init` on first - /// use. Concurrent callers for the same key share 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 + fn max_values(&self, column: &Column) -> Option { + self.columns.get(column.name()).map(|c| Arc::clone(&c.max)) } -} -impl Default for PruneCache { - fn default() -> Self { - Self::new() + fn null_counts(&self, column: &Column) -> Option { + self.columns + .get(column.name()) + .map(|c| Arc::clone(&c.null_count)) } -} -impl std::fmt::Debug for PruneCache { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("PruneCache").finish_non_exhaustive() + fn row_counts(&self, column: &Column) -> Option { + self.columns + .get(column.name()) + .map(|c| Arc::clone(&c.row_count)) } -} -/// 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, + fn num_containers(&self) -> usize { + self.rows } -} -/// 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) + 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 + } } -async fn try_candidates( - 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 - }; - - let referenced = collect_columns(pruning_predicate.orig_expr()); - if referenced.is_empty() { - return Ok(None); - } +// ─── Building it ───────────────────────────────────────────────────────────── - // 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 +/// 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(|col| column_key(&merged, col.name()).map(|k| (col.name().to_string(), k))) + .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) + } + Some(AtlasColumnView::GlobalAttribute { map }) + | Some(AtlasColumnView::VariableAttribute { map }) => { + pack_attribute_column(map, names, target) + } + }; + Some((column.clone(), packed)) + }) .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, + PruningIndex { + rows: names.len(), columns, - }; - 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()); - } } - 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. +/// One array column, from the segment that holds the variable. /// -/// 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, - 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))); - } 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); +/// 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 (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; + }; + 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); } - let min = ScalarValue::iter_to_array(mins).ok()?; - let max = ScalarValue::iter_to_array(maxes).ok()?; - Some(PackedStats { - min, - max, + 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)), - }) + } } -/// 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)); - } - } +/// One attribute column. +/// +/// An attribute's value is exact, so it is both the minimum and the maximum of +/// 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( + values: &IndexMap, + names: &[String], + target: &DataType, +) -> 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 row_counts: Vec> = vec![None; rows]; + + 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).and_then(|scalar| scalar.cast_to(target).ok()) + else { + continue; + }; + bounds[row] = scalar; + null_counts[row] = Some(0); + row_counts[row] = Some(1); + } + + 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)), } - 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 { +/// 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(|_| 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())), + _ => 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 arrow::datatypes::{Field, Schema}; + use atlas::Atlas; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{ + BinaryExpr, Column as ColumnExpr, IsNullExpr, Literal, + }; -impl PruningStatistics for AtlasPruningStatistics { - fn min_values(&self, column: &Column) -> Option { - self.columns.get(column.name()).map(|c| c.min.clone()) - } + use super::*; + use crate::datafusion::view::column_views; + use crate::test_support; - fn max_values(&self, column: &Column) -> Option { - self.columns.get(column.name()).map(|c| c.max.clone()) + fn schema(name: &str, data_type: DataType) -> SchemaRef { + Arc::new(Schema::new(vec![Field::new(name, data_type, true)])) } - fn null_counts(&self, column: &Column) -> Option { - self.columns.get(column.name()).map(|c| c.null_count.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 row_counts(&self, column: &Column) -> Option { - self.columns.get(column.name()).map(|c| c.row_count.clone()) + fn is_null(column: &str) -> Arc { + Arc::new(IsNullExpr::new(Arc::new(ColumnExpr::new(column, 0)))) } - fn num_containers(&self) -> usize { - self.num_containers + /// 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 views = Arc::new(column_views(atlas, &schema).await.unwrap()); + prune_datasets(&views, atlas.list_datasets(), &predicate, &schema).await } - fn contained( - &self, - _column: &Column, - _values: &HashSet, - ) -> Option { - None - } -} + // ── the index over array statistics ───────────────────────────────── -#[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; + /// 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 survivors = kept( + &atlas, + binary( + "temperature", + Operator::Gt, + ScalarValue::Float32(Some(45.0)), + ), + schema("temperature", DataType::Float32), + ) + .await; + assert_eq!(survivors, vec!["d5", "d6", "d7", "d8", "d9"]); + } - 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()); } + /// 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 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 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 survivors = kept( + &atlas, + binary( + "temperature", + Operator::GtEq, + ScalarValue::Float32(Some(0.0)), + ), + schema("temperature", DataType::Float32), + ) + .await; + assert_eq!(survivors, vec!["d1", "d2", "d3", "d4", "d5"]); } + // ── 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 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_mixed_dtype_column_is_cast_before_it_is_compared() { + let tmp = tempfile::tempdir().unwrap(); + 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); + + // a holds [1, 2] and b holds [3.5, 4.5]. + assert_eq!( + kept( + &atlas, + binary("value", Operator::Gt, ScalarValue::Float64(Some(3.0))), + Arc::clone(&schema) + ) + .await, + vec!["b"] + ); + assert_eq!( + kept( + &atlas, + binary("value", Operator::Lt, ScalarValue::Float64(Some(3.0))), + Arc::clone(&schema) + ) + .await, + vec!["a"] + ); + assert!( + 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 + /// what is in memory alone. #[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 an_attribute_predicate_prunes() { 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| { + test_support::ranged(tmp.path(), 6).await; + let atlas = test_support::open(tmp.path()).await; + + let survivors = kept( + &atlas, binary( - col("value", &schema).unwrap(), - Operator::Gt, - lit(ScalarValue::Float32(Some(t))), - &schema, + ".platform", + Operator::Eq, + ScalarValue::Utf8(Some("p3".to_string())), + ), + schema(".platform", DataType::Utf8), + ) + .await; + assert_eq!(survivors, vec!["d3"]); + } + + // ── 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_missing_attribute_reads_as_null_and_prunes_as_null() { + let tmp = tempfile::tempdir().unwrap(); + test_support::two_datasets(tmp.path()).await; + let atlas = test_support::open(tmp.path()).await; + let schema = schema(".year", DataType::Int64); + + assert_eq!( + kept( + &atlas, + binary(".year", Operator::Eq, ScalarValue::Int64(Some(2024))), + Arc::clone(&schema) ) - .unwrap() - }; - let names = atlas.list_datasets(); // [a, b] + .await, + vec!["winter"] + ); + assert_eq!(kept(&atlas, is_null(".year"), schema).await, vec!["summer"]); + } + + /// 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_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); - // `> 3`: a (max 2, cast from Int16) is pruned; b (max 4.5) survives. assert_eq!( - retain_candidates(&atlas, names.clone(), &value_gt(3.0), &schema).await, - vec!["b"] + kept( + &atlas, + binary("flag", Operator::Gt, ScalarValue::Int32(Some(5))), + Arc::clone(&schema) + ) + .await, + vec!["a"] ); - // `> 10`: neither can match. assert!( - retain_candidates(&atlas, names.clone(), &value_gt(10.0), &schema) - .await - .is_empty() + 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() ); - // `> 1.5`: a's cast max (2.0) still qualifies, so both survive. assert_eq!( - retain_candidates(&atlas, names.clone(), &value_gt(1.5), &schema).await, - names + 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 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, + 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; + + let survivors = kept( + &atlas, + binary("value", Operator::Eq, ScalarValue::Int32(Some(0))), + schema("value", DataType::Int32), + ) + .await; + 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 survivors = kept( + &atlas, + binary("temperature", Operator::Gt, ScalarValue::Float32(Some(0.0))), + schema("temperature", DataType::Float32), + ) + .await; + assert!(survivors.is_empty(), "nothing in, nothing out"); + } + + // ── 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, + ); + 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!( + attr_to_scalar(&Attr::String("p1".into())), + Some(ScalarValue::Utf8(Some("p1".into()))) + ); + } + + /// 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); } } 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..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 @@ -1,119 +1,88 @@ -//! 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 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 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 +//! +//! 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, 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 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 beacon_datafusion_ext::scan_adapt::batch_adapter_factory; use datafusion::{ config::ConfigOptions, datasource::{ - listing::PartitionedFile, - physical_plan::{FileOpenFuture, FileOpener, FileScanConfig, FileSource}, - schema_adapter::SchemaAdapterFactory, + physical_plan::{FileOpener, FileScanConfig, FileSource}, table_schema::TableSchema, }, + error::Result, physical_expr::{PhysicalExpr, conjunction, projection::ProjectionExprs}, physical_plan::{ filter_pushdown::{FilterPushdownPropagation, PushedDown}, metrics::ExecutionPlanMetricsSet, }, }; -use futures::future; -use futures::{StreamExt, TryStreamExt, stream::BoxStream}; -use object_store::{ObjectMeta, ObjectStore}; +use object_store::ObjectStore; -use crate::datafusion::cache::AtlasReaderCache; -use crate::datafusion::metrics::AtlasScanMetrics; -use crate::datafusion::pruning::PruneCache; +use beacon_datafusion_ext::nd::logical_schema; -/// 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::{metrics::AtlasScanMetrics, opener::AtlasOpener, pool::AtlasReaderPool}; +use crate::store::AtlasReaderCache; -/// The slice of a store's dataset names assigned to one scan partition. -/// -/// 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. -#[derive(Debug, Clone)] -pub struct AtlasDatasetSlice { - pub names: Vec, -} - -/// 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. - cache: Option, - /// Whether to prune non-matching datasets before reading them. - 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. - prune_cache: PruneCache, - /// Projection pushed down by the scan, applied on top of the table schema. projection: Option, + /// The reader cache every open goes through. + cache: AtlasReaderCache, + /// The scan's pools, one per collection, shared by every partition. + reader_pool: Arc, } 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 { - schema_adapter_factory: None, table_schema, execution_plan_metrics: ExecutionPlanMetricsSet::new(), - batch_size: usize::MAX, predicate: None, read_dimensions, - cache: None, - use_pruning: false, - prune_cache: PruneCache::new(), projection: None, + cache, + reader_pool: Arc::new(AtlasReaderPool::new()), } } - /// 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. - pub fn with_cache(mut self, cache: Option) -> Self { - self.cache = cache; - self - } - - /// Enable or disable dataset pruning for this scan. - 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 @@ -126,24 +95,23 @@ impl FileSource for AtlasSource { object_store: Arc, base_config: &FileScanConfig, partition: usize, - ) -> datafusion::error::Result> { + ) -> Result> { let projected_schema = base_config.projected_schema()?; - Ok(Arc::new(AtlasOpener { object_store, + 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, - 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(), - use_pruning: self.use_pruning, - prune_cache: self.prune_cache.clone(), + scan_metrics: AtlasScanMetrics::new(&self.execution_plan_metrics, partition), + reader_pool: Arc::clone(&self.reader_pool), })) } - fn as_any(&self) -> &dyn std::any::Any { + fn as_any(&self) -> &dyn Any { self } @@ -151,30 +119,15 @@ 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()) } - /// Whether a scan may split one file across partitions. It may not. - /// - /// 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. - /// - /// 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 + /// A container is one unit. A byte range of it names nothing a reader can + /// 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 } @@ -187,20 +140,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 +147,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 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>, _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()), @@ -246,453 +187,3 @@ impl FileSource for AtlasSource { .with_updated_node(Arc::new(source))) } } - -// ─── FileOpener ──────────────────────────────────────────────────────────── - -struct AtlasOpener { - object_store: Arc, - projected_schema: SchemaRef, - batch_size: usize, - metrics: ExecutionPlanMetricsSet, - partition: usize, - read_dimensions: Option>, - predicate: Option>, - cache: Option, - use_pruning: bool, - prune_cache: PruneCache, -} - -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(); - - Ok(stream) - } -} - -/// 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, - projected_schema: SchemaRef, - read_dimensions: Option>, - predicate: Option>, - batch_size: usize, - metrics: Option, - 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() - } - } - }; - - scan_metrics.dataset_build_time.add_elapsed(build_start); - Ok(stream) -} - -/// 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(); - - Ok(stream) -} - -/// 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() - .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}")) - }) -} - -impl FileOpener for AtlasOpener { - fn open(&self, file: PartitionedFile) -> datafusion::error::Result { - let assigned_names = 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(), - 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)) - } -} - -#[cfg(test)] -mod repartition_tests { - //! Atlas divides its own work, by dataset name, so it must refuse - //! DataFusion's byte-range split. - - use std::sync::Arc; - - use datafusion::datasource::listing::PartitionedFile; - use datafusion::datasource::physical_plan::{FileScanConfigBuilder, FileSource}; - use datafusion::datasource::table_schema::TableSchema; - use datafusion::execution::object_store::ObjectStoreUrl; - - use super::AtlasSource; - - /// 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. - #[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); - 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)) - .build(); - - assert!(!source.supports_repartitioning()); - assert!( - source.repartitioned(4, 1, None, &config).unwrap().is_none(), - "an atlas store must not split by byte range" - ); - } -} - -#[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]))], - ) - .expect("source batch"); - - 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]); - } - - #[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"); - } - - #[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"); - - let out = adapt(source, target, batch).expect("Int64 -> Utf8 must be castable"); - let col = out - .column(0) - .as_any() - .downcast_ref::() - .expect("Utf8 column"); - assert_eq!( - (0..col.len()).map(|i| col.value(i)).collect::>(), - vec!["1", "2"], - ); - } -} 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/datafusion/view.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs new file mode 100644 index 00000000..21760bac --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/view.rs @@ -0,0 +1,524 @@ +//! 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}, +}; +use datafusion::physical_plan::PhysicalExpr; +use indexmap::IndexMap; +use object_store::{ObjectMeta, ObjectStore}; + +use crate::{ + compat, + datafusion::{metrics::AtlasScanMetrics, pruning::prune_datasets}, + 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 { + /// 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 + .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, + }) + } + + /// The table schema the view resolves columns for, in field order. + pub fn table_schema(&self) -> &SchemaRef { + &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>, + 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) + } + + /// 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) + .with_context(|| { + format!( + "reading array '{}' of dataset '{dataset_name}'", + field.name() + ) + })?, + ), + 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); + } + } + 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)) + } +} + +/// 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) +} + +/// 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 +/// 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" + ); + } + + /// 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] + 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/lib.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/lib.rs index 31aae712..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 @@ -1,27 +1,86 @@ -//! `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 │ temperature │ salinity │ … │ footer │ trailer +//! └── deleted.mask optional: ordinals of deleted datasets +//! ``` +//! +//! **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 +//! +//! [`store`] finds a collection's marker and opens it, through a reader cache. +//! [`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 +//! +//! 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 +//! +//! 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. +//! +//! 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. +//! +//! # 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 datafusion; -pub mod reader; -pub mod util; +pub mod store; + +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 deleted file mode 100644 index ffb5fb9a..00000000 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/reader.rs +++ /dev/null @@ -1,492 +0,0 @@ -//! High-level atlas reader that produces [`AnyDataset`] values. -//! -//! 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}`. - -use std::sync::Arc; - -use beacon_nd_array::{ - NdArrayD, - dataset::{AnyDataset, Dataset}, -}; -use indexmap::IndexMap; -use object_store::{ObjectStore, path::Path as OsPath}; - -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`]. -/// -/// 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`: -/// - `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, - 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)); - - // 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(); - - // ── 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) { - 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); - } - Err(e) => { - tracing::warn!("Skipping atlas array '{array_name}' in dataset '{dataset_name}': {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); - } - 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; - AnyDataset::try_from_dataset(dataset) - .await - .map_err(|e| anyhow::anyhow!("Failed to wrap atlas dataset as AnyDataset: {}", e)) -} - -#[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; - } - - /// 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"); - } - - // ── 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"); - } - - // Persist the metadata marker + array files to the store. - atlas.flush().await.expect("flush atlas store"); - } - - /// 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"); - } - - /// 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"); - } - - /// 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"); - } - atlas.flush().await.expect("flush atlas store"); - } -} - -#[cfg(test)] -mod tests { - use super::test_support::build_two_dataset_store; - use super::*; - use beacon_nd_array::NdArray; - use object_store::local::LocalFileSystem; - - /// 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") - } - - #[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:?}"); - } - - #[tokio::test] - async fn open_dataset_reads_array_values() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; - - let winter = open_fixture_dataset(tmp.path(), "winter").await; - let temp = winter - .get_array("temperature") - .expect("temperature array") - .as_any() - .downcast_ref::>() - .expect("downcast f32"); - assert_eq!(temp.clone_into_raw_vec().await, vec![1.0f32, 2.0, 3.0, 4.0]); - - let cycle = winter - .get_array("cycle") - .expect("cycle array") - .as_any() - .downcast_ref::>() - .expect("downcast i32"); - assert_eq!(cycle.clone_into_raw_vec().await, vec![10i32, 20, 30, 40]); - } - - #[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") - .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()]); - - let year = winter - .get_array("year") - .expect("year attribute") - .as_any() - .downcast_ref::>() - .expect("downcast i64"); - assert_eq!(year.clone_into_raw_vec().await, vec![2024i64]); - } - - #[tokio::test] - async fn open_dataset_propagates_array_fill_value() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; - - let winter = open_fixture_dataset(tmp.path(), "winter").await; - let cycle = winter - .get_array("cycle") - .expect("cycle array") - .as_any() - .downcast_ref::>() - .expect("downcast i32"); - assert_eq!(cycle.fill_value().await, Some(-1i32)); - - let temperature = winter - .get_array("temperature") - .expect("temperature array") - .as_any() - .downcast_ref::>() - .expect("downcast f32"); - assert_eq!(temperature.fill_value().await, None); - } - - #[tokio::test] - async fn open_dataset_distinguishes_between_dataset_views() { - let tmp = tempfile::tempdir().expect("temp dir"); - build_two_dataset_store(tmp.path()).await; - - let winter = open_fixture_dataset(tmp.path(), "winter").await; - let summer = open_fixture_dataset(tmp.path(), "summer").await; - - assert_eq!( - winter.dataset().get_array("temperature").unwrap().shape(), - &[4] - ); - assert_eq!( - summer.dataset().get_array("temperature").unwrap().shape(), - &[3] - ); - 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; - - let store: Arc = - Arc::new(LocalFileSystem::new_with_prefix(tmp.path()).unwrap()); - let err = open_dataset(store, OsPath::from(""), "ghost") - .await - .expect_err("should fail for unknown dataset"); - let msg = format!("{err:#}"); - assert!( - msg.contains("ghost") || msg.contains("DatasetNotFound"), - "error should mention missing dataset name: {msg}" - ); - } -} 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..f0a27a25 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/store.rs @@ -0,0 +1,406 @@ +//! 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. +#[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 + // 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)] +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..3e181a90 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/test_support.rs @@ -0,0 +1,387 @@ +//! 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"); +} + +/// 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"); +} 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/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/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-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..1f1ce303 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/default.rs @@ -0,0 +1,411 @@ +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 DefaultDataset { + pub name: String, + pub dimensions: Vec, + pub shape: Vec, + pub chunk_shape: Vec, + pub arrays: IndexMap>, +} + +impl DefaultDataset { + 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 DefaultDataset { + /// 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() + } + + /// 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 + /// 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 { + DefaultDataset::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]) -> DefaultDataset { + 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)); + DefaultDataset { + name: "gridded".to_string(), + dimensions: names(&["time", "lat"]), + shape: vec![4, 3], + chunk_shape: chunk.to_vec(), + arrays, + } + } + + /// 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(); + 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/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}"); + } +} 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..a08aedb0 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-nd-array/src/dataset/source.rs @@ -0,0 +1,81 @@ +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>; + + /// 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, + ) -> 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 + } +} 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/src/main.rs b/beacon-server/beacon-server/src/main.rs index 2b7b5efb..aceaf527 100644 --- a/beacon-server/beacon-server/src/main.rs +++ b/beacon-server/beacon-server/src/main.rs @@ -133,6 +133,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/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..20d54d6a 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,120 @@ 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: + +```text +my_collection/ +├── data.atlas one segment per variable, then a footer describing them all +└── deleted.mask optional: the datasets a delete has hidden +``` + +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: + +- **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.** 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.** 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. + +### Columns + +| Atlas | Column | +| --- | --- | +| array `temperature` | `temperature` | +| 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"`. ```sql -SELECT * FROM read_atlas(['collections/sensor/atlas.json']) +SELECT temperature, "temperature.units", ".platform" +FROM read_atlas('collections/sensor/data.atlas') +LIMIT 1 ``` -### External tables over Atlas +### Types and decoding -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: +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). -```sql -CREATE EXTERNAL TABLE sensor_atlas -STORED AS ATLAS -LOCATION 'collections/sensor/atlas.json'; +A cell nobody wrote reads as the array's fill value, and the fill reads as null. Two consequences +are worth knowing: -SELECT time, temperature -FROM sensor_atlas -WHERE time >= '2024-01-01'; +- 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, 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 +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.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 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 +163,40 @@ 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 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. | + +```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 5faf18f3..4184c387 100644 --- a/docs/docs/2.0.0-rc5/server/configuration.md +++ b/docs/docs/2.0.0-rc5/server/configuration.md @@ -331,13 +331,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` | 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 Binary Format (BBF) | Variable | Default | Description | 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 558e280d..0e775bce 100644 --- a/docs/docs/2.0.0-rc5/server/performance-tuning.md +++ b/docs/docs/2.0.0-rc5/server/performance-tuning.md @@ -227,20 +227,17 @@ 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. 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 Atlas reader again. It therefore does not parse the -`atlas.json` registry for every query. - -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. +`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 d04e3149..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 @@ -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` | [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..67dd58c7 --- /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_stale_collection_is_passed_over(con, datasets): + """Atlas before 0.16 was a directory behind an `atlas.json` registry. + + 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) + (legacy / "atlas.json").write_text("{}") + + try: + rows = con.sql("SELECT * FROM read_atlas('legacy/atlas.json')").fetchall() + except Exception: + return + 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 dc8f050a..7f8cec72 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.17.0 h5py==3.14.0 geopandas==1.1.4 rasterio==1.5.1