From f10af1b0376a4438285a3897bdb5b6312eb1bffb Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Sun, 9 Aug 2026 14:23:50 +0200 Subject: [PATCH 1/2] Read GeoTIFF and TIFF through the nd pipeline beacon-arrow-tiff was the last nd format still on the v1 broadcast, which built the pixel cross-product per column inside the opener. It now runs the same spine as netCDF, HDF5 and Zarr. - The opener emits `beacon.nd`-encoded batches through `any_dataset_as_encoded_stream` and adapts them in the encoded struct domain. - `create_physical_plan` sets the file source schema to `encoded_schema(file_schema)` and returns `NdBroadcastExec(NdSourceExec(DataSourceExec))`. - `COUNT(*)` keeps the v1 zero-column driver path, and now also drives with the predicate columns, which the old `any_dataset_as_row_size` path ignored. A raster is a grid over `y` (image rows) and `x` (image columns): bands on (y, x), `geo.lat` on y, `geo.lon` on x, TIFF tags as rank-0 scalars. The two nd optimizer rules therefore now reach a raster. `read_tiff` also gains the optional `dimensions` argument the other nd readers take, with `OPTIONS (read_dimensions '...')` as the external-table equivalent, and `resolve_read_dimensions` in both schema inference and the opener. --- CHANGELOG.md | 13 + beacon-db/beacon-core/tests/read_functions.rs | 77 +++ .../beacon-arrow-tiff/src/datafusion/mod.rs | 572 ++++++++++++------ .../src/datafusion/reader.rs | 23 + .../src/datafusion/source.rs | 175 ++++-- .../src/datafusion/table_function.rs | 43 +- docs/docs/2.0.0-rc2/formats/geotiff.md | 10 + docs/docs/2.0.0-rc2/sql/table-functions.md | 10 +- 8 files changed, 682 insertions(+), 241 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1681e0e9..74f92d4b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,6 +37,19 @@ tag. Releases before 2.0.0 are recorded in the ### Changed +- **GeoTIFF and TIFF read through the nd pipeline**, as netCDF, HDF5 and Zarr already did. A raster + is a grid over the axes `y` (image rows) and `x` (image columns), and the reader now hands the + plan its columns un-broadcast — the bands on `y × x`, `geo.lat` on `y`, `geo.lon` on `x`, and the + TIFF tags as scalars — for a decode and a broadcast above the scan. The pixel cross-product is + therefore built once, at the top of the plan, instead of per column in the reader, and the two nd + optimizer rules now reach a raster: a `WHERE` on a coordinate selects the grid before it is + materialized, and an element-wise `SELECT` expression on a coordinate is evaluated on that axis + alone rather than on every pixel. `read_tiff` also takes the optional `dimensions` argument the + other nd readers take: `read_tiff('raster.tif', ['y'])` returns one row per image row. The same + list is accepted as `OPTIONS (read_dimensions 'y')` on `CREATE EXTERNAL TABLE … STORED AS TIFF`. + As for every nd format, the projected columns define the grid, so selecting only 1-d coordinate + columns returns their own axes rather than the full pixel grid. See + [GeoTIFF](docs/docs/2.0.0-rc2/formats/geotiff.md). - **Minimum supported Rust is 1.94**, up from 1.91. `iceberg` and `iceberg-datafusion` 0.10 — the only release line built against the DataFusion 53 and Arrow 58 this workspace unifies on — declare `rust-version = "1.94"`, so the workspace floor follows. Beacon's own code uses no diff --git a/beacon-db/beacon-core/tests/read_functions.rs b/beacon-db/beacon-core/tests/read_functions.rs index 14d3c515..0cf0fae9 100644 --- a/beacon-db/beacon-core/tests/read_functions.rs +++ b/beacon-db/beacon-core/tests/read_functions.rs @@ -34,6 +34,8 @@ fn seed(rt: TestRuntime) -> TestRuntime { .expect("copy netcdf fixture as hdf5"); std::fs::copy(nested_hdf5_fixture(), rt.datasets_dir().join("nested.h5")) .expect("copy nested hdf5 fixture"); + std::fs::copy(tiff_fixture(), rt.datasets_dir().join("raster.tif")) + .expect("copy geotiff fixture"); rt } @@ -62,6 +64,19 @@ fn nested_hdf5_fixture() -> std::path::PathBuf { .join("beacon-file-formats/beacon-arrow-hdf5/test_files/nested-groups.h5") } +/// A stripped single-band GeoTIFF, shipped with the TIFF reader: 1287 x 380 +/// float32 pixels on the axes `x` and `y`. +fn tiff_fixture() -> std::path::PathBuf { + std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .parent() + .unwrap() + .join("beacon-file-formats/beacon-arrow-tiff/test-files/test.tif") +} + +/// The GeoTIFF fixture's grid: `y` (image rows) x `x` (image columns). +const TIFF_HEIGHT: i64 = 380; +const TIFF_WIDTH: i64 = 1287; + #[tokio::test(flavor = "multi_thread")] async fn read_csv_scans_filters_and_projects() { let rt = seeded("read-csv").await; @@ -197,6 +212,68 @@ async fn read_hdf5_schema_lists_the_nested_columns() { ); } +/// A GeoTIFF rides the same nd pipeline as netCDF, HDF5 and zarr: the raster is +/// a `y` x `x` grid, and the 1-d coordinate axes broadcast over it. +#[tokio::test(flavor = "multi_thread")] +async fn read_tiff_scans_the_fixture_as_a_grid() { + let rt = seeded("read-tiff").await; + + assert_eq!( + scalar_i64(&rt.sql("SELECT count(*) FROM read_tiff('raster.tif')").await), + TIFF_HEIGHT * TIFF_WIDTH, + "the raster is one row per pixel" + ); + + // `geo.lat` lives on `y` alone, so it is a broadcast column of the grid the + // full-rank band establishes. The band is co-selected because, as for every + // nd format, the projected columns are what define the grid. + let broadcast = rt + .sql( + r#"SELECT count("geo.lat") AS lat_rows, + count(DISTINCT "geo.lat") AS lat_values, + count("band.0") AS band_values + FROM read_tiff('raster.tif')"#, + ) + .await; + let column = |name: &str| { + let index = broadcast[0].schema().index_of(name).expect(name); + arrow::array::as_primitive_array::(broadcast[0].column(index)) + .value(0) + }; + assert_eq!(column("lat_rows"), TIFF_HEIGHT * TIFF_WIDTH); + assert_eq!(column("lat_values"), TIFF_HEIGHT); + // The band's nodata pixels come back as nulls, so it counts fewer. + assert!(column("band_values") > 0); + assert!(column("band_values") < column("lat_rows")); +} + +/// The optional second argument sets the grid for a raster too: `['y']` keeps +/// the latitude axis and drops the band, which needs both axes. +#[tokio::test(flavor = "multi_thread")] +async fn read_tiff_takes_a_dimensions_argument() { + let rt = seeded("read-tiff-dimensions").await; + + let narrowed = rt + .sql("SELECT * FROM read_tiff(['raster.tif'], ['y'])") + .await; + let columns: Vec = narrowed[0] + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(); + assert!(columns.contains(&"geo.lat".to_string()), "{columns:?}"); + assert!( + !columns.contains(&"band.0".to_string()), + "a 2-d band does not fit a 1-d grid: {columns:?}" + ); + assert_eq!( + total_rows(&narrowed) as i64, + TIFF_HEIGHT, + "one row for each image row" + ); +} + #[tokio::test(flavor = "multi_thread")] async fn read_of_a_missing_file_is_an_error_not_an_empty_result() { let rt = seeded("read-missing").await; diff --git a/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/mod.rs index cc8c2141..f50269b5 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/mod.rs @@ -39,9 +39,19 @@ impl FileFormatFactory for TiffFormatFactory { fn create( &self, _state: &dyn Session, - _format_options: &std::collections::HashMap, + format_options: &std::collections::HashMap, ) -> datafusion::error::Result> { - Ok(Arc::new(TiffFormat::new(self.options.clone()))) + // Per-table override from `CREATE EXTERNAL TABLE ... OPTIONS (...)`. + let read_dimensions = format_options.get("read_dimensions").map(|value| { + value + .split(',') + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect() + }); + Ok(Arc::new( + TiffFormat::new(self.options.clone()).with_read_dimensions(read_dimensions), + )) } fn default(&self) -> Arc { @@ -86,17 +96,48 @@ impl FileFormatFactoryExt for TiffFormatFactory { } } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub struct TiffFormat { pub options: TiffOptions, + /// Explicit dimensions requested via `read_tiff(paths, ['dims'])` or a + /// `CREATE EXTERNAL TABLE ... OPTIONS (read_dimensions '...')`. When set, + /// only variables whose dimensions are a subset of these are read; when + /// `None`, a broadcast-compatible default is auto-selected. + pub read_dimensions: Option>, } impl TiffFormat { pub fn new(options: TiffOptions) -> Self { - Self { options } + Self { + options, + read_dimensions: None, + } + } + + /// Returns a copy of this format that reads only the variables belonging to + /// `read_dimensions` (or auto-selects a default when `None`). + pub fn with_read_dimensions(mut self, read_dimensions: Option>) -> Self { + self.read_dimensions = read_dimensions; + self } } +/// Wrap a TIFF file scan in the nd spine: `NdBroadcastExec` → `NdSourceExec` → +/// `DataSourceExec`. +/// +/// The scan carries nd data as `beacon.nd`-encoded struct columns, so +/// `NdSourceExec` decodes it and `NdBroadcastExec` broadcasts it back to the +/// logical table schema above the scan. +fn nd_scan_plan(conf: FileScanConfig) -> datafusion::error::Result> { + let data_source: Arc = DataSourceExec::from_data_source(conf); + let nd_source = Arc::new(beacon_datafusion_ext::nd::exec::NdSourceExec::try_new( + data_source, + )?); + Ok(Arc::new( + beacon_datafusion_ext::nd::exec::NdBroadcastExec::try_new(nd_source)?, + )) +} + #[async_trait::async_trait] impl FileFormat for TiffFormat { fn as_any(&self) -> &dyn Any { @@ -126,7 +167,8 @@ impl FileFormat for TiffFormat { ) -> datafusion::error::Result { let mut tasks = vec![]; for object in objects { - let task = reader::fetch_schema(store.clone(), object.clone()); + let task = + reader::fetch_schema(store.clone(), object.clone(), self.read_dimensions.clone()); tasks.push(task); } @@ -159,42 +201,52 @@ impl FileFormat for TiffFormat { _state: &dyn Session, conf: FileScanConfig, ) -> datafusion::error::Result> { + // The scan carries nd data as `beacon.nd`-encoded struct columns, so the + // file source's schema is the encoded form of the logical table schema. + let encoded_file_schema = Arc::new(beacon_datafusion_ext::nd::encoded_schema( + conf.file_schema(), + )); let table_schema = datafusion::datasource::table_schema::TableSchema::new( - conf.file_schema().clone(), + encoded_file_schema, 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. let projection = conf.file_source().projection().cloned(); - let source = TiffSource::new(table_schema).with_projection(projection); + let source = TiffSource::new(table_schema) + .with_read_dimensions(self.read_dimensions.clone()) + .with_projection(projection); let conf = FileScanConfigBuilder::from(conf) .with_source(Arc::new(source)) .build(); - Ok(DataSourceExec::from_data_source(conf)) + nd_scan_plan(conf) } fn file_source( &self, table_schema: datafusion::datasource::table_schema::TableSchema, ) -> Arc { - Arc::new(TiffSource::new(table_schema)) + Arc::new(TiffSource::new(table_schema).with_read_dimensions(self.read_dimensions.clone())) } } #[cfg(test)] mod tests { use super::*; - use datafusion::datasource::physical_plan::{FileScanConfigBuilder, FileSource}; - use datafusion::execution::object_store::ObjectStoreUrl; - use futures::StreamExt; + use datafusion::physical_plan::displayable; + use datafusion::prelude::{SessionConfig, SessionContext}; + use object_store::ObjectStoreExt; use object_store::memory::InMemory; use object_store::path::Path; - use object_store::ObjectStoreExt; const TEST_TIF_BYTES: &[u8] = include_bytes!("../../test-files/test.tif"); + /// The bundled `test.tif` is 1287 × 380, single band, float32. + const WIDTH: usize = 1287; + const HEIGHT: usize = 380; + async fn put_fixture(store: &Arc, path: &Path, bytes: &[u8]) -> ObjectMeta { store .put(path, bytes::Bytes::copy_from_slice(bytes).into()) @@ -206,6 +258,50 @@ mod tests { .expect("should fetch object metadata") } + /// Register the bundled `test.tif` as a DataFusion table backed by + /// [`TiffFormat`] + `ListingTable` over the local filesystem. + async fn register_example_with(ctx: &SessionContext, format: TiffFormat) { + use datafusion::datasource::file_format::FileFormat; + use datafusion::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, + }; + + let file = concat!(env!("CARGO_MANIFEST_DIR"), "/test-files/test.tif"); + let table_path = ListingTableUrl::parse(format!("file://{file}")).unwrap(); + let format: Arc = Arc::new(format); + let listing_options = ListingOptions::new(format).with_file_extension("tif"); + 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("tiff_t", Arc::new(table)).unwrap(); + } + + async fn register_example(ctx: &SessionContext) { + register_example_with(ctx, TiffFormat::new(Default::default())).await; + } + + /// A session with the nd pushdown rules registered — the same wiring + /// beacon-core installs. Single partition so row order is deterministic + /// (the differential tests compare results positionally). + fn ctx_with_pushdown() -> SessionContext { + use datafusion::execution::session_state::SessionStateBuilder; + + let state = SessionStateBuilder::new() + .with_config(SessionConfig::new().with_target_partitions(1)) + .with_default_features() + .with_physical_optimizer_rule(Arc::new( + beacon_datafusion_ext::nd::NdProjectionPushdown::new(), + )) + .with_physical_optimizer_rule(Arc::new( + beacon_datafusion_ext::nd::NdFilterPushdown::new(), + )) + .build(); + SessionContext::new_with_state(state) + } + #[tokio::test] async fn infer_schema_reads_real_stripped_geotiff_fixture() { let store = Arc::new(InMemory::new()); @@ -213,221 +309,319 @@ mod tests { let path = Path::from("tests/datafusion/test.tif"); let object = put_fixture(&store, &path, TEST_TIF_BYTES).await; - let schema = reader::fetch_schema(object_store, object) + let schema = reader::fetch_schema(object_store, object, None) .await .expect("real stripped GeoTIFF should produce a schema"); let field_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - assert!( - field_names.contains(&"band.0"), - "schema should contain band.0" - ); - assert!( - field_names.contains(&"geo.lat"), - "schema should contain geo.lat" - ); - assert!( - field_names.contains(&"geo.lon"), - "schema should contain geo.lon" - ); - assert!( - field_names.contains(&"image.width"), - "schema should contain image.width" - ); - println!("Schema is: {:?}", schema); + for expected in ["band.0", "geo.lat", "geo.lon", "image.width"] { + assert!( + field_names.contains(&expected), + "schema should contain {expected}: {field_names:?}" + ); + } } + /// Explicit `read_dimensions` narrows the schema to the variables living on + /// the requested axis: `geo.lat` is on `y`, `geo.lon` on `x`, and the band on + /// both. Scalar metadata (rank-0) survives every narrowing. #[tokio::test] - async fn opener_streams_record_batches_for_real_fixture() { + async fn read_dimensions_narrows_the_schema_to_one_axis() { let store = Arc::new(InMemory::new()); let object_store: Arc = store.clone(); - let path = Path::from("tests/datafusion/test2.tif"); + let path = Path::from("tests/datafusion/test_dims.tif"); let object = put_fixture(&store, &path, TEST_TIF_BYTES).await; - let table_schema = reader::fetch_schema(object_store.clone(), object.clone()) + let schema = reader::fetch_schema(object_store, object, Some(vec!["y".to_string()])) .await - .expect("schema"); - - let ts = datafusion::datasource::table_schema::TableSchema::from_file_schema(table_schema); - let source = source::TiffSource::new(ts); - let file_opener = { - let conf = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("memory://").expect("url"), - Arc::new(source.clone()) as Arc, - ) - .build(); - source - .create_file_opener(object_store, &conf, 0) - .expect("file opener") - }; + .expect("narrowed schema"); - let stream = file_opener - .open(datafusion::datasource::listing::PartitionedFile::from(object)) - .expect("open") + let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + assert!(names.contains(&"geo.lat"), "geo.lat is on y: {names:?}"); + assert!(names.contains(&"image.width"), "scalars survive: {names:?}"); + assert!(!names.contains(&"geo.lon"), "geo.lon is on x: {names:?}"); + assert!(!names.contains(&"band.0"), "the band is on y,x: {names:?}"); + } + + /// Reading on the `y` axis alone makes the table one row per image row, + /// instead of the full `y × x` grid `count_star_counts_the_full_grid` sees. + #[tokio::test] + async fn read_dimensions_narrows_the_row_count_to_one_axis() { + use arrow::array::Int64Array; + + let ctx = SessionContext::new(); + register_example_with( + &ctx, + TiffFormat::new(Default::default()).with_read_dimensions(Some(vec!["y".to_string()])), + ) + .await; + + let batches = ctx + .sql("SELECT COUNT(*) AS n FROM tiff_t") + .await + .unwrap() + .collect() .await - .expect("stream"); + .unwrap(); + let n = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!( + n as usize, HEIGHT, + "the y axis alone has one row per image row" + ); + } + + // ── nd pipeline: plan shape ────────────────────────────────────────── + + /// The physical plan is the nd spine over the standard file scan: + /// `NdBroadcastExec` → `NdSourceExec` → `DataSourceExec`, in that nesting + /// order (parent above child in the indented render). + #[tokio::test] + async fn physical_plan_is_nd_spine_over_scan() { + let ctx = SessionContext::new(); + register_example(&ctx).await; - let batches: Vec<_> = stream - .collect::>() + let plan = ctx + .sql("SELECT \"band.0\" FROM tiff_t") .await - .into_iter() - .collect::, _>>() - .expect("all batches should be ok"); + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); - assert!(!batches.is_empty(), "should produce at least one batch"); + let broadcast = rendered.find("NdBroadcastExec"); + let source = rendered.find("NdSourceExec"); + let scan = rendered.find("DataSourceExec"); + assert!( + broadcast.is_some() && source.is_some() && scan.is_some(), + "plan must contain the nd spine over a DataSourceExec:\n{rendered}" + ); + assert!( + broadcast < source && source < scan, + "expected NdBroadcastExec → NdSourceExec → DataSourceExec nesting:\n{rendered}" + ); + } - // Concatenate into a single batch for easy column access. - let full = arrow::compute::concat_batches(&batches[0].schema(), &batches).expect("concat"); + /// With the projection rule registered, `SELECT "geo.lat" * 2` plans with an + /// `NdProjectionExec` *below* the `NdBroadcastExec` — so the arithmetic runs + /// on the 380-element latitude axis, not on all 380 × 1287 grid cells — and + /// produces the same values as a session without the rule. + #[tokio::test] + async fn projection_pushdown_fires_end_to_end() { + use arrow::compute::concat_batches; - let schema = full.schema(); + let sql = "SELECT \"geo.lat\" * 2 AS lat2 FROM tiff_t"; - // geo.lat column — values span ~30°N to ~46°N - let lat_idx = schema.index_of("geo.lat").expect("geo.lat column"); - let lat_col = full - .column(lat_idx) - .as_any() - .downcast_ref::() - .expect("geo.lat should be Float64"); - assert!(lat_col.len() > 0); - // First value: lat[0] = 0.04166667002172143 * 0 + 30.16666666498914 + let on = ctx_with_pushdown(); + register_example(&on).await; + let df = on.sql(sql).await.unwrap(); + let plan = df.clone().create_physical_plan().await.unwrap(); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); + + let broadcast = rendered.find("NdBroadcastExec"); + let projection = rendered.find("NdProjectionExec"); + let source = rendered.find("NdSourceExec"); assert!( - (lat_col.value(0) - 30.166_666_664_989_14).abs() < 1e-6, - "lat[0]={}", - lat_col.value(0) + projection.is_some() && broadcast < projection && projection < source, + "expected NdBroadcastExec → NdProjectionExec → NdSourceExec:\n{rendered}" ); - // All values should be within the expected geographic range. - for i in 0..lat_col.len() { - let v = lat_col.value(i); - assert!(v >= 30.0 && v <= 47.0, "lat[{i}]={v} out of range"); - } + let actual = df.collect().await.unwrap(); - // geo.lon column — values span ~-17°E to ~36°E - let lon_idx = schema.index_of("geo.lon").expect("geo.lon column"); - let lon_col = full - .column(lon_idx) - .as_any() - .downcast_ref::() - .expect("geo.lon should be Float64"); - assert!(lon_col.len() > 0); - // First value: lon[0] = 0.0416666671610546 * 0 + -17.312499364464315 + // Same single-partition config so row order matches positionally. + let off = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + register_example(&off).await; + let expected = off.sql(sql).await.unwrap().collect().await.unwrap(); + + let schema = expected[0].schema(); + assert_eq!( + concat_batches(&schema, &actual).unwrap(), + concat_batches(&schema, &expected).unwrap(), + ); + } + + /// With the filter rule registered, `WHERE "geo.lat" > 40` sinks into an + /// `NdFilterExec` below the broadcast — the grid is selected before it is + /// materialized — and the rows match the unoptimized session. + #[tokio::test] + async fn filter_pushdown_fires_end_to_end() { + use arrow::compute::concat_batches; + + let sql = "SELECT \"geo.lat\" FROM tiff_t WHERE \"geo.lat\" > 40"; + + let on = ctx_with_pushdown(); + register_example(&on).await; + let df = on.sql(sql).await.unwrap(); + let plan = df.clone().create_physical_plan().await.unwrap(); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); + + let broadcast = rendered.find("NdBroadcastExec"); + let filter = rendered.find("NdFilterExec"); + let source = rendered.find("NdSourceExec"); assert!( - (lon_col.value(0) - -17.312_499_364_464_315).abs() < 1e-6, - "lon[0]={}", - lon_col.value(0) + filter.is_some() && broadcast < filter && filter < source, + "expected NdBroadcastExec → NdFilterExec → NdSourceExec:\n{rendered}" + ); + let actual = df.collect().await.unwrap(); + + let off = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + register_example(&off).await; + let expected = off.sql(sql).await.unwrap().collect().await.unwrap(); + + let schema = expected[0].schema(); + assert_eq!( + concat_batches(&schema, &actual).unwrap(), + concat_batches(&schema, &expected).unwrap(), ); - // All values should be within the expected geographic range. - for i in 0..lon_col.len() { - let v = lon_col.value(i); - assert!(v >= -18.0 && v <= 37.0, "lon[{i}]={v} out of range"); - } } + // ── end-to-end reads ───────────────────────────────────────────────── + + /// The two coordinate axes — `geo.lat` on `y`, `geo.lon` on `x` — broadcast + /// against each other into their full cross product, with the values the + /// GeoTIFF's ModelTransformation tag defines. + /// + /// Assertions are order-independent on purpose: the nd spine derives the + /// grid's axis order from the widest projected column, so two same-rank + /// coordinates leave the row order unspecified (as for every nd format). #[tokio::test] - async fn opener_with_predicate_filters_rows() { - use datafusion::config::ConfigOptions; - use datafusion::datasource::physical_plan::FileSource; - use datafusion::logical_expr::Operator; - use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; - use datafusion::scalar::ScalarValue; + async fn end_to_end_reads_broadcast_coordinates() { + use arrow::array::{Float64Array, Int64Array}; - let store = Arc::new(InMemory::new()); - let object_store: Arc = store.clone(); - let path = Path::from("tests/datafusion/test_pred.tif"); - let object = put_fixture(&store, &path, TEST_TIF_BYTES).await; + let ctx = SessionContext::new(); + register_example(&ctx).await; - let table_schema = reader::fetch_schema(object_store.clone(), object.clone()) + let batches = ctx + .sql( + r#"SELECT COUNT(*) AS rows, + COUNT(DISTINCT "geo.lat") AS lats, + COUNT(DISTINCT "geo.lon") AS lons, + MIN("geo.lat") AS lat_min, + MAX("geo.lat") AS lat_max, + MIN("geo.lon") AS lon_min + FROM tiff_t"#, + ) .await - .expect("schema"); - - // Build predicate: geo.lat > 40.0 - // The Column index must match geo.lat's position in the file schema. - let lat_idx = table_schema.index_of("geo.lat").expect("geo.lat field"); - let predicate: Arc = - Arc::new(BinaryExpr::new( - Arc::new(Column::new("geo.lat", lat_idx)), - Operator::Gt, - Arc::new(Literal::new(ScalarValue::Float64(Some(40.0)))), - )); - - // Push the predicate into a TiffSource via try_pushdown_filters. - let source_with_predicate: Arc = { - let ts = datafusion::datasource::table_schema::TableSchema::from_file_schema( - table_schema.clone(), - ); - let base_source = source::TiffSource::new(ts); - let pushdown = base_source - .try_pushdown_filters(vec![predicate], &ConfigOptions::default()) - .expect("try_pushdown_filters"); - pushdown.updated_node.expect("updated node with predicate") - }; + .unwrap() + .collect() + .await + .unwrap(); - let file_opener = { - let conf = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("memory://").expect("url"), - source_with_predicate.clone(), - ) - .build(); - source_with_predicate - .create_file_opener(object_store, &conf, 0) - .expect("file opener") + let row = &batches[0]; + let int = |name: &str| { + row.column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + }; + let float = |name: &str| { + row.column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0) }; - let stream = file_opener - .open(datafusion::datasource::listing::PartitionedFile::from(object)) - .expect("open") - .await - .expect("stream"); + // Each axis contributes its own extent, and the table is their product. + assert_eq!(int("lats") as usize, HEIGHT); + assert_eq!(int("lons") as usize, WIDTH); + assert_eq!(int("rows") as usize, HEIGHT * WIDTH); - let batches: Vec<_> = stream - .collect::>() - .await - .into_iter() - .collect::, _>>() - .expect("all batches should be ok"); + // ModelTransformationTag: lat[y] = 0.04166667002172143 * y + 30.16666666498914 + // lon[x] = 0.0416666671610546 * x + -17.312499364464315 + assert!((float("lat_min") - 30.166_666_664_989_14).abs() < 1e-6); + assert!((float("lat_max") - 45.958_334_603_221_566).abs() < 1e-6); + assert!((float("lon_min") - -17.312_499_364_464_315).abs() < 1e-6); + } - assert!(!batches.is_empty(), "should produce at least one batch"); + /// A rank-0 metadata scalar (`image.width`) rides the nd encoding and + /// broadcasts to a constant column over every row of the grid its + /// co-selected variable establishes — here `band.0`, the full `y × x` band. + #[tokio::test] + async fn end_to_end_broadcasts_scalar_metadata() { + use arrow::array::Int64Array; + + let ctx = SessionContext::new(); + register_example(&ctx).await; - let full = arrow::compute::concat_batches(&batches[0].schema(), &batches).expect("concat"); + let batches = ctx + .sql( + r#"SELECT COUNT(DISTINCT "image.width") AS distinct_widths, + COUNT("image.width") AS scalar_rows, + COUNT("geo.lat") AS coord_rows, + COUNT("band.0") AS band_values + FROM tiff_t"#, + ) + .await + .unwrap() + .collect() + .await + .unwrap(); - // Predicate pushdown here is coarse-grained: entire chunks whose coordinate range - // falls entirely outside the predicate are skipped (no I/O). Chunks that partially - // overlap are emitted in full. We therefore only verify that I/O was reduced, not - // that every row satisfies the predicate. - let total_rows = 380 * 1287; + let int = |name: &str| { + batches[0] + .column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + }; + assert_eq!(int("distinct_widths"), 1, "a scalar is a single constant"); + // `band.0` is the widest column, so the grid is the full image. + assert_eq!( + int("scalar_rows") as usize, + HEIGHT * WIDTH, + "the scalar must be broadcast onto every grid row" + ); + assert_eq!( + int("coord_rows") as usize, + HEIGHT * WIDTH, + "the latitude axis must be broadcast onto every grid row" + ); + // The band's nodata pixels come back as nulls, so it counts fewer. + assert!(int("band_values") > 0); assert!( - full.num_rows() < total_rows, - "predicate should skip at least one chunk, reducing row count below {total_rows} (got {})", - full.num_rows() + int("band_values") < int("scalar_rows"), + "the fixture's nodata pixels must be null" ); - assert!(full.num_rows() > 0, "predicate should keep some rows"); } - // ── End-to-end via SessionContext (projection + predicate pushdown) ── + /// `COUNT(*)` projects no columns, so the opener drives the read with the + /// highest-volume variable and reports the full broadcast row count. + #[tokio::test] + async fn count_star_counts_the_full_grid() { + use arrow::array::Int64Array; - /// Register the bundled `test.tif` as a DataFusion table backed by - /// [`TiffFormat`] + `ListingTable` over the local filesystem. - async fn register_example(ctx: &datafusion::prelude::SessionContext) { - use datafusion::datasource::file_format::FileFormat; - use datafusion::datasource::listing::{ - ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, - }; + let ctx = SessionContext::new(); + register_example(&ctx).await; - let file = concat!(env!("CARGO_MANIFEST_DIR"), "/test-files/test.tif"); - let table_path = ListingTableUrl::parse(format!("file://{file}")).unwrap(); - let format: Arc = Arc::new(TiffFormat::new(Default::default())); - let listing_options = ListingOptions::new(format).with_file_extension("tif"); - let config = ListingTableConfig::new(table_path) - .with_listing_options(listing_options) - .infer_schema(&ctx.state()) + let batches = ctx + .sql("SELECT COUNT(*) AS n FROM tiff_t") + .await + .unwrap() + .collect() .await .unwrap(); - let table = ListingTable::try_new(config).unwrap(); - ctx.register_table("tiff_t", Arc::new(table)).unwrap(); + let n = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(n as usize, HEIGHT * WIDTH); } #[tokio::test] async fn projection_pushdown_through_datafusion() { - let ctx = datafusion::prelude::SessionContext::new(); + let ctx = SessionContext::new(); register_example(&ctx).await; let df = ctx @@ -447,12 +641,12 @@ mod tests { let batches = df.collect().await.unwrap(); assert_eq!(batches[0].num_columns(), 2); let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert!(rows > 0); + assert_eq!(rows, HEIGHT * WIDTH); } #[tokio::test] - async fn predicate_pushdown_prunes_through_datafusion() { - let ctx = datafusion::prelude::SessionContext::new(); + async fn predicate_prunes_every_row_through_datafusion() { + let ctx = SessionContext::new(); register_example(&ctx).await; // Latitude never exceeds ~47°, so this predicate excludes every row. @@ -466,12 +660,15 @@ mod tests { .iter() .map(|b| b.num_rows()) .sum(); - assert_eq!(rows, 0, "impossible latitude predicate should yield no rows"); + assert_eq!( + rows, 0, + "impossible latitude predicate should yield no rows" + ); } #[tokio::test] - async fn predicate_pushdown_selects_subset_through_datafusion() { - let ctx = datafusion::prelude::SessionContext::new(); + async fn predicate_selects_subset_through_datafusion() { + let ctx = SessionContext::new(); register_example(&ctx).await; let batches = ctx @@ -490,12 +687,15 @@ mod tests { .downcast_ref::() .expect("geo.lat is Float64"); for i in 0..col.len() { - assert!(col.value(i) > 40.0, "every returned lat must satisfy the predicate"); + assert!( + col.value(i) > 40.0, + "every returned lat must satisfy the predicate" + ); } total += b.num_rows(); } assert!(total > 0, "satisfiable predicate should keep some rows"); - assert!(total < 380 * 1287, "predicate should drop some rows"); + assert!(total < HEIGHT * WIDTH, "predicate should drop some rows"); } } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/reader.rs b/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/reader.rs index e36192bb..a4c65241 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/reader.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/reader.rs @@ -12,9 +12,16 @@ pub async fn open_dataset( } /// Fetch the Arrow schema for a TIFF object. +/// +/// When `read_dimensions` is provided the dataset is projected to only include +/// variables that belong to those dimensions before deriving the Arrow schema. +/// When it is absent, a broadcast-compatible default dimension set is +/// auto-selected (see [`beacon_nd_array::dataset::resolve_read_dimensions`]) so +/// the schema matches what `SELECT *` can actually return. pub async fn fetch_schema( object_store: Arc, object: ObjectMeta, + read_dimensions: Option>, ) -> datafusion::error::Result { let dataset = open_dataset(object_store, object).await.map_err(|e| { datafusion::error::DataFusionError::Execution(format!( @@ -22,6 +29,22 @@ pub async fn fetch_schema( )) })?; + let dataset = if let Some(dims) = beacon_nd_array::dataset::resolve_read_dimensions( + &dataset, + read_dimensions, + Some("read_tiff"), + ) { + dataset + .project(&DatasetProjection::new_with_dimension_projection(dims)) + .map_err(|e| { + datafusion::error::DataFusionError::Execution(format!( + "Failed to project TIFF dataset with dimensions: {e}" + )) + })? + } else { + dataset + }; + let schema = beacon_nd_array::arrow::schema::any_dataset_to_arrow_schema(&dataset).map_err(|e| { datafusion::error::DataFusionError::Execution(format!( diff --git a/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/source.rs b/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/source.rs index d28ee23b..2cb9e89a 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/source.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/source.rs @@ -1,16 +1,24 @@ +//! DataFusion [`FileSource`]/[`FileOpener`] for TIFF/GeoTIFF files. +//! +//! The opener builds an [`AnyDataset`](beacon_nd_array::dataset::AnyDataset) +//! for the (projected) columns and emits `beacon.nd`-encoded batches, which the +//! `NdSourceExec`/`NdBroadcastExec` pair above the scan decodes and broadcasts. +//! This mirrors the netCDF, HDF5 and zarr sources. + use std::sync::Arc; -use arrow::{datatypes::SchemaRef, record_batch::RecordBatch}; +use arrow::{ + datatypes::SchemaRef, + record_batch::{RecordBatch, RecordBatchOptions}, +}; use beacon_nd_array::{ arrow::{ - batch::{any_dataset_as_record_batch_stream, any_dataset_as_row_size}, - metrics::DatasetReadMetrics, - pushdown_filter::PushdownFilter, + batch::any_dataset_as_record_batch_stream, metrics::DatasetReadMetrics, + nd_provider::any_dataset_as_encoded_stream, pushdown_filter::PushdownFilter, }, projection::DatasetProjection, }; use datafusion::{ - common::Statistics, config::ConfigOptions, datasource::{ listing::PartitionedFile, @@ -18,13 +26,12 @@ use datafusion::{ schema_adapter::SchemaAdapterFactory, table_schema::TableSchema, }, - execution::SendableRecordBatchStream, + error::DataFusionError, physical_expr::{PhysicalExpr, conjunction, projection::ProjectionExprs}, physical_expr_adapter::BatchAdapterFactory, physical_plan::{ filter_pushdown::{FilterPushdownPropagation, PushedDown}, - metrics::{ExecutionPlanMetricsSet, SplitMetrics}, - stream::{BatchSplitStream, RecordBatchStreamAdapter}, + metrics::ExecutionPlanMetricsSet, }, }; use futures::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream}; @@ -39,6 +46,8 @@ pub struct TiffSource { execution_plan_metrics: ExecutionPlanMetricsSet, batch_size: usize, predicate: Option>, + /// Explicit dimensions to read, or `None` to auto-select a default. + read_dimensions: Option>, /// Projection pushed down by the scan, applied on top of the table schema. projection: Option, } @@ -51,10 +60,18 @@ impl TiffSource { execution_plan_metrics: ExecutionPlanMetricsSet::new(), batch_size: 128 * 1024, predicate: None, + read_dimensions: None, projection: None, } } + /// Returns a copy of this source that reads only the variables belonging to + /// `read_dimensions` (or auto-selects a default when `None`). + pub fn with_read_dimensions(mut self, read_dimensions: Option>) -> Self { + self.read_dimensions = read_dimensions; + 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`. @@ -71,13 +88,12 @@ impl FileSource for TiffSource { base_config: &FileScanConfig, partition: usize, ) -> datafusion::error::Result> { - let file_schema = self.table_schema.file_schema().clone(); let projected_schema = base_config.projected_schema()?; Ok(Arc::new(TiffOpener::new( - file_schema, object_store, projected_schema, + self.read_dimensions.clone(), self.batch_size, self.predicate.clone(), self.execution_plan_metrics.clone(), @@ -164,10 +180,12 @@ impl FileSource for TiffSource { } } +// ─── FileOpener ────────────────────────────────────────────────────────────── + struct TiffOpener { - table_schema: SchemaRef, object_store: Arc, projected_schema: SchemaRef, + read_dimensions: Option>, batch_size: usize, predicate: Option>, partition: usize, @@ -175,19 +193,20 @@ struct TiffOpener { } impl TiffOpener { + #[allow(clippy::too_many_arguments)] fn new( - table_schema: SchemaRef, object_store: Arc, projected_schema: SchemaRef, + read_dimensions: Option>, batch_size: usize, predicate: Option>, metrics: ExecutionPlanMetricsSet, partition: usize, ) -> Self { Self { - table_schema, object_store, projected_schema, + read_dimensions, batch_size, predicate, partition, @@ -195,10 +214,12 @@ impl TiffOpener { } } + #[allow(clippy::too_many_arguments)] async fn read_task( object: ObjectMeta, object_store: Arc, projected_schema: SchemaRef, + read_dimensions: Option>, batch_size: usize, predicate: Option>, metrics: Option, @@ -206,16 +227,34 @@ impl TiffOpener { let dataset = reader::open_dataset(object_store, object.clone()) .await .map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( + DataFusionError::Execution(format!( "Failed to open TIFF dataset {}: {e}", object.location, )) })?; + // Apply the dimension projection before deriving the file schema. With + // no explicit dimensions, fall back to the dataset's auto-selected + // default (matching `fetch_schema`). No log label here: this runs per + // file/partition, so logging would spam. + let read_dimensions = + beacon_nd_array::dataset::resolve_read_dimensions(&dataset, read_dimensions, None); + let dataset = if let Some(dims) = read_dimensions { + dataset + .project(&DatasetProjection::new_with_dimension_projection(dims)) + .map_err(|e| { + DataFusionError::Execution(format!( + "Failed to project TIFF dataset with dimensions: {e}" + )) + })? + } else { + dataset + }; + let file_schema: SchemaRef = beacon_nd_array::arrow::schema::any_dataset_to_arrow_schema(&dataset) .map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( + DataFusionError::Execution(format!( "Failed to derive Arrow schema from TIFF dataset: {e}" )) })? @@ -232,51 +271,92 @@ impl TiffOpener { .collect(); if projection.is_empty() { - return Ok(any_dataset_as_row_size(dataset) - .map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to compute row size for empty projection on TIFF dataset: {e}" - )) - })? + // No output columns are needed (e.g. `COUNT(*)`). Reading zero + // columns would yield an empty stream and an incorrect count of 0. + // Drive the read with the highest-volume variable so the row count + // equals the full broadcast row count (a scalar like `image.width` + // would give just 1 row), plus any predicate columns so a + // pushed-down filter still applies (PushdownFilter matches by + // name). Emit zero-column batches carrying the correct row counts. + let driver_idx = dataset + .fields() + .keys() + .max_by_key(|name| { + dataset + .get_array(name) + .map(|a| a.shape().iter().product::()) + .unwrap_or(0) + }) + .and_then(|name| file_schema.index_of(name).ok()) + .unwrap_or(0); + let mut driver: Vec = vec![driver_idx]; + if let Some(pred) = &predicate { + for col in datafusion::physical_expr::utils::collect_columns(pred) { + if let Ok(idx) = file_schema.index_of(col.name()) { + driver.push(idx); + } + } + } + driver.sort_unstable(); + driver.dedup(); + + let dataset = dataset + .project(&DatasetProjection::new_with_index_projection(driver)) .map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to read TIFF dataset with empty projection: {e}" + DataFusionError::Execution(format!( + "Failed to project TIFF dataset for count: {e}" )) - }) - .boxed()); + })?; + + let pushdown_filter = predicate.map(PushdownFilter::new); + let count_schema = projected_schema.clone(); + let stream = + any_dataset_as_record_batch_stream(dataset, batch_size, pushdown_filter, metrics) + .map(move |batch| { + let batch = batch.map_err(|e| { + DataFusionError::Execution(format!( + "Error reading TIFF as Arrow stream: {e}" + )) + })?; + RecordBatch::try_new_with_options( + count_schema.clone(), + vec![], + &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + ) + .map_err(|e| { + DataFusionError::Execution(format!("Failed to build count batch: {e}")) + }) + }) + .boxed(); + return Ok(stream); } - // Adapt batches (read with `projection`) onto the projected output - // schema: reorder, cast, and null-fill columns this file lacks. - let source_schema: SchemaRef = Arc::new(file_schema.project(&projection)?); + // The opener emits nd-encoded batches, so adaptation happens in the + // encoded (struct) domain: reorder, and null-fill columns this file + // lacks, onto the projected encoded schema. + let source_schema: SchemaRef = Arc::new(beacon_datafusion_ext::nd::encoded_schema( + &file_schema.project(&projection)?, + )); let adapter = BatchAdapterFactory::new(projected_schema).make_adapter(&source_schema)?; let dataset = if projection.len() < file_schema.fields().len() { - let proj = DatasetProjection { - dimension_projection: None, - index_projection: Some(projection), - }; - dataset.project(&proj).map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to project TIFF dataset: {e}" - )) - })? + dataset + .project(&DatasetProjection::new_with_index_projection(projection)) + .map_err(|e| { + DataFusionError::Execution(format!("Failed to project TIFF dataset: {e}")) + })? } else { dataset }; - let pushdown_filter = 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 TIFF as Arrow stream: {e}" - )) - }) + // Emit nd-encoded batches (decoded/broadcast by the NdSourceExec / + // NdBroadcastExec above the scan), adapted onto the projected encoded + // schema. + let _ = metrics; + let stream = any_dataset_as_encoded_stream(dataset, batch_size) .and_then(move |batch| { let mapped = adapter.adapt_batch(&batch).map_err(|e| { - datafusion::error::DataFusionError::Execution(format!( - "Failed to adapt TIFF batch schema: {e}" - )) + DataFusionError::Execution(format!("Failed to adapt TIFF batch schema: {e}")) }); futures::future::ready(mapped) }) @@ -293,6 +373,7 @@ impl FileOpener for TiffOpener { file.object_meta, self.object_store.clone(), self.projected_schema.clone(), + self.read_dimensions.clone(), self.batch_size, self.predicate.clone(), metrics, diff --git a/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/table_function.rs b/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/table_function.rs index 9df21268..413bb1f6 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/table_function.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-tiff/src/datafusion/table_function.rs @@ -5,7 +5,8 @@ use arrow::datatypes::{DataType, Field}; use beacon_common::super_table::SuperListingTable; use beacon_datafusion_ext::listing_factory::ListingFactory; use datafusion::{ - catalog::TableFunctionImpl, execution::object_store::ObjectStoreUrl, prelude::SessionContext, + catalog::TableFunctionImpl, common::plan_err, prelude::Expr, prelude::SessionContext, + scalar::ScalarValue, }; use beacon_common::table_function::BeaconTableFunctionImpl; @@ -44,14 +45,40 @@ impl BeaconTableFunctionImpl for ReadTiffFunc { } fn arguments(&self) -> Option> { - Some(vec![Field::new( - "glob_paths", - DataType::List(Arc::new(Field::new("glob_path", DataType::Utf8, false))), - false, - )]) + Some(vec![ + Field::new( + "glob_paths", + DataType::List(Arc::new(Field::new("glob_path", DataType::Utf8, false))), + false, + ), + Field::new( + "dimensions", + DataType::List(Arc::new(Field::new("dimension", DataType::Utf8, false))), + true, + ), + ]) } } +/// The optional second argument: the dimensions to read. +fn parse_dimensions_arg(args: &[Expr]) -> datafusion::error::Result> { + let Some(Expr::Literal(ScalarValue::List(values), _)) = args.get(1) else { + return Ok(vec![]); + }; + let Some(strings) = values + .as_ref() + .values() + .as_any() + .downcast_ref::() + else { + return plan_err!("read_tiff second argument must be a List of dimension names"); + }; + Ok(strings + .iter() + .filter_map(|value| value.map(|s| s.to_string())) + .collect()) +} + impl TableFunctionImpl for ReadTiffFunc { fn call( &self, @@ -70,13 +97,15 @@ impl TableFunctionImpl for ReadTiffFunc { ) })?; let glob_paths = beacon_common::table_function::parse_glob_paths_arg(args, "read_tiff")?; + let dimensions = parse_dimensions_arg(args)?; let mut listing_urls = vec![]; for path in &glob_paths { listing_urls.push(listing_factory.parse_listing_table_url(&state, path)?); } - let file_format = TiffFormat::new(Default::default()); + let file_format = TiffFormat::new(Default::default()) + .with_read_dimensions((!dimensions.is_empty()).then_some(dimensions)); let super_listing_table = tokio::task::block_in_place(|| { self.runtime_handle.block_on(async move { diff --git a/docs/docs/2.0.0-rc2/formats/geotiff.md b/docs/docs/2.0.0-rc2/formats/geotiff.md index dbfedcda..5392b327 100644 --- a/docs/docs/2.0.0-rc2/formats/geotiff.md +++ b/docs/docs/2.0.0-rc2/formats/geotiff.md @@ -8,12 +8,22 @@ description: Read GeoTIFF and Cloud-Optimized GeoTIFF rasters with read_tiff(). ```text read_tiff(glob_paths) +read_tiff(glob_paths, dimensions) ``` Beacon reads GeoTIFF and Cloud-Optimized GeoTIFF files. +A raster is a grid over the axes `y` (image rows) and `x` (image columns). Beacon returns the bands +on that grid, and broadcasts the coordinate axes `geo.lat` and `geo.lon` over it. + +The optional `dimensions` argument selects the columns. Beacon returns a column only if the list +holds all of its dimensions. Use `['y']` to read the latitude axis alone, without the full grid. + ```sql SELECT * FROM read_tiff('rasters/elevation.tif') + +-- One row per image row, not one row per pixel +SELECT "geo.lat" FROM read_tiff('rasters/elevation.tif', ['y']) ``` ## Inspect the schema diff --git a/docs/docs/2.0.0-rc2/sql/table-functions.md b/docs/docs/2.0.0-rc2/sql/table-functions.md index 6586645c..e58058cf 100644 --- a/docs/docs/2.0.0-rc2/sql/table-functions.md +++ b/docs/docs/2.0.0-rc2/sql/table-functions.md @@ -214,12 +214,20 @@ SELECT * FROM read_bbf('bbf/**/*.bbf') ```text read_tiff(glob_paths) +read_tiff(glob_paths, dimensions) ``` -Beacon reads GeoTIFF and Cloud-Optimized GeoTIFF files. +Beacon reads GeoTIFF and Cloud-Optimized GeoTIFF files. A raster is a grid over the axes `y` (image +rows) and `x` (image columns). + +The optional `dimensions` argument selects the columns. Beacon returns a column only if the list +holds all of its dimensions. Use `['y']` to read the latitude axis alone, without the full grid. ```sql SELECT * FROM read_tiff('rasters/elevation.tif') + +-- One row per image row, not one row per pixel +SELECT "geo.lat" FROM read_tiff('rasters/elevation.tif', ['y']) ``` ### Tag attributes From b504c34000686c25ba98c566a792eb0be58d473e Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 12 Aug 2026 10:50:48 +0200 Subject: [PATCH 2/2] Name GeoTIFF among the formats that stack nodes over their scan The scan-source docs list which formats return more than a bare `DataSourceExec`, because that is why both the metric recorder and the pruner descend the single-child chain. GeoTIFF joined that set in this branch. --- beacon-db/beacon-datafusion-ext/src/fast_object/mod.rs | 8 ++++---- beacon-db/beacon-datafusion-ext/src/fast_object/table.rs | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/beacon-db/beacon-datafusion-ext/src/fast_object/mod.rs b/beacon-db/beacon-datafusion-ext/src/fast_object/mod.rs index aa943dab..bab2cc90 100644 --- a/beacon-db/beacon-datafusion-ext/src/fast_object/mod.rs +++ b/beacon-db/beacon-datafusion-ext/src/fast_object/mod.rs @@ -38,10 +38,10 @@ //! # The format still plans its own scan //! //! `create_physical_plan` is what turns the config into a plan, so every format -//! keeps the shape it wants: netCDF and HDF5 stack decode and broadcast nodes -//! over their scan, and Zarr and Atlas expand a store directory into partitions -//! and reduce it to the marker at its root. Nothing here knows about any of -//! that. +//! keeps the shape it wants: netCDF, HDF5, Zarr and GeoTIFF stack decode and +//! broadcast nodes over their scan, and Zarr and Atlas expand a store directory +//! into partitions and reduce it to the marker at its root. Nothing here knows +//! about any of that. //! //! # What `EXPLAIN` shows //! diff --git a/beacon-db/beacon-datafusion-ext/src/fast_object/table.rs b/beacon-db/beacon-datafusion-ext/src/fast_object/table.rs index 5ab597a6..07ee8bde 100644 --- a/beacon-db/beacon-datafusion-ext/src/fast_object/table.rs +++ b/beacon-db/beacon-datafusion-ext/src/fast_object/table.rs @@ -281,8 +281,8 @@ impl FastObjectTable { /// blocking a later repartition or limit pushdown. /// /// The scan is found by descending the single-child chain, because an nd format -/// returns a stack: netCDF and HDF5 hand back decode and broadcast nodes above -/// their scan. +/// returns a stack: netCDF, HDF5, Zarr and GeoTIFF hand back decode and +/// broadcast nodes above their scan. fn record_counters(plan: &Arc, considered: usize, dropped: usize) { let mut node: &dyn ExecutionPlan = plan.as_ref(); loop {