From 8356a58e0c14a9645c05a91e3b8ed5e509a198d1 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Mon, 17 Aug 2026 14:47:56 +0200 Subject: [PATCH] Apply the whole pushed-down projection in the ODV and BBF scans (#382) A FileSource that accepts a projection must apply all of it. OdvSource and BBFSource read only the column names out of it. An alias resolves to no file column, so the scan reads nothing. A partition value stays null. Hold a SplitProjection instead. Read the plain file columns, and let ProjectionOpener apply the rest. This is the shape that #378 gave GeoParquet. An empty projection is COUNT(*). The ODV decoder now states the row count on a batch that holds no column. Add a scan test per crate. Each test renames a column after the first one. Add beacon-core/tests/nd_pipeline_filters.rs. It shows that the nd formats stay correct without this change, because NdSourceExec and NdBroadcastExec refuse a projection. It also pins the nd filter and projection behaviour. Assisted-by: Claude:claude-opus-5 --- Cargo.lock | 2 + .../beacon-core/tests/nd_pipeline_filters.rs | 338 ++++++++++++++++++ .../beacon-arrow-bbf/Cargo.toml | 1 + .../beacon-arrow-bbf/src/datafusion/mod.rs | 148 ++++++++ .../beacon-arrow-bbf/src/datafusion/opener.rs | 21 +- .../beacon-arrow-bbf/src/datafusion/source.rs | 61 ++-- .../beacon-arrow-odv/Cargo.toml | 1 + .../beacon-arrow-odv/src/datafusion/source.rs | 62 ++-- .../beacon-arrow-odv/src/reader.rs | 42 +-- .../beacon-arrow-odv/tests/scan.rs | 159 ++++++++ 10 files changed, 753 insertions(+), 82 deletions(-) create mode 100644 beacon-db/beacon-core/tests/nd_pipeline_filters.rs create mode 100644 beacon-db/beacon-file-formats/beacon-arrow-odv/tests/scan.rs diff --git a/Cargo.lock b/Cargo.lock index a8c5e42c..a749f3c5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1402,6 +1402,7 @@ dependencies = [ "beacon-common", "beacon-datafusion-ext", "datafusion 53.1.0", + "datafusion-datasource 53.1.0", "futures", "nd-arrow-array", "object_store 0.13.2", @@ -1539,6 +1540,7 @@ dependencies = [ "bytes", "csv", "datafusion 53.1.0", + "datafusion-datasource 53.1.0", "futures", "indexmap 2.14.0", "object_store 0.13.2", diff --git a/beacon-db/beacon-core/tests/nd_pipeline_filters.rs b/beacon-db/beacon-core/tests/nd_pipeline_filters.rs new file mode 100644 index 00000000..b4eef3fd --- /dev/null +++ b/beacon-db/beacon-core/tests/nd_pipeline_filters.rs @@ -0,0 +1,338 @@ +//! Projections and filters over an nd scan, end to end. +//! +//! A NetCDF/Zarr/TIFF/HDF5 scan puts `NdSourceExec` and `NdBroadcastExec` +//! between the file scan and the rest of the plan. Two things ride on that: +//! +//! * The nodes refuse a projection, so an alias or a computed column never +//! reaches the file source — the scan below only ever gets a plain column +//! list. That is why [#382](https://github.com/maris-development/beacon/issues/382) +//! left the nd formats alone, and the first test here pins it. +//! * With the nd pipeline enabled, `NdFilterPushdown` sinks the element-wise +//! conjuncts of a `WHERE` below the broadcast into `NdFilterExec`, which +//! evaluates each conjunct on the sub-grid its inputs span and records the +//! surviving cells as a grid selection. A fully-sunk predicate leaves *no* +//! `FilterExec` in the plan, so a conjunct applied wrongly would silently +//! change the answer. +//! +//! The second group therefore runs each query twice — once with the nd pipeline +//! on, once with it off — and requires the same rows both ways. The pipeline-off +//! run is the oracle: there the predicate is a plain `FilterExec` over fully +//! broadcast columns. + +mod common; + +use arrow::record_batch::RecordBatch; +use common::{runtime_with, TestRuntime}; + +/// The WOD CTD fixture shipped with the NetCDF reader: 418 rows, every column +/// full-rank on one axis. Good for checking *what* the pipeline answers; it +/// cannot show footprint reduction, because every footprint is the whole grid. +/// [`gridded_fixture`] is the one with real coordinate axes. +fn netcdf_fixture() -> std::path::PathBuf { + std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .parent() + .unwrap() + .join("beacon-file-formats/beacon-arrow-netcdf/test_files/wod_ctd_1964.nc") +} + +/// A real gridded SST file: `lat` (1208) × `lon` (1920) × `time` (1) = 2,319,360 +/// cells, with `lat` and `lon` on their own axes. A predicate over one axis has a +/// footprint of that axis alone; one over both spans the whole plane. +fn gridded_fixture() -> std::path::PathBuf { + std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .parent() + .unwrap() + .join("beacon-file-formats/beacon-arrow-netcdf/test_files/gridded-example.nc") +} + +/// A box over the north-east Atlantic — it holds part of the fixture, not all of +/// it, so a predicate that was dropped or inverted shows up as a row count. +const BOX_WKT: &str = "POLYGON((-13 40, 32 40, 32 60, -13 60, -13 40))"; + +async fn nd_runtime(tag: &str, nd_pipeline: bool) -> TestRuntime { + let rt = runtime_with(tag, |b| if nd_pipeline { b.with_nd_pipeline() } else { b }).await; + std::fs::copy(netcdf_fixture(), rt.datasets_dir().join("nd.nc")).expect("copy fixture"); + rt.sql("CREATE EXTERNAL TABLE nd STORED AS NC LOCATION 'nd.nc'") + .await; + rt +} + +/// A runtime over [`gridded_fixture`], table `g`, nd pipeline on. +async fn gridded_runtime(tag: &str) -> TestRuntime { + let rt = runtime_with(tag, |b| b.with_nd_pipeline()).await; + std::fs::copy(gridded_fixture(), rt.datasets_dir().join("g.nc")).expect("copy fixture"); + rt.sql("CREATE EXTERNAL TABLE g STORED AS NC LOCATION 'g.nc'") + .await; + rt +} + +/// The result rendered as text, so two runtimes' answers compare directly. +fn rendered(batches: &[RecordBatch]) -> String { + arrow::util::pretty::pretty_format_batches(batches) + .expect("format") + .to_string() +} + +/// The physical plan of `sql` as one string. +async fn plan(rt: &TestRuntime, sql: &str) -> String { + rendered(&rt.sql(&format!("EXPLAIN {sql}")).await) +} + +/// The plan of `sql` annotated with the metrics the run collected. +async fn analyzed(rt: &TestRuntime, sql: &str) -> String { + rendered(&rt.sql(&format!("EXPLAIN ANALYZE {sql}")).await) +} + +/// The value of `metric` on the plan line naming `node`, as printed — so `"0"`, +/// `"2.42 K"`, `"2.32 M"`. Compared as text, because the exact humanized figure +/// is what the reader of a plan sees. +fn metric(plan: &str, node: &str, metric: &str) -> String { + let line = plan + .lines() + .find(|l| l.contains(node)) + .unwrap_or_else(|| panic!("no {node} in plan:\n{plan}")); + let rest = line + .split_once(&format!("{metric}=")) + .unwrap_or_else(|| panic!("no {metric} on the {node} line:\n{line}")) + .1; + rest.split(',') + .next() + .expect("a metric is comma-terminated") + .trim() + .to_string() +} + +// ── the nd nodes refuse a projection ──────────────────────────────────────── + +/// An alias and a computed column stay in a `ProjectionExec` above the +/// broadcast; the file source below is only ever handed a plain column list. +/// +/// This is the exemption #382 relies on. If a future change let a projection +/// through to `NetCdfSource`, the alias would reach a scan that resolves file +/// columns by name and the query would break — so assert on the plan, not just +/// the values. +#[tokio::test(flavor = "multi_thread")] +async fn an_alias_never_reaches_the_file_source() { + let rt = nd_runtime("nd-alias-plan", false).await; + + let plan = plan(&rt, "SELECT lat AS easting, lat * 2 AS doubled FROM nd").await; + assert!( + plan.contains("ProjectionExec: expr=[lat@0 as easting"), + "the projection must stay above the broadcast:\n{plan}" + ); + assert!( + plan.contains("projection=[lat], file_type=netcdf"), + "the scan must get a plain column list:\n{plan}" + ); + + // And it answers: the alias names the column, the values are `lat`'s. + let aliased = rt.sql("SELECT lat AS easting FROM nd LIMIT 3").await; + let plain = rt.sql("SELECT lat FROM nd LIMIT 3").await; + assert_eq!(aliased[0].schema().field(0).name(), "easting"); + assert_eq!(aliased[0].column(0), plain[0].column(0)); +} + +// ── a sunk filter answers what the plain filter answers ───────────────────── + +/// Every query that the nd filter pushdown rewrites must answer exactly what it +/// answers without the rewrite. +#[tokio::test(flavor = "multi_thread")] +async fn a_sunk_filter_answers_what_the_plain_filter_answers() { + let with_nd = nd_runtime("nd-filter-on", true).await; + let without = nd_runtime("nd-filter-off", false).await; + + // Each query is ordered, so the two runtimes' rows line up positionally. + let queries = [ + // A spatial predicate — a non-volatile scalar function over two columns. + format!( + "SELECT lat, lon FROM nd \ + WHERE st_within_point('{BOX_WKT}', CAST(lon AS DOUBLE), CAST(lat AS DOUBLE)) \ + ORDER BY lat, lon" + ), + // The same predicate conjoined with a plain one: two conjuncts whose + // masks have to intersect on the target grid rather than replace it. + format!( + "SELECT lat, lon, z FROM nd \ + WHERE st_within_point('{BOX_WKT}', CAST(lon AS DOUBLE), CAST(lat AS DOUBLE)) \ + AND z > 100 \ + ORDER BY lat, lon, z" + ), + // A predicate over a computed value. + "SELECT lat, z FROM nd WHERE lat * 2 > 90.0 ORDER BY lat, z".to_string(), + // A disjunction inside one conjunct — it is not split, so the nd filter + // has to union the branches rather than intersect them. + "SELECT lat, lon FROM nd WHERE lat > 60 OR lon < -10 ORDER BY lat, lon".to_string(), + // A projection with an alias on top of a sunk filter. + "SELECT lat AS easting, z AS depth FROM nd WHERE lat > 0 AND z < 50 ORDER BY 1, 2" + .to_string(), + // A predicate combining two columns arithmetically. + "SELECT lat, z FROM nd WHERE lat + z > 200 ORDER BY lat, z".to_string(), + ]; + + // The fixture's full row count, so each predicate can be shown to select a + // strict subset — a filter that matched everything, or nothing, would agree + // across the two runtimes while proving nothing. + const TOTAL_ROWS: usize = 418; + + for sql in &queries { + let with = with_nd.sql(sql).await; + let plain = without.sql(sql).await; + assert_eq!( + rendered(&with), + rendered(&plain), + "the nd pipeline changed the answer to:\n{sql}" + ); + + let rows: usize = with.iter().map(|b| b.num_rows()).sum(); + assert!( + rows > 0 && rows < TOTAL_ROWS, + "this query should select a strict subset, got {rows} of {TOTAL_ROWS}:\n{sql}" + ); + } +} + +/// The comparison above is only worth something if the rewrite actually fires. +#[tokio::test(flavor = "multi_thread")] +async fn the_spatial_predicate_is_sunk_below_the_broadcast() { + let rt = nd_runtime("nd-filter-fires", true).await; + + let sql = format!( + "SELECT lat, lon FROM nd \ + WHERE st_within_point('{BOX_WKT}', CAST(lon AS DOUBLE), CAST(lat AS DOUBLE))" + ); + let plan = plan(&rt, &sql).await; + assert!( + plan.contains("NdFilterExec: predicate=[st_within_point("), + "the spatial predicate must sink below the broadcast:\n{plan}" + ); + assert!( + !plan.contains("FilterExec: st_within_point"), + "a fully-sunk predicate leaves no residual filter:\n{plan}" + ); +} + +/// A volatile function is not element-wise under broadcast — evaluating it on a +/// sub-grid would repeat one draw across a whole slice — so it must stay in a +/// `FilterExec` above the broadcast. +#[tokio::test(flavor = "multi_thread")] +async fn a_volatile_predicate_stays_above_the_broadcast() { + let rt = nd_runtime("nd-filter-volatile", true).await; + + let plan = plan(&rt, "SELECT lat FROM nd WHERE random() < 0.5 AND lat > 0").await; + assert!( + plan.contains("NdFilterExec: predicate=[lat@0 > 0]"), + "the deterministic conjunct still sinks:\n{plan}" + ); + assert!( + plan.contains("FilterExec: random()"), + "the volatile conjunct must stay above the broadcast:\n{plan}" + ); +} + +// ── what the footprint actually buys ──────────────────────────────────────── + +/// A predicate over one coordinate axis is evaluated over *that axis*, not the +/// grid it selects from. +/// +/// This is the whole point of `NdFilterExec`: `lat > 40` is 2.42 K comparisons +/// on the lat axis, and the resulting mask is lifted to the 2.32 M-cell grid. +#[tokio::test(flavor = "multi_thread")] +async fn a_single_axis_predicate_is_evaluated_over_that_axis() { + let rt = gridded_runtime("nd-footprint-axis").await; + + let plan = analyzed(&rt, "SELECT lat, analysed_sst FROM g WHERE lat > 40").await; + assert_eq!(metric(&plan, "NdFilterExec", "input_rows"), "2.32 M"); + assert_eq!( + metric(&plan, "NdFilterExec", "elements_evaluated"), + "2.42 K", + "the predicate runs on the lat axis, not on the grid:\n{plan}" + ); + assert_eq!(metric(&plan, "NdFilterExec", "elements_saved"), "2.32 M"); +} + +/// A predicate over *two* axes has both in its footprint, so it is evaluated +/// once per cell of their plane — there is no reduction to a smaller axis. +/// +/// This is the shape every spatial predicate takes: `lon` is on the lon axis and +/// `lat` on the lat axis, so `ST_Within(ST_Point(lon, lat), …)` spans lat × lon, +/// which on this single-time-step file *is* the whole grid. The saving on a file +/// with further axes is that the mask is computed once for the plane and reused +/// across them — not that the function runs fewer times per cell of the plane. +#[tokio::test(flavor = "multi_thread")] +async fn a_lon_lat_predicate_spans_the_whole_plane() { + let rt = gridded_runtime("nd-footprint-plane").await; + + let plan = analyzed( + &rt, + "SELECT lat, lon FROM g \ + WHERE ST_Within(ST_Point(CAST(lon AS DOUBLE), CAST(lat AS DOUBLE)), \ + ST_GeomFromText('POLYGON((0 0, 10 0, 10 10, 0 10, 0 0))'))", + ) + .await; + + assert!( + plan.contains("NdFilterExec: predicate=[st_within("), + "the spatial predicate is evaluated before the broadcast:\n{plan}" + ); + assert_eq!( + metric(&plan, "NdFilterExec", "elements_evaluated"), + "2.32 M", + "lat ∪ lon is the whole plane here:\n{plan}" + ); + assert_eq!(metric(&plan, "NdFilterExec", "elements_saved"), "0"); +} + +/// The geometry itself is built before the broadcast too: `ST_Point` becomes an +/// `NdProjectionExec` output column on the un-broadcast nd columns. +#[tokio::test(flavor = "multi_thread")] +async fn st_point_is_constructed_below_the_broadcast() { + let rt = gridded_runtime("nd-footprint-point").await; + + let plan = plan( + &rt, + "SELECT ST_Point(CAST(lon AS DOUBLE), CAST(lat AS DOUBLE)) AS geom FROM g WHERE lat > 40", + ) + .await; + + assert!( + plan.contains("NdProjectionExec: exprs=[geom]"), + "the point is constructed below the broadcast:\n{plan}" + ); + assert!( + !plan.contains("ProjectionExec: expr=[st_point("), + "…and not left above it:\n{plan}" + ); +} + +/// Whether a predicate sinks depends on the *select list*, not the predicate. +/// +/// Narrowing the output to fewer columns than the predicate reads makes +/// DataFusion fold the narrowing into the filter (`FilterExec: …, +/// projection=[…]`), and `NdFilterPushdown` declines any filter carrying one. So +/// the most natural spatial query — project a couple of columns, filter on +/// lon/lat — is exactly the one that still filters *after* the broadcast. +#[tokio::test(flavor = "multi_thread")] +async fn a_narrowing_select_list_keeps_the_filter_above_the_broadcast() { + let rt = gridded_runtime("nd-footprint-narrow").await; + + const WITHIN: &str = "ST_Within(ST_Point(CAST(lon AS DOUBLE), CAST(lat AS DOUBLE)), \ + ST_GeomFromText('POLYGON((0 0, 10 0, 10 10, 0 10, 0 0))'))"; + + // Both predicate columns projected: the filter sinks. + let both = plan(&rt, &format!("SELECT lat, lon FROM g WHERE {WITHIN}")).await; + assert!( + both.contains("NdFilterExec: predicate=[st_within("), + "{both}" + ); + + // One of them projected: the filter carries the narrowing and stays put. + let one = plan(&rt, &format!("SELECT lat FROM g WHERE {WITHIN}")).await; + assert!( + one.contains("FilterExec: st_within(") && one.contains("projection=[lat@0]"), + "the narrowing is folded into the filter:\n{one}" + ); + assert!( + !one.contains("NdFilterExec"), + "so nothing sinks below the broadcast:\n{one}" + ); +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-bbf/Cargo.toml b/beacon-db/beacon-file-formats/beacon-arrow-bbf/Cargo.toml index 9877069d..4566cb21 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-bbf/Cargo.toml +++ b/beacon-db/beacon-file-formats/beacon-arrow-bbf/Cargo.toml @@ -7,6 +7,7 @@ rust-version.workspace = true [dependencies] arrow = { workspace = true } datafusion = { workspace = true } +datafusion-datasource = { workspace = true } object_store = { workspace = true } async-trait = { workspace = true } futures = { workspace = true } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/mod.rs index 763586ed..b3d490c9 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/mod.rs @@ -488,3 +488,151 @@ mod tests { ); } } + +#[cfg(test)] +mod scan_tests { + //! End-to-end checks on the columns the BBF scan returns. + //! + //! `BBFSource` accepts a pushed-down projection, so it has to apply the + //! whole of it. The projections here rename a column after the first one — + //! the shape [#382](https://github.com/maris-development/beacon/issues/382) + //! reported: the scan looked for a file column under the alias, found none, + //! and read nothing. + + use std::sync::Arc; + + use arrow::array::{Array, AsArray}; + use arrow::record_batch::RecordBatch; + use datafusion::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, + }; + use datafusion::prelude::{SessionConfig, SessionContext}; + + use super::BBFFormat; + + /// A session over the two-entry fixture. The temp dir is returned so it + /// outlives the queries. + async fn table() -> (SessionContext, tempfile::TempDir) { + let dir = tempfile::tempdir().expect("tempdir"); + super::test_util::write_bbf_fixture(dir.path(), "scan.bbf").await; + + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + let options = + ListingOptions::new(Arc::new(BBFFormat::default())).with_file_extension(".bbf"); + let url = ListingTableUrl::parse(dir.path().to_str().expect("utf-8 path")).expect("url"); + let schema = options + .infer_schema(&ctx.state(), &url) + .await + .expect("schema"); + let config = ListingTableConfig::new(url) + .with_listing_options(options) + .with_schema(schema); + ctx.register_table("t", Arc::new(ListingTable::try_new(config).expect("table"))) + .expect("register"); + (ctx, dir) + } + + async fn query(ctx: &SessionContext, sql: &str) -> Vec { + ctx.sql(sql) + .await + .unwrap_or_else(|e| panic!("planning {sql}: {e}")) + .collect() + .await + .unwrap_or_else(|e| panic!("running {sql}: {e}")) + } + + fn one_batch(batches: Vec) -> RecordBatch { + let schema = batches.first().expect("at least one batch").schema(); + arrow::compute::concat_batches(&schema, &batches).expect("concat") + } + + /// Every value of `ints`, in the fixture's own order. + fn ints(batch: &RecordBatch, column: usize) -> Vec { + batch + .column(column) + .as_primitive::() + .values() + .to_vec() + } + + /// An aliased projection is pushed into the scan whole. The scan has to + /// rename the column it read, not look for a file column under the alias. + /// + /// `ints` sits after the entry-key column, so a scan that kept the file's + /// own column order would also answer with the wrong values here. + #[tokio::test] + async fn applies_an_aliased_projection() { + let (ctx, _dir) = table().await; + let batch = one_batch(query(&ctx, "SELECT ints AS measurement FROM t").await); + + assert_eq!(batch.schema().field(0).name(), "measurement"); + assert_eq!(batch.column(0).null_count(), 0); + assert_eq!( + ints(&batch, 0), + vec![1, 2, 3, 10, 20], + "3 rows from entry_a, then 2 from entry_b" + ); + } + + /// A projection that reorders columns and renames one of them returns both, + /// in the order asked for and with their own values. + #[tokio::test] + async fn applies_a_reordered_projection() { + let (ctx, _dir) = table().await; + let batch = one_batch( + query( + &ctx, + "SELECT ints AS measurement, __entry_key AS entry FROM t", + ) + .await, + ); + + assert_eq!( + batch + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect::>(), + vec!["measurement", "entry"] + ); + assert_eq!(ints(&batch, 0), vec![1, 2, 3, 10, 20]); + + let entries = batch.column(1).as_string::(); + assert_eq!( + (0..5).map(|i| entries.value(i)).collect::>(), + vec!["entry_a", "entry_a", "entry_a", "entry_b", "entry_b"] + ); + } + + /// A computed column is a projection too, and it is pushed down whole. + #[tokio::test] + async fn applies_a_computed_projection() { + let (ctx, _dir) = table().await; + let batch = one_batch(query(&ctx, "SELECT ints + 1 AS bumped FROM t").await); + + assert_eq!(batch.schema().field(0).name(), "bumped"); + assert_eq!( + batch + .column(0) + .as_primitive::() + .values() + .to_vec(), + vec![2, 3, 4, 11, 21] + ); + } + + /// An alias must change nothing but the column's name. A scan that resolved + /// file columns by the *output* name would silently return a different set + /// of rows here, which is what #382 reported. + #[tokio::test] + async fn an_alias_changes_only_the_name() { + let (ctx, _dir) = table().await; + let plain = one_batch(query(&ctx, "SELECT names FROM t").await); + let aliased = one_batch(query(&ctx, "SELECT names AS label FROM t").await); + + assert_eq!(aliased.schema().field(0).name(), "label"); + assert_eq!(aliased.num_rows(), plain.num_rows()); + assert_eq!(aliased.column(0), plain.column(0)); + } +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/opener.rs b/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/opener.rs index 4be758b0..12ef6404 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/opener.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/opener.rs @@ -31,7 +31,10 @@ use parking_lot::Mutex; use crate::datafusion::{metrics::BBFGlobalMetrics, stream_share::StreamShare}; pub struct BBFOpener { - projected_schema: SchemaRef, + /// The plain file columns to read, in file order. Every mapped batch is + /// produced in this schema; anything else the query asks for is applied by + /// the `ProjectionOpener` wrapped around this one. + read_schema: SchemaRef, pruning_predicate: Option, object_store: Arc, table_schema: Arc, @@ -50,7 +53,7 @@ impl FileOpener for BBFOpener { fn open(&self, file: PartitionedFile) -> datafusion::error::Result { let async_reader = ArrowBBFObjectReader::new(file.object_meta.location.clone(), self.object_store.clone()); - let projected_schema = self.projected_schema.clone(); + let read_schema = self.read_schema.clone(); let pruning_predicate = self.pruning_predicate.clone(); let table_schema = self.table_schema.clone(); let file_tracer = self.file_tracer.clone(); @@ -64,7 +67,7 @@ impl FileOpener for BBFOpener { .clone() }; let metrics = self.metrics.clone(); - let fut_projected_schema = projected_schema.clone(); + let fut_read_schema = read_schema.clone(); let split_streams_slice = self.split_streams_slice; let split_batch_size = self.split_batch_size; @@ -86,12 +89,12 @@ impl FileOpener for BBFOpener { .fields() .iter() .enumerate() - .filter(|(_, f)| fut_projected_schema.index_of(f.name()).is_ok()) + .filter(|(_, f)| fut_read_schema.index_of(f.name()).is_ok()) .map(|(i, _)| i) .collect(); let source_schema: SchemaRef = Arc::new(file_schema.project(&projection)?); let schema_mapper = Arc::new( - BatchAdapterFactory::new(fut_projected_schema.clone()) + BatchAdapterFactory::new(fut_read_schema.clone()) .make_adapter(&source_schema)?, ); let mut selection: Option = None; @@ -149,8 +152,8 @@ impl FileOpener for BBFOpener { RecordBatch::new_empty(file_schema.clone()) }); let batch_schema = arrow_batch.schema(); - // Map the batch schema to the table schema. - let schema_mapper = BatchAdapterFactory::new(projected_schema.clone()) + // Map the batch schema to the read schema. + let schema_mapper = BatchAdapterFactory::new(read_schema.clone()) .make_adapter(&batch_schema)?; let mapped_batch = schema_mapper .adapt_batch(&arrow_batch) @@ -342,7 +345,7 @@ mod split_tests { impl BBFOpener { #[allow(clippy::too_many_arguments)] pub fn new( - projected_schema: SchemaRef, + read_schema: SchemaRef, pruning_predicate: Option, object_store: Arc, table_schema: Arc, @@ -353,7 +356,7 @@ impl BBFOpener { split_batch_size: usize, ) -> Self { Self { - projected_schema, + read_schema, object_store, pruning_predicate, table_schema, diff --git a/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/source.rs index 68f01022..d2b4156c 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/source.rs @@ -17,6 +17,7 @@ use datafusion::{ metrics::ExecutionPlanMetricsSet, }, }; +use datafusion_datasource::projection::{ProjectionOpener, SplitProjection}; use object_store::ObjectStore; use parking_lot::Mutex; @@ -42,8 +43,13 @@ pub struct BBFSource { stream_partition_shares: Arc>>>, /// Global Metrics global_metrics: BBFGlobalMetrics, - /// Projection pushed down by the scan, applied on top of the table schema. - projection: Option, + /// The projection the scan pushed down, split into the file columns the + /// reader selects and a remainder applied on top of them. + /// + /// A `FileSource` that accepts a projection must apply it in full, so this + /// source only reads plain columns and leaves everything else — aliases, + /// computed expressions, partition columns — to [`ProjectionOpener`]. + projection: SplitProjection, } impl BBFSource { @@ -52,6 +58,7 @@ impl BBFSource { let global_metrics = BBFGlobalMetrics::new(base_metrics.clone()); Self { schema_adapter_factory: None, + projection: SplitProjection::unprojected(&table_schema), table_schema, execution_plan_metrics: base_metrics, batch_size: 32 * 1024, @@ -60,7 +67,6 @@ impl BBFSource { file_tracer: Arc::new(Mutex::new(Arc::new(Mutex::new(vec![])))), stream_partition_shares: Arc::new(Mutex::new(HashMap::new())), global_metrics, - projection: None, } } @@ -75,7 +81,10 @@ impl BBFSource { /// preserve a pushed-down projection when the format rebuilds the source /// in `create_physical_plan`. pub fn with_projection(mut self, projection: Option) -> Self { - self.projection = projection; + self.projection = match projection { + Some(projection) => SplitProjection::new(self.table_schema.file_schema(), &projection), + None => SplitProjection::unprojected(&self.table_schema), + }; self } @@ -90,27 +99,31 @@ impl FileSource for BBFSource { fn create_file_opener( &self, object_store: Arc, - base_config: &FileScanConfig, + _base_config: &FileScanConfig, _partition: usize, ) -> datafusion::error::Result> { - let table_schema = self.table_schema.file_schema().clone(); - let projected_schema = base_config.projected_schema()?; + let file_schema = self.table_schema.file_schema(); + // The columns the reader selects, in file order. `ProjectionOpener` + // derives its input schema the same way, so the two always agree. + let read_schema = Arc::new(file_schema.project(&self.projection.file_indices)?); let pruning_predicate = self .predicate .clone() - .map(|p| PruningPredicate::try_new(p, table_schema.clone())) + .map(|p| PruningPredicate::try_new(p, file_schema.clone())) .transpose()?; - Ok(Arc::new(BBFOpener::new( - projected_schema, + let opener = Arc::new(BBFOpener::new( + read_schema, pruning_predicate, object_store, - table_schema, + file_schema.clone(), self.file_tracer.lock().clone(), self.stream_partition_shares.clone(), self.global_metrics.clone(), self.split_streams_slice, self.batch_size, - ))) + )) as Arc; + + ProjectionOpener::try_new(self.projection.clone(), opener, file_schema) } /// Any @@ -139,19 +152,16 @@ impl FileSource for BBFSource { } fn projection(&self) -> Option<&ProjectionExprs> { - self.projection.as_ref() + Some(&self.projection.source) } fn try_pushdown_projection( &self, projection: &ProjectionExprs, ) -> datafusion::error::Result>> { - let merged = match &self.projection { - Some(existing) => existing.try_merge(projection)?, - None => projection.clone(), - }; + let merged = self.projection.source.try_merge(projection)?; let source = BBFSource { - projection: Some(merged), + projection: SplitProjection::new(self.table_schema.file_schema(), &merged), ..self.clone() }; Ok(Some(Arc::new(source))) @@ -223,13 +233,20 @@ mod tests { .expect("should still be a BBFSource") } - /// A fresh source must not prune, project or slice anything; those only appear - /// once the optimizer pushes them down. + /// A fresh source must not prune or slice anything, and it must report every + /// column: a source that accepts a projection always states one, and before + /// the optimizer pushes anything down that projection is the whole table. #[test] - fn new_source_starts_without_predicate_or_projection() { + fn new_source_starts_without_predicate_and_projects_every_column() { let source = source(); assert!(source.predicate.is_none()); - assert!(source.projection().is_none()); + assert_eq!( + source + .projection() + .expect("a source that accepts a projection always states one") + .column_indices(), + vec![0, 1, 2] + ); assert!(!source.split_streams_slice); assert_eq!(source.file_type(), "bbf"); } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-odv/Cargo.toml b/beacon-db/beacon-file-formats/beacon-arrow-odv/Cargo.toml index a83e36ff..1258c95e 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-odv/Cargo.toml +++ b/beacon-db/beacon-file-formats/beacon-arrow-odv/Cargo.toml @@ -11,6 +11,7 @@ indexmap = { workspace = true } arrow = { workspace=true} arrow-csv = {workspace = true} datafusion = { workspace = true } +datafusion-datasource = { workspace = true } tracing = { workspace = true } beacon-common = { path = "../../beacon-common" } beacon-datafusion-ext = { path = "../../beacon-datafusion-ext" } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-odv/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-odv/src/datafusion/source.rs index fb2436cc..925b166f 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-odv/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-odv/src/datafusion/source.rs @@ -17,6 +17,7 @@ use datafusion::{ physical_expr_adapter::BatchAdapterFactory, physical_plan::metrics::ExecutionPlanMetricsSet, }; +use datafusion_datasource::projection::{ProjectionOpener, SplitProjection}; use futures::{StreamExt, TryFutureExt, TryStreamExt}; use object_store::{ObjectStore, ObjectStoreExt}; @@ -37,8 +38,13 @@ pub struct OdvSource { table_schema: TableSchema, /// Execution plan metrics. execution_plan_metrics: ExecutionPlanMetricsSet, - /// Projection pushed down by the scan, applied on top of the table schema. - projection: Option, + /// The projection the scan pushed down, split into the file columns the + /// decoder reads and a remainder applied on top of them. + /// + /// A `FileSource` that accepts a projection must apply it in full, so this + /// source only reads plain columns and leaves everything else — aliases, + /// computed expressions, partition columns — to [`ProjectionOpener`]. + projection: SplitProjection, } impl OdvSource { @@ -46,9 +52,9 @@ impl OdvSource { pub fn new(table_schema: TableSchema) -> Self { Self { schema_adapter_factory: None, + projection: SplitProjection::unprojected(&table_schema), table_schema, execution_plan_metrics: ExecutionPlanMetricsSet::new(), - projection: None, } } @@ -56,7 +62,10 @@ impl OdvSource { /// preserve a pushed-down projection when the format rebuilds the source /// in `create_physical_plan`. pub fn with_projection(mut self, projection: Option) -> Self { - self.projection = projection; + self.projection = match projection { + Some(projection) => SplitProjection::new(self.table_schema.file_schema(), &projection), + None => SplitProjection::unprojected(&self.table_schema), + }; self } } @@ -67,15 +76,20 @@ impl FileSource for OdvSource { fn create_file_opener( &self, object_store: Arc, - base_config: &FileScanConfig, + _base_config: &FileScanConfig, _partition: usize, ) -> datafusion::error::Result> { - let projected_schema = base_config.projected_schema()?; + let file_schema = self.table_schema.file_schema(); + // The columns the decoder reads, in file order. `ProjectionOpener` + // derives its input schema the same way, so the two always agree. + let read_schema = Arc::new(file_schema.project(&self.projection.file_indices)?); - Ok(Arc::new(OdvOpener { - projected_schema, + let opener = Arc::new(OdvOpener { + read_schema, object_store, - })) + }) as Arc; + + ProjectionOpener::try_new(self.projection.clone(), opener, file_schema) } fn table_schema(&self) -> &TableSchema { @@ -118,10 +132,8 @@ impl FileSource for OdvSource { factory: Arc, ) -> datafusion::error::Result> { Ok(Arc::new(Self { - table_schema: self.table_schema.clone(), - execution_plan_metrics: self.execution_plan_metrics.clone(), schema_adapter_factory: Some(factory), - projection: self.projection.clone(), + ..self.clone() })) } @@ -131,19 +143,16 @@ impl FileSource for OdvSource { } fn projection(&self) -> Option<&ProjectionExprs> { - self.projection.as_ref() + Some(&self.projection.source) } fn try_pushdown_projection( &self, projection: &ProjectionExprs, ) -> datafusion::error::Result>> { - let merged = match &self.projection { - Some(existing) => existing.try_merge(projection)?, - None => projection.clone(), - }; + let merged = self.projection.source.try_merge(projection)?; let source = Self { - projection: Some(merged), + projection: SplitProjection::new(self.table_schema.file_schema(), &merged), ..self.clone() }; Ok(Some(Arc::new(source))) @@ -154,8 +163,10 @@ impl FileSource for OdvSource { /// /// It uses a schema adapter and handles file compression. struct OdvOpener { - /// The projected output schema each mapped batch is produced in. - projected_schema: SchemaRef, + /// The plain file columns to read, in file order. Every mapped batch is + /// produced in this schema; anything else the query asks for is applied by + /// the [`ProjectionOpener`] wrapped around this one. + read_schema: SchemaRef, /// Object store for file access. object_store: Arc, } @@ -163,7 +174,7 @@ struct OdvOpener { impl FileOpener for OdvOpener { /// Opens an ODV file and returns a stream of record batches. fn open(&self, file: PartitionedFile) -> datafusion::error::Result { - let projected_schema = self.projected_schema.clone(); + let read_schema = self.read_schema.clone(); let object_store = self.object_store.clone(); let compression = OdvFormat::infer_compression(&file.object_meta); @@ -188,15 +199,14 @@ impl FileOpener for OdvOpener { .fields() .iter() .enumerate() - .filter(|(_, f)| projected_schema.index_of(f.name()).is_ok()) + .filter(|(_, f)| read_schema.index_of(f.name()).is_ok()) .map(|(i, _)| i) .collect(); - // Adapt decoded batches onto the projected output schema: reorder, - // cast, and null-fill columns the file lacks. + // Adapt decoded batches onto the read schema: reorder, cast, and + // null-fill columns this file lacks. let source_schema: SchemaRef = Arc::new(file_schema.project(&projection)?); - let adapter = - BatchAdapterFactory::new(projected_schema).make_adapter(&source_schema)?; + let adapter = BatchAdapterFactory::new(read_schema).make_adapter(&source_schema)?; // Open and decode the file body let body_stream = object_store diff --git a/beacon-db/beacon-file-formats/beacon-arrow-odv/src/reader.rs b/beacon-db/beacon-file-formats/beacon-arrow-odv/src/reader.rs index 5e261f83..bd0c5425 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-odv/src/reader.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-odv/src/reader.rs @@ -1,7 +1,7 @@ use std::{collections::HashMap, io::Read, sync::Arc, task::Poll}; use arrow::{ - array::{RecordBatch, StringArray}, + array::{RecordBatch, RecordBatchOptions, StringArray}, datatypes::{DataType, Field, SchemaRef}, error::ArrowError, }; @@ -104,28 +104,12 @@ impl OdvSchemaMapper { batch: RecordBatch, projection: Option>, ) -> Result { - let mut schema = self.output_schema.clone(); - let mut arrays = batch.columns().to_vec(); - for (_, value) in self.metadata_fields.iter() { - let array = Arc::new(StringArray::from_iter_values(std::iter::repeat_n( - value.clone(), - batch.num_rows(), - ))); - - arrays.push(array); - } - - //Apply the projection - if let Some(projection) = projection { - let projection = projection.as_ref(); - arrays = projection - .iter() - .map(|&idx| arrays[idx].clone()) - .collect::>(); - schema = Arc::new(schema.project(projection)?); - } - - RecordBatch::try_new(schema, arrays) + AsyncOdvDecoder::decode_batch( + self.output_schema.clone(), + &self.metadata_fields, + projection, + batch, + ) } } @@ -360,11 +344,12 @@ impl AsyncOdvDecoder { batch: RecordBatch, ) -> Result { let mut schema = output_schema; + let rows = batch.num_rows(); let mut arrays = batch.columns().to_vec(); for (_, value) in metadata_fields.iter() { let array = Arc::new(StringArray::from_iter_values(std::iter::repeat_n( value.clone(), - batch.num_rows(), + rows, ))); arrays.push(array); @@ -380,7 +365,14 @@ impl AsyncOdvDecoder { schema = Arc::new(schema.project(projection)?); } - RecordBatch::try_new(schema, arrays) + // A projection of no columns at all is `COUNT(*)`: the batch then holds + // only its row count, which a column-less `RecordBatch` cannot state on + // its own. + RecordBatch::try_new_with_options( + schema, + arrays, + &RecordBatchOptions::new().with_row_count(Some(rows)), + ) } fn decode_byte_stream> + Unpin>( diff --git a/beacon-db/beacon-file-formats/beacon-arrow-odv/tests/scan.rs b/beacon-db/beacon-file-formats/beacon-arrow-odv/tests/scan.rs new file mode 100644 index 00000000..1d135d17 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-arrow-odv/tests/scan.rs @@ -0,0 +1,159 @@ +//! End-to-end checks on the ODV scan: the columns it returns. +//! +//! `OdvSource` accepts a pushed-down projection, so it has to apply the whole +//! of it. Every projection here renames a column after the first one — the +//! shape [#382](https://github.com/maris-development/beacon/issues/382) +//! reported: the scan looked for a file column under the alias, found none, and +//! decoded nothing. + +use std::sync::Arc; + +use arrow::array::{Array, AsArray}; +use arrow::record_batch::RecordBatch; +use beacon_arrow_odv::datafusion::OdvFormat; +use datafusion::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, +}; +use datafusion::prelude::{SessionConfig, SessionContext}; + +/// A session over the crate's ODV fixture. +async fn table() -> SessionContext { + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + + let options = + ListingOptions::new(Arc::new(OdvFormat::new())).with_file_extension("test_file.txt"); + let url = ListingTableUrl::parse( + std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test-data") + .to_str() + .expect("utf-8 path"), + ) + .expect("listing url"); + let schema = options + .infer_schema(&ctx.state(), &url) + .await + .expect("schema"); + let config = ListingTableConfig::new(url) + .with_listing_options(options) + .with_schema(schema); + ctx.register_table("t", Arc::new(ListingTable::try_new(config).expect("table"))) + .expect("register"); + ctx +} + +async fn query(ctx: &SessionContext, sql: &str) -> Vec { + ctx.sql(sql) + .await + .unwrap_or_else(|e| panic!("planning {sql}: {e}")) + .collect() + .await + .unwrap_or_else(|e| panic!("running {sql}: {e}")) +} + +fn one_batch(batches: Vec) -> RecordBatch { + let schema = batches.first().expect("at least one batch").schema(); + arrow::compute::concat_batches(&schema, &batches).expect("concat") +} + +/// An aliased projection is pushed into the scan whole. The scan has to rename +/// the column it read, not look for a file column under the alias. +/// +/// `Station` is the second column of the file, so a scan that kept the file's +/// own column order would also answer with the wrong values here. +#[tokio::test] +async fn applies_an_aliased_projection() { + let ctx = table().await; + let batch = one_batch(query(&ctx, r#"SELECT "Station" AS s FROM t LIMIT 3"#).await); + + assert_eq!(batch.schema().field(0).name(), "s"); + assert_eq!(batch.num_rows(), 3); + assert_eq!(batch.column(0).null_count(), 0); + let stations = batch.column(0).as_string::(); + assert!( + (0..stations.len()).all(|i| !stations.value(i).is_empty()), + "every station name is a value from the file" + ); +} + +/// A projection that reorders columns returns them in the order asked for. +#[tokio::test] +async fn applies_a_reordered_projection() { + let ctx = table().await; + let batch = one_batch(query(&ctx, r#"SELECT "Station", "Cruise" AS c FROM t LIMIT 3"#).await); + + assert_eq!( + batch + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect::>(), + vec!["Station", "c"] + ); + assert_eq!(batch.column(0).null_count(), 0); + assert_eq!(batch.column(1).null_count(), 0); +} + +/// A computed column is a projection too, and it is pushed down whole. +#[tokio::test] +async fn applies_a_computed_projection() { + let ctx = table().await; + let batch = one_batch(query(&ctx, r#"SELECT "Depth" + 1.0 AS deeper FROM t LIMIT 3"#).await); + + assert_eq!(batch.schema().field(0).name(), "deeper"); + assert_eq!(batch.num_rows(), 3); +} + +/// A partition column is part of the projection the scan accepts, so the scan +/// has to fill it from the file's path. Dropping it leaves the column null. +#[tokio::test] +async fn fills_a_partition_column() { + let dir = tempfile::tempdir().expect("tempdir"); + let partition = dir.path().join("basin=atlantic"); + std::fs::create_dir(&partition).expect("partition dir"); + std::fs::copy( + std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test-data") + .join("test_file.txt"), + partition.join("test_file.txt"), + ) + .expect("copy fixture"); + + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + let options = ListingOptions::new(Arc::new(OdvFormat::new())) + .with_file_extension("test_file.txt") + .with_table_partition_cols(vec![( + "basin".to_string(), + arrow::datatypes::DataType::Utf8, + )]); + let url = ListingTableUrl::parse(dir.path().to_str().expect("utf-8 path")).expect("url"); + let schema = options + .infer_schema(&ctx.state(), &url) + .await + .expect("schema"); + let config = ListingTableConfig::new(url) + .with_listing_options(options) + .with_schema(schema); + ctx.register_table("t", Arc::new(ListingTable::try_new(config).expect("table"))) + .expect("register"); + + let batch = one_batch(query(&ctx, r#"SELECT basin, "Station" FROM t LIMIT 3"#).await); + assert_eq!(batch.num_rows(), 3); + let basin = batch.column(0).as_string::(); + assert_eq!( + (0..3).map(|i| basin.value(i)).collect::>(), + vec!["atlantic"; 3] + ); +} + +/// `count(*)` selects no column at all. +#[tokio::test] +async fn counts_rows_without_reading_a_column() { + let ctx = table().await; + let batch = one_batch(query(&ctx, "SELECT count(*) FROM t").await); + let count = batch + .column(0) + .as_primitive::() + .value(0); + assert!(count > 0, "the fixture holds rows"); +}