From be785e9620e29993d624fcbc8678ad06ea0ad864 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Wed, 9 Sep 2026 17:38:39 -0700 Subject: [PATCH 1/3] feat(rust/sedona-raster): metadata-only byte estimates for rasters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `sedona_raster::size`: `estimated_band_bytes`, `estimated_raster_bytes` and `estimated_row_bytes` compute Σ_bands Π raw_source_shape × pixel bytes from band metadata alone. For an InDb band that equals the buffer length; for an OutDb band it is the size loading it will allocate — which is what a memory budget needs to know before anything is loaded. The raw source shape (not the visible shape) is used so a broadcast view doesn't inflate the figure and a slice doesn't shrink it; bands sharing one buffer are each counted in full, the safe direction for a budget. Refs apache/sedona-db#1220 --- rust/sedona-raster/src/array.rs | 25 +++ rust/sedona-raster/src/lib.rs | 1 + rust/sedona-raster/src/size.rs | 320 ++++++++++++++++++++++++++++++++ 3 files changed, 346 insertions(+) create mode 100644 rust/sedona-raster/src/size.rs diff --git a/rust/sedona-raster/src/array.rs b/rust/sedona-raster/src/array.rs index 90e339f38d..16690eea8e 100644 --- a/rust/sedona-raster/src/array.rs +++ b/rust/sedona-raster/src/array.rs @@ -932,6 +932,31 @@ impl<'a> RasterStructArray<'a> { pub fn band_data_row(&self, raster_idx: usize, band_idx: usize) -> usize { self.bands_list.value_offsets()[raster_idx] as usize + band_idx } + + /// Rows of `raster_idx`'s bands in the flattened band columns. Lets + /// metadata-only passes (e.g. `crate::size`) walk every band without + /// constructing a `RasterRef` / boxed `BandRef` per row. + pub(crate) fn band_rows(&self, raster_idx: usize) -> std::ops::Range { + let offsets = self.bands_list.value_offsets(); + offsets[raster_idx] as usize..offsets[raster_idx + 1] as usize + } + + /// Raw source shape of the band at flattened row `band_row`. + pub(crate) fn band_source_shape_at(&self, band_row: usize) -> &[i64] { + let offsets = self.band_source_shape_list.value_offsets(); + &self.band_source_shape_values.values() + [offsets[band_row] as usize..offsets[band_row + 1] as usize] + } + + /// Pixel type of the band at flattened row `band_row`. + pub(crate) fn band_data_type_at(&self, band_row: usize) -> Result { + let code = self.band_datatype_array.value(band_row); + BandDataType::try_from_u32(code).ok_or_else(|| { + RasterError::Invalid(format!( + "invalid band data type code {code} at band row {band_row}" + )) + }) + } } #[cfg(test)] diff --git a/rust/sedona-raster/src/lib.rs b/rust/sedona-raster/src/lib.rs index 154dc62888..4070e7b669 100644 --- a/rust/sedona-raster/src/lib.rs +++ b/rust/sedona-raster/src/lib.rs @@ -23,5 +23,6 @@ pub mod display; pub mod error; pub mod geo_transform; pub mod raster_loader; +pub mod size; pub mod traits; pub mod view_entries; diff --git a/rust/sedona-raster/src/size.rs b/rust/sedona-raster/src/size.rs new file mode 100644 index 0000000000..d0b6eaeb77 --- /dev/null +++ b/rust/sedona-raster/src/size.rs @@ -0,0 +1,320 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Metadata-only byte estimates for rasters. +//! +//! Every estimate here is `Σ_bands Π raw_source_shape × pixel bytes`, +//! computed from band metadata alone — the `data` column is never read. +//! That makes one formula serve both storage kinds: for an InDb band it is +//! an identity with the actual buffer length (the read boundary rejects a +//! mismatch as corruption), and for an OutDb band it is the size loading +//! it will allocate, which is what a memory budget needs to know *before* +//! anything is loaded. +//! +//! The estimate uses the **raw source shape, not the visible shape**: +//! resident bytes follow the source buffer, so a broadcast view does not +//! inflate the figure and a slice view does not shrink it. Bands that share +//! one buffer (zero-copy derivations) are each counted in full — an +//! over-estimate, which is the safe direction for a budget. + +use crate::array::RasterStructArray; +use crate::error::RasterError; +use crate::traits::{BandRef, RasterRef}; +use sedona_schema::raster::BandDataType; + +/// `Π source_shape × data_type.byte_size()`, overflow-checked. +fn band_bytes(source_shape: &[i64], data_type: BandDataType) -> Result { + source_shape + .iter() + .try_fold(data_type.byte_size() as u64, |acc, &dim| { + u64::try_from(dim).ok().and_then(|dim| acc.checked_mul(dim)) + }) + .ok_or_else(|| { + RasterError::Invalid(format!( + "band byte count overflows u64 for source_shape {source_shape:?} × {data_type:?}" + )) + }) +} + +/// Bytes the band's source buffer holds (InDb) or will hold once loaded +/// (OutDb): `Π raw_source_shape × data_type.byte_size()`. +pub fn estimated_band_bytes(band: &dyn BandRef) -> Result { + band_bytes(band.raw_source_shape(), band.data_type()) +} + +/// Sum of [`estimated_band_bytes`] over the raster's bands. +pub fn estimated_raster_bytes(raster: &dyn RasterRef) -> Result { + let mut total = 0u64; + for band_idx in 0..raster.num_bands() { + let band = raster.band(band_idx)?; + total = total + .checked_add(estimated_band_bytes(band.as_ref())?) + .ok_or_else(|| RasterError::Invalid("raster byte count overflows u64".to_string()))?; + } + Ok(total) +} + +/// [`estimated_raster_bytes`] for every row of a raster column; null rows +/// estimate to 0. +/// +/// Reads the flattened band shape and data-type columns directly rather +/// than going through `RasterStructArray::get` and a boxed `BandRef` per +/// band, so a batch of 8192 rows costs a few hundred microseconds rather +/// than milliseconds — this runs once per input batch on the hot path of +/// byte-bounded batching. +pub fn estimated_row_bytes(rasters: &RasterStructArray<'_>) -> Result, RasterError> { + (0..rasters.len()) + .map(|idx| { + if rasters.is_null(idx) { + return Ok(0); + } + let mut total = 0u64; + for band_row in rasters.band_rows(idx) { + let bytes = band_bytes( + rasters.band_source_shape_at(band_row), + rasters.band_data_type_at(band_row)?, + )?; + total = total.checked_add(bytes).ok_or_else(|| { + RasterError::Invalid("raster byte count overflows u64".to_string()) + })?; + } + Ok(total) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::builder::{RasterBuilder, StartBandArgs}; + use crate::view_entries::{ViewEntries, ViewEntry}; + use arrow_array::StructArray; + + const TRANSFORM: [f64; 6] = [0.0, 1.0, 0.0, 0.0, 0.0, -1.0]; + + enum Storage { + InDb(Vec), + OutDb, + } + + /// One-row raster with a single `[y, x]` band of `source_shape` and + /// `data_type`, optionally viewed. + fn one_band( + spatial_shape: &[i64], + source_shape: &[i64], + data_type: BandDataType, + view: Option<&ViewEntries>, + storage: Storage, + ) -> StructArray { + let mut b = RasterBuilder::new(1); + b.start_raster_nd(&TRANSFORM, &["y", "x"], spatial_shape, None) + .unwrap(); + let outdb = matches!(storage, Storage::OutDb); + b.start_band(StartBandArgs { + view, + outdb_uri: outdb.then_some("mock://band"), + outdb_format: outdb.then_some("mock"), + ..StartBandArgs::new(&["y", "x"], source_shape, data_type) + }) + .unwrap(); + match storage { + Storage::InDb(bytes) => b.band_data_writer().append_value(bytes), + Storage::OutDb => b.band_data_writer().append_value([0u8; 0]), + } + b.finish_band().unwrap(); + b.finish_raster().unwrap(); + b.finish().unwrap() + } + + fn row_bytes(array: &StructArray) -> u64 { + let rasters = RasterStructArray::try_new(array).unwrap(); + estimated_raster_bytes(&rasters.get(0).unwrap()).unwrap() + } + + #[test] + fn indb_estimate_equals_buffer_length() { + let pixels: Vec = (0..32).collect(); + let array = one_band( + &[4, 8], + &[4, 8], + BandDataType::UInt8, + None, + Storage::InDb(pixels), + ); + let rasters = RasterStructArray::try_new(&array).unwrap(); + let raster = rasters.get(0).unwrap(); + let band = raster.band(0).unwrap(); + assert_eq!( + estimated_band_bytes(band.as_ref()).unwrap(), + band.nd_buffer().unwrap().buffer.len() as u64 + ); + assert_eq!(row_bytes(&array), 32); + } + + #[test] + fn outdb_estimate_comes_from_metadata_without_loading() { + // 512 × 512 float32 — the spatialbench tile — is 1 MiB once loaded. + let array = one_band( + &[512, 512], + &[512, 512], + BandDataType::Float32, + None, + Storage::OutDb, + ); + assert_eq!(row_bytes(&array), 1024 * 1024); + } + + #[test] + fn estimate_uses_raw_source_shape_not_visible_shape() { + // Source [4, 8] viewed down to row 1 (visible [1, 8]): the buffer is + // still 32 bytes, so the estimate must not shrink to 8. + let view = ViewEntries::new(vec![ + ViewEntry { + source_axis: 0, + start: 1, + step: 1, + steps: 1, + }, + ViewEntry { + source_axis: 1, + start: 0, + step: 1, + steps: 8, + }, + ]); + let array = one_band( + &[1, 8], + &[4, 8], + BandDataType::UInt8, + Some(&view), + Storage::OutDb, + ); + let rasters = RasterStructArray::try_new(&array).unwrap(); + let raster = rasters.get(0).unwrap(); + assert_eq!(raster.band(0).unwrap().shape(), &[1, 8]); + assert_eq!(row_bytes(&array), 32); + } + + #[test] + fn multi_band_and_null_rows() { + let mut b = RasterBuilder::new(3); + // Row 0: two bands, 2×3 UInt8 + 2×3 Int16 = 6 + 12. + b.start_raster_nd(&TRANSFORM, &["y", "x"], &[2, 3], None) + .unwrap(); + for dt in [BandDataType::UInt8, BandDataType::Int16] { + b.start_band(StartBandArgs { + outdb_uri: Some("mock://band"), + outdb_format: Some("mock"), + ..StartBandArgs::new(&["y", "x"], &[2, 3], dt) + }) + .unwrap(); + b.band_data_writer().append_value([0u8; 0]); + b.finish_band().unwrap(); + } + b.finish_raster().unwrap(); + // Row 1: null. + b.append_null().unwrap(); + // Row 2: no bands. + b.start_raster_nd(&TRANSFORM, &["y", "x"], &[2, 3], None) + .unwrap(); + b.finish_raster().unwrap(); + let array = b.finish().unwrap(); + + let rasters = RasterStructArray::try_new(&array).unwrap(); + assert_eq!(estimated_row_bytes(&rasters).unwrap(), vec![18, 0, 0]); + } + + #[test] + fn flat_row_estimates_match_the_per_row_trait_path() { + // Mixed bands, a null row, an empty raster and a viewed band: the + // flattened fast path must agree with `estimated_raster_bytes` row + // by row. + let mut b = RasterBuilder::new(4); + b.start_raster_nd(&TRANSFORM, &["y", "x"], &[2, 3], None) + .unwrap(); + for dt in [BandDataType::UInt8, BandDataType::Float64] { + b.start_band(StartBandArgs { + outdb_uri: Some("mock://band"), + outdb_format: Some("mock"), + ..StartBandArgs::new(&["y", "x"], &[2, 3], dt) + }) + .unwrap(); + b.band_data_writer().append_value([0u8; 0]); + b.finish_band().unwrap(); + } + b.finish_raster().unwrap(); + b.append_null().unwrap(); + b.start_raster_nd(&TRANSFORM, &["y", "x"], &[2, 3], None) + .unwrap(); + b.finish_raster().unwrap(); + let view = ViewEntries::new(vec![ + ViewEntry { + source_axis: 0, + start: 1, + step: 1, + steps: 1, + }, + ViewEntry { + source_axis: 1, + start: 0, + step: 1, + steps: 8, + }, + ]); + b.start_raster_nd(&TRANSFORM, &["y", "x"], &[1, 8], None) + .unwrap(); + b.start_band(StartBandArgs { + view: Some(&view), + ..StartBandArgs::new(&["y", "x"], &[4, 8], BandDataType::Int16) + }) + .unwrap(); + b.band_data_writer().append_value(vec![0u8; 64]); + b.finish_band().unwrap(); + b.finish_raster().unwrap(); + let array = b.finish().unwrap(); + + let rasters = RasterStructArray::try_new(&array).unwrap(); + let fast = estimated_row_bytes(&rasters).unwrap(); + let slow: Vec = (0..rasters.len()) + .map(|i| { + if rasters.is_null(i) { + 0 + } else { + estimated_raster_bytes(&rasters.get(i).unwrap()).unwrap() + } + }) + .collect(); + assert_eq!(fast, slow); + assert_eq!(fast, vec![6 + 48, 0, 0, 64]); + } + + #[test] + fn estimate_rejects_overflow_instead_of_wrapping() { + // 2^31 × 2^31 float64 = 2^65 bytes. + let big = 1i64 << 31; + let array = one_band( + &[big, big], + &[big, big], + BandDataType::Float64, + None, + Storage::OutDb, + ); + let rasters = RasterStructArray::try_new(&array).unwrap(); + let err = estimated_raster_bytes(&rasters.get(0).unwrap()).unwrap_err(); + assert!(err.to_string().contains("overflows"), "{err}"); + } +} From d02171c8d8948c6c62d5e66c15d6d56d1a6dbdd3 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Wed, 9 Sep 2026 17:38:39 -0700 Subject: [PATCH 2/3] feat(rust/sedona-common): add sedona.raster.max_batch_bytes New `raster` namespace on `SedonaOptions` with `max_batch_bytes` (default 256 MiB): the byte budget RS_EnsureLoaded materializes per batch, sized from the metadata-only raster estimate. `0` disables the slicing. Consumed by the EnsureLoadedExec in the following commit. Refs apache/sedona-db#1220 --- rust/sedona-common/src/option.rs | 45 ++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/rust/sedona-common/src/option.rs b/rust/sedona-common/src/option.rs index b566e2323c..65716a1099 100644 --- a/rust/sedona-common/src/option.rs +++ b/rust/sedona-common/src/option.rs @@ -48,6 +48,51 @@ config_namespace! { /// Options for configuring GDAL usage pub gdal: GdalOptions, default = GdalOptions::default() + + /// Options for raster execution + pub raster: RasterOptions, default = RasterOptions::default() + } +} + +/// Default byte budget for a batch of materialized rasters when no memory +/// limit is configured. `SedonaContext` lowers it to a fraction of the +/// per-partition memory limit when one is. +pub const DEFAULT_RASTER_MAX_BATCH_BYTES: usize = 256 * 1024 * 1024; + +config_namespace! { + /// Configuration options for raster execution. + pub struct RasterOptions { + /// Byte budget for one batch of materialized rasters, **per partition**. + /// `RS_EnsureLoaded` slices each input batch so that the estimated bytes + /// of the rasters it materializes (sum over bands of source shape × pixel + /// size, taken from band metadata) stay within this budget, emitting each + /// slice as its own batch instead of trusting the row-count + /// `datafusion.execution.batch_size`. A single row larger than the budget + /// is still processed, on its own. `0` disables the slicing. + /// + /// It is not a global cap: every partition streams its own batches, and a + /// pipeline typically holds two or three of them at once (the one being + /// built, the one the consumer is working on, and whatever a downstream + /// operator buffers), so the loaded pixels in flight across a query are + /// roughly `target_partitions × 3 × max_batch_bytes`. + /// + /// The default is 256 MiB. When a memory limit is configured, + /// `SedonaContext` derives a lower default at session construction: + /// `min(256 MiB, max(16 MiB, (memory_limit / target_partitions) / 8))`, + /// i.e. one eighth of each partition's share of the limit, floored at + /// 16 MiB. For example: + /// + /// | memory limit | target_partitions | per-partition share | budget | + /// |--------------|-------------------|---------------------|---------| + /// | none | any | — | 256 MiB | + /// | 16 GiB | 8 | 2 GiB | 256 MiB | + /// | 16 GiB | 16 | 1 GiB | 128 MiB | + /// | 16 GiB | 64 | 256 MiB | 32 MiB | + /// | 2 GiB | 16 | 128 MiB | 16 MiB | + /// + /// A `SET sedona.raster.max_batch_bytes = …` after connecting overrides + /// whichever default was derived. + pub max_batch_bytes: usize, default = DEFAULT_RASTER_MAX_BATCH_BYTES } } From 261debe32427fdaa13485706b09bb31488956168 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Wed, 9 Sep 2026 17:38:39 -0700 Subject: [PATCH 3/3] feat(rust/sedona-query-planner): re-batch RS_EnsureLoaded by estimated raster bytes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `datafusion.execution.batch_size` counts rows, but a row holding a raster can carry megabytes of pixels, so the batch DataFusion hands RS_EnsureLoaded can be gigabytes once loaded (8192 × 1 MiB spatialbench tiles = 8 GiB per partition). Nothing upstream can shrink it: DataFusion's AsyncFuncExec re-coalesces its input to exactly `batch_size` rows before evaluating, and the UDF's `ideal_batch_size` only chunks the invocation — the results are concatenated back into one array. Add `RasterBatchBudgetRule`, a physical optimizer rule appended after DataFusion's own, that replaces every AsyncFuncExec carrying an `rs_ensureloaded` call with `EnsureLoadedExec`: the same expressions, schema and plan properties (including the `__async_fn_N` output columns), but each input batch is sliced so the metadata-only byte estimate of the rasters about to be materialized stays within `sedona.raster.max_batch_bytes`, and each slice is evaluated and emitted as its own batch. A row over budget goes alone, null rows cost nothing, empty batches flow through, and other async expressions in the same node are evaluated per slice. `SedonaContext` installs the rule alongside the existing planner rules and, when a memory limit is configured, lowers the budget to 1/8 of the per-partition limit (floor 16 MiB), mirroring the spill threshold. Streaming 2048 one-MiB OutDb rasters through `SELECT RS_EnsureLoaded(rast)` (mock loader, single partition, debug build) peaks at 2.20 GB RSS unbounded versus 779 MB at the 256 MiB default, 591 MB at 64 MiB and 145 MB at 16 MiB. Refs apache/sedona-db#1220 --- Cargo.lock | 4 + rust/sedona-query-planner/Cargo.toml | 8 + .../benches/raster_batch_budget.rs | 210 +++++ rust/sedona-query-planner/src/lib.rs | 1 + .../src/raster_batch_budget.rs | 746 ++++++++++++++++++ rust/sedona/src/context.rs | 135 ++++ 6 files changed, 1104 insertions(+) create mode 100644 rust/sedona-query-planner/benches/raster_batch_budget.rs create mode 100644 rust/sedona-query-planner/src/raster_batch_budget.rs diff --git a/Cargo.lock b/Cargo.lock index 969745fc75..1cdfb88eb0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6074,8 +6074,10 @@ name = "sedona-query-planner" version = "0.5.0" dependencies = [ "arrow", + "arrow-array", "arrow-schema", "async-trait", + "criterion", "datafusion", "datafusion-common", "datafusion-execution", @@ -6083,8 +6085,10 @@ dependencies = [ "datafusion-optimizer", "datafusion-physical-expr", "datafusion-physical-plan", + "futures", "sedona-common", "sedona-expr", + "sedona-raster", "sedona-schema", "tokio", ] diff --git a/rust/sedona-query-planner/Cargo.toml b/rust/sedona-query-planner/Cargo.toml index eee2e66dfa..7613246c33 100644 --- a/rust/sedona-query-planner/Cargo.toml +++ b/rust/sedona-query-planner/Cargo.toml @@ -27,6 +27,7 @@ edition = "2024" rust-version.workspace = true [dependencies] +arrow-array = { workspace = true } arrow-schema = { workspace = true } async-trait = { workspace = true } datafusion = { workspace = true } @@ -36,12 +37,19 @@ datafusion-expr = { workspace = true } datafusion-optimizer = { workspace = true } datafusion-physical-expr = { workspace = true } datafusion-physical-plan = { workspace = true } +futures = { workspace = true } sedona-common = { workspace = true } sedona-expr = { workspace = true } +sedona-raster = { workspace = true } sedona-schema = { workspace = true } [dev-dependencies] arrow = { workspace = true } +criterion = { workspace = true } datafusion = { workspace = true, features = ["nested_expressions", "sql"] } sedona-schema = { workspace = true } tokio = { workspace = true, features = ["macros", "rt"] } + +[[bench]] +name = "raster_batch_budget" +harness = false diff --git a/rust/sedona-query-planner/benches/raster_batch_budget.rs b/rust/sedona-query-planner/benches/raster_batch_budget.rs new file mode 100644 index 0000000000..05dd088299 --- /dev/null +++ b/rust/sedona-query-planner/benches/raster_batch_budget.rs @@ -0,0 +1,210 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Overhead of `EnsureLoadedExec` in the non-materialized case. +//! +//! The async UDF is an identity that stands in for `rs_ensureloaded` over +//! rasters that are already in memory, so the only work measured is the +//! operator's own: the per-row byte estimate, the slicing, and one async +//! invocation plus one output batch per slice. DataFusion's stock +//! `AsyncFuncExec` over the same plan is the baseline, and the budget is +//! swept so the 8192-row input batch is emitted as 1, 8, 64, 512 and 8192 +//! slices. + +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow_array::{ArrayRef, RecordBatch}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use async_trait::async_trait; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::SessionStateBuilder; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::prelude::SessionConfig; +use datafusion_common::Result; +use datafusion_common::config::ConfigOptions; +use datafusion_execution::TaskContext; +use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; +use datafusion_physical_expr::ScalarFunctionExpr; +use datafusion_physical_expr::async_scalar_function::AsyncFuncExpr; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::async_func::AsyncFuncExec; +use datafusion_physical_plan::common::collect; +use sedona_common::option::{RasterOptions, SedonaOptions}; +use sedona_query_planner::raster_batch_budget::RasterBatchBudgetRule; +use sedona_raster::builder::{RasterBuilder, StartBandArgs}; +use sedona_schema::datatypes::SedonaType; +use sedona_schema::raster::BandDataType; + +const ROWS: usize = 8192; +/// 16 × 16 UInt8: 256 bytes per raster, so a budget of `256 × rows_per_slice` +/// bytes yields slices of exactly that many rows. +const SIDE: i64 = 16; +const BYTES_PER_ROW: usize = (SIDE * SIDE) as usize; + +/// Identity stand-in for `rs_ensureloaded` (same name, so the rule picks it +/// up), costing nothing beyond the argument array clone. +#[derive(Debug)] +struct IdentityEnsureLoaded { + signature: Signature, +} + +impl PartialEq for IdentityEnsureLoaded { + fn eq(&self, _other: &Self) -> bool { + true + } +} +impl Eq for IdentityEnsureLoaded {} +impl Hash for IdentityEnsureLoaded { + fn hash(&self, state: &mut H) { + "rs_ensureloaded".hash(state); + } +} + +impl ScalarUDFImpl for IdentityEnsureLoaded { + fn name(&self) -> &str { + "rs_ensureloaded" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) + } + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + unreachable!("async only") + } +} + +#[async_trait] +impl AsyncScalarUDFImpl for IdentityEnsureLoaded { + async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result { + Ok(ColumnarValue::Array( + args.args[0].to_array(args.number_rows)?, + )) + } +} + +/// One batch of `ROWS` in-db 16 × 16 UInt8 rasters. +fn raster_batch() -> (SchemaRef, RecordBatch) { + let pixels = vec![7u8; BYTES_PER_ROW]; + let mut b = RasterBuilder::new(ROWS); + for _ in 0..ROWS { + b.start_raster_nd( + &[0.0, 1.0, 0.0, 0.0, 0.0, -1.0], + &["y", "x"], + &[SIDE, SIDE], + None, + ) + .unwrap(); + b.start_band(StartBandArgs::new( + &["y", "x"], + &[SIDE, SIDE], + BandDataType::UInt8, + )) + .unwrap(); + b.band_data_writer().append_value(&pixels); + b.finish_band().unwrap(); + b.finish_raster().unwrap(); + } + let rasters: ArrayRef = Arc::new(b.finish().unwrap()); + let raster_field = SedonaType::Raster.to_storage_field("rast", true).unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("rast", rasters.data_type().clone(), true) + .with_metadata(raster_field.metadata().clone()), + ])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![rasters]).unwrap(); + (schema, batch) +} + +/// `AsyncFuncExec(rs_ensureloaded(rast))` over the batch. +fn async_func_plan(schema: &SchemaRef, batch: &RecordBatch) -> Arc { + let input = + MemorySourceConfig::try_new_exec(&[vec![batch.clone()]], Arc::clone(schema), None).unwrap(); + let udf = Arc::new( + AsyncScalarUDF::new(Arc::new(IdentityEnsureLoaded { + signature: Signature::any(1, Volatility::Stable), + })) + .into_scalar_udf(), + ); + let func = Arc::new( + ScalarFunctionExpr::try_new( + udf, + vec![col("rast", schema).unwrap()], + schema, + Arc::new(ConfigOptions::default()), + ) + .unwrap(), + ); + let expr = Arc::new(AsyncFuncExpr::try_new("__async_fn_0", func, schema).unwrap()); + Arc::new(AsyncFuncExec::try_new(vec![expr], input).unwrap()) +} + +fn task_context(max_batch_bytes: usize) -> Arc { + let config = SessionConfig::new().with_option_extension(SedonaOptions { + raster: RasterOptions { max_batch_bytes }, + ..Default::default() + }); + SessionStateBuilder::new() + .with_config(config) + .build() + .task_ctx() +} + +fn bench(c: &mut Criterion) { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let (schema, batch) = raster_batch(); + let stock = async_func_plan(&schema, &batch); + let budgeted = RasterBatchBudgetRule + .optimize(async_func_plan(&schema, &batch), &ConfigOptions::default()) + .unwrap(); + let ctx = task_context(0); + + let mut group = c.benchmark_group("ensure_loaded_exec_overhead"); + group.throughput(Throughput::Elements(ROWS as u64)); + + group.bench_function("AsyncFuncExec (baseline)", |b| { + b.iter(|| { + let stream = stock.execute(0, Arc::clone(&ctx)).unwrap(); + rt.block_on(collect(stream)).unwrap() + }) + }); + + for rows_per_slice in [ROWS, 1024, 128, 16, 1] { + let ctx = task_context(BYTES_PER_ROW * rows_per_slice); + let slices = ROWS / rows_per_slice; + group.bench_with_input( + BenchmarkId::new("EnsureLoadedExec/slices", slices), + &slices, + |b, _| { + b.iter(|| { + let stream = budgeted.execute(0, Arc::clone(&ctx)).unwrap(); + rt.block_on(collect(stream)).unwrap() + }) + }, + ); + } + group.finish(); +} + +criterion_group!(benches, bench); +criterion_main!(benches); diff --git a/rust/sedona-query-planner/src/lib.rs b/rust/sedona-query-planner/src/lib.rs index 40b30df706..c9a7d5ce20 100644 --- a/rust/sedona-query-planner/src/lib.rs +++ b/rust/sedona-query-planner/src/lib.rs @@ -25,6 +25,7 @@ pub mod probe_shuffle_exec; // https://github.com/apache/sedona-db/issues/1232. pub mod push_down_leaf_projections; pub mod query_planner; +pub mod raster_batch_budget; mod restore_metadata; mod spatial_expr_utils; pub mod spatial_join_physical_planner; diff --git a/rust/sedona-query-planner/src/raster_batch_budget.rs b/rust/sedona-query-planner/src/raster_batch_budget.rs new file mode 100644 index 0000000000..ec21b15b17 --- /dev/null +++ b/rust/sedona-query-planner/src/raster_batch_budget.rs @@ -0,0 +1,746 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Byte-bounded batches for raster materialization. +//! +//! `datafusion.execution.batch_size` counts rows, but a row holding a +//! raster can carry megabytes of pixels, so the batch DataFusion hands +//! `RS_EnsureLoaded` can be gigabytes once loaded. DataFusion runs async +//! UDFs in an `AsyncFuncExec` that re-coalesces its input to exactly +//! `batch_size` rows before evaluating, so nothing upstream can shrink that +//! batch, and the UDF's `ideal_batch_size` only chunks the *invocation* — +//! the results are concatenated back into one array. +//! +//! [`RasterBatchBudgetRule`] therefore replaces every `AsyncFuncExec` that +//! carries an `rs_ensureloaded` call with an [`EnsureLoadedExec`]: same +//! expressions, same output schema, but each input batch is sliced so that +//! the *estimated* bytes of the rasters about to be materialized stay +//! within `sedona.raster.max_batch_bytes`, and each slice is evaluated and +//! emitted as its own batch. The estimate is metadata-only +//! ([`sedona_raster::size`]), so it is available before any loading +//! happens and is identical for InDb and OutDb bands. A single row larger +//! than the budget still goes through, alone. +//! +//! Slices are contiguous row ranges of the input batch (`RecordBatch::slice` +//! is zero-copy), so output order is the input order. +//! +//! This bounds the bytes *created* at the materialization point. Operators +//! that *merge* batches downstream — `FilterExec`'s coalescer, +//! `CoalesceBatchesExec`, `RepartitionExec` — still re-merge by row count +//! and can rebuild an oversized batch from these slices; keeping them from +//! doing so on raster-bearing streams is follow-up work here. Longer term +//! that half belongs upstream: if arrow's `BatchCoalescer` (and DataFusion's +//! `LimitedBatchCoalescer` on top of it) accepted a byte target next to the +//! row target, every merger would become byte-aware on its own and only the +//! materialization points would need Sedona-owned operators. See +//! , which tracks the +//! coalescers' lack of memory awareness. + +use std::fmt; +use std::sync::Arc; + +use arrow_array::{Array, RecordBatch, StructArray}; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_physical_expr::async_scalar_function::AsyncFuncExpr; +use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr}; +use datafusion_physical_plan::async_func::AsyncFuncExec; +use datafusion_physical_plan::execution_plan::CardinalityEffect; +use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use datafusion_physical_plan::stream::RecordBatchStreamAdapter; +use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; +use futures::{StreamExt, stream}; +use sedona_common::option::{DEFAULT_RASTER_MAX_BATCH_BYTES, SedonaOptions}; +use sedona_common::sedona_internal_datafusion_err; +use sedona_raster::array::RasterStructArray; +use sedona_raster::size::estimated_row_bytes; + +/// Name of the async UDF whose raster argument sizes the slices. Kept in +/// sync with `sedona_raster_functions::rs_ensure_loaded` (this crate can't +/// depend on it); `ensure_loaded.rs` resolves the same name. +const ENSURE_LOADED_NAME: &str = "rs_ensureloaded"; + +/// Physical optimizer rule that swaps DataFusion's `AsyncFuncExec` for an +/// [`EnsureLoadedExec`] wherever the async expressions include +/// `rs_ensureloaded`. Runs after DataFusion's own physical rules (it is +/// appended via `SessionStateBuilder::with_physical_optimizer_rule`), so it +/// sees the final placement of the async node. Plans without raster +/// materialization are untouched. +#[derive(Debug, Default)] +pub struct RasterBatchBudgetRule; + +impl PhysicalOptimizerRule for RasterBatchBudgetRule { + fn optimize( + &self, + plan: Arc, + _config: &ConfigOptions, + ) -> Result> { + plan.transform_up(|node| { + let Some(async_exec) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + if !async_exec.async_exprs().iter().any(|e| is_ensure_loaded(e)) { + return Ok(Transformed::no(node)); + } + Ok(Transformed::yes(Arc::new( + EnsureLoadedExec::from_async_func_exec(async_exec)?, + ))) + }) + .data() + } + + fn name(&self) -> &str { + "sedona.raster_batch_budget" + } + + fn schema_check(&self) -> bool { + true + } +} + +fn scalar_function(expr: &AsyncFuncExpr) -> Option<&ScalarFunctionExpr> { + expr.func.downcast_ref::() +} + +fn is_ensure_loaded(expr: &AsyncFuncExpr) -> bool { + scalar_function(expr).is_some_and(|f| f.fun().name() == ENSURE_LOADED_NAME) +} + +/// The raster argument of an `rs_ensureloaded` call. +fn raster_arg(expr: &AsyncFuncExpr) -> Result<&Arc> { + scalar_function(expr) + .and_then(|f| f.args().first()) + .ok_or_else(|| { + sedona_internal_datafusion_err!( + "EnsureLoadedExec: {} has no raster argument to size by", + expr.name() + ) + }) +} + +/// Contiguous `[start, end)` row ranges whose summed `estimates` stay within +/// `budget`. A row that alone exceeds the budget gets a range of its own. An +/// empty input yields one empty range so an empty batch still flows through. +fn slice_ranges(estimates: &[u64], budget: u64) -> Vec<(usize, usize)> { + if estimates.is_empty() { + return vec![(0, 0)]; + } + let mut ranges = Vec::new(); + let mut start = 0; + let mut acc = 0u64; + for (idx, &bytes) in estimates.iter().enumerate() { + if idx > start && acc.saturating_add(bytes) > budget { + ranges.push((start, idx)); + start = idx; + acc = 0; + } + acc = acc.saturating_add(bytes); + } + ranges.push((start, estimates.len())); + ranges +} + +/// Byte budget for one slice, from the session's +/// `sedona.raster.max_batch_bytes`. It applies **per batch, per partition**: +/// each partition's stream is sliced independently, so the loaded pixels in +/// flight across a query are roughly `target_partitions × (batches a +/// pipeline holds, typically 2–3) × budget`. +/// +/// The rule never consults the memory limit itself. The value read here is +/// whatever the session holds: the 256 MiB crate default, the lower default +/// `SedonaContext` derives at construction when a memory limit is configured +/// (`(memory_limit / target_partitions) / 8`, floored at 16 MiB, capped at +/// 256 MiB — worked examples on the option's docs in `sedona-common`), or an +/// explicit `SET`. `0` disables slicing. A session without the extension (a +/// bare DataFusion context) gets the crate default. +fn budget_bytes(context: &TaskContext) -> u64 { + let bytes = context + .session_config() + .options() + .extensions + .get::() + .map(|opts| opts.raster.max_batch_bytes) + .unwrap_or(DEFAULT_RASTER_MAX_BATCH_BYTES); + if bytes == 0 { u64::MAX } else { bytes as u64 } +} + +/// DataFusion's `AsyncFuncExec`, re-batched by estimated raster bytes. See +/// the module docs. +#[derive(Debug)] +pub struct EnsureLoadedExec { + async_exprs: Vec>, + input: Arc, + /// Indices into `async_exprs` of the `rs_ensureloaded` calls whose + /// raster arguments size the slices (summed per row when several). + sized_by: Vec, + cache: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl EnsureLoadedExec { + /// Build from an `AsyncFuncExec`, inheriting its schema and plan + /// properties verbatim — including the `__async_fn_N` output column + /// names DataFusion's planner projects on top of this node. + pub fn from_async_func_exec(exec: &AsyncFuncExec) -> Result { + let async_exprs = exec.async_exprs().to_vec(); + let sized_by: Vec = async_exprs + .iter() + .enumerate() + .filter(|(_, e)| is_ensure_loaded(e)) + .map(|(idx, _)| idx) + .collect(); + if sized_by.is_empty() { + return internal_err!( + "EnsureLoadedExec requires at least one {ENSURE_LOADED_NAME} expression" + ); + } + for &idx in &sized_by { + raster_arg(&async_exprs[idx])?; + } + Ok(Self { + async_exprs, + input: Arc::clone(exec.input()), + sized_by, + cache: Arc::clone(exec.properties()), + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + pub fn try_new( + async_exprs: Vec>, + input: Arc, + ) -> Result { + Self::from_async_func_exec(&AsyncFuncExec::try_new(async_exprs, input)?) + } + + pub fn async_exprs(&self) -> &[Arc] { + &self.async_exprs + } + + pub fn input(&self) -> &Arc { + &self.input + } + + /// Per-row estimated bytes of everything `rs_ensureloaded` will + /// materialize from `batch`: the sum over the sized expressions of the + /// metadata-only estimate of their raster argument. + fn estimate_rows( + async_exprs: &[Arc], + sized_by: &[usize], + batch: &RecordBatch, + ) -> Result> { + let mut totals = vec![0u64; batch.num_rows()]; + for &idx in sized_by { + let array = raster_arg(&async_exprs[idx])? + .evaluate(batch)? + .into_array(batch.num_rows())?; + let struct_array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + sedona_internal_datafusion_err!( + "EnsureLoadedExec: expected a Raster (Struct) argument, got {:?}", + array.data_type() + ) + })?; + let rasters = RasterStructArray::try_new(struct_array)?; + for (total, bytes) in totals.iter_mut().zip(estimated_row_bytes(&rasters)?) { + *total = total.saturating_add(bytes); + } + } + Ok(totals) + } +} + +impl DisplayAs for EnsureLoadedExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + let exprs = self + .async_exprs + .iter() + .map(|e| e.to_string()) + .collect::>() + .join(", "); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "EnsureLoadedExec: async_expr=[{exprs}]") + } + DisplayFormatType::TreeRender => { + writeln!(f, "format=async_expr")?; + writeln!(f, "async_expr={exprs}") + } + } + } +} + +impl ExecutionPlan for EnsureLoadedExec { + fn name(&self) -> &str { + "EnsureLoadedExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!( + "EnsureLoadedExec expects exactly 1 child, got {}", + children.len() + ); + } + Ok(Arc::new(Self::try_new( + self.async_exprs.clone(), + children.remove(0), + )?)) + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + // `benefits_from_input_partitioning` is deliberately left at the default + // (true), matching DataFusion's `AsyncFuncExec`: the planner then puts any + // round-robin repartition *below* this node, over cheap OutDb references, + // so loads run in parallel and no repartition coalescer sits above the + // byte-bounded output re-merging it. + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let input = self.input.execute(partition, Arc::clone(&context))?; + let budget = budget_bytes(&context); + let config_options = Arc::clone(context.session_config().options()); + let async_exprs = Arc::new(self.async_exprs.clone()); + let sized_by = Arc::new(self.sized_by.clone()); + let schema = self.schema(); + let baseline = BaselineMetrics::new(&self.metrics, partition); + + // One input batch fans out into as many output batches as the + // budget requires; each slice is evaluated in turn so at most one + // slice's worth of loaded pixels is in flight per partition. + let output = input.flat_map(move |batch| { + let batch = match batch { + Ok(batch) => batch, + Err(e) => return stream::once(async move { Err(e) }).boxed(), + }; + let ranges = match Self::estimate_rows(&async_exprs, &sized_by, &batch) { + Ok(estimates) => slice_ranges(&estimates, budget), + Err(e) => return stream::once(async move { Err(e) }).boxed(), + }; + let slices: Vec = ranges + .into_iter() + .map(|(start, end)| batch.slice(start, end - start)) + .collect(); + + let async_exprs = Arc::clone(&async_exprs); + let schema = Arc::clone(&schema); + let config_options = Arc::clone(&config_options); + let baseline = baseline.clone(); + stream::iter(slices) + .then(move |slice| { + let async_exprs = Arc::clone(&async_exprs); + let schema = Arc::clone(&schema); + let config_options = Arc::clone(&config_options); + let baseline = baseline.clone(); + async move { + let mut columns = slice.columns().to_vec(); + for expr in async_exprs.iter() { + let value = expr + .invoke_with_args(&slice, Arc::clone(&config_options)) + .await?; + columns.push(value.to_array(slice.num_rows())?); + } + let out = RecordBatch::try_new(schema, columns)?; + baseline.record_output(out.num_rows()); + Ok(out) + } + }) + .boxed() + }); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + output, + ))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Mutex; + + use arrow_array::{ArrayRef, Int32Array}; + use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use async_trait::async_trait; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::execution::SessionStateBuilder; + use datafusion::prelude::SessionConfig; + use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; + use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, + }; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_plan::common::collect; + use sedona_common::option::RasterOptions; + use sedona_raster::builder::{RasterBuilder, StartBandArgs}; + use sedona_schema::raster::BandDataType; + + /// Stand-in for `RS_EnsureLoaded`: same name, identity on its argument, + /// records the row count of every invocation. + #[derive(Debug)] + struct MockEnsureLoaded { + signature: Signature, + calls: Arc>>, + } + + impl MockEnsureLoaded { + fn new(calls: Arc>>) -> Self { + Self { + signature: Signature::any(1, Volatility::Stable), + calls, + } + } + } + + // DataFusion dedups UDFs by equality/hash; identity by name is enough here. + impl PartialEq for MockEnsureLoaded { + fn eq(&self, _other: &Self) -> bool { + true + } + } + impl Eq for MockEnsureLoaded {} + impl std::hash::Hash for MockEnsureLoaded { + fn hash(&self, state: &mut H) { + ENSURE_LOADED_NAME.hash(state); + } + } + + impl ScalarUDFImpl for MockEnsureLoaded { + fn name(&self) -> &str { + ENSURE_LOADED_NAME + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) + } + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + internal_err!("async only") + } + } + + #[async_trait] + impl AsyncScalarUDFImpl for MockEnsureLoaded { + async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result { + self.calls.lock().unwrap().push(args.number_rows); + Ok(ColumnarValue::Array( + args.args[0].to_array(args.number_rows)?, + )) + } + } + + /// Async UDF that is *not* rs_ensureloaded; the rule must leave it alone. + #[derive(Debug)] + struct OtherAsync { + signature: Signature, + } + + impl PartialEq for OtherAsync { + fn eq(&self, _other: &Self) -> bool { + true + } + } + impl Eq for OtherAsync {} + impl std::hash::Hash for OtherAsync { + fn hash(&self, state: &mut H) { + "other_async".hash(state); + } + } + + impl ScalarUDFImpl for OtherAsync { + fn name(&self) -> &str { + "other_async" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int32) + } + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + internal_err!("async only") + } + } + + #[async_trait] + impl AsyncScalarUDFImpl for OtherAsync { + async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result { + // One value per row: DataFusion expands a scalar async result to a + // length-1 array, so async UDFs must return arrays. + let rows = args.number_rows; + Ok(ColumnarValue::Array(Arc::new(Int32Array::from(vec![ + rows as i32; + rows + ])))) + } + } + + /// One OutDb `[side, side]` UInt8 band per row, so row `i` estimates to + /// `sides[i]²` bytes without anything to load; `None` is a null row. + fn raster_batch(sides: &[Option]) -> (SchemaRef, RecordBatch) { + let mut b = RasterBuilder::new(sides.len()); + for side in sides { + let Some(side) = side else { + b.append_null().unwrap(); + continue; + }; + b.start_raster_nd( + &[0.0, 1.0, 0.0, 0.0, 0.0, -1.0], + &["y", "x"], + &[*side, *side], + None, + ) + .unwrap(); + b.start_band(StartBandArgs { + outdb_uri: Some("mock://tile"), + outdb_format: Some("mock"), + ..StartBandArgs::new(&["y", "x"], &[*side, *side], BandDataType::UInt8) + }) + .unwrap(); + b.band_data_writer().append_value([0u8; 0]); + b.finish_band().unwrap(); + b.finish_raster().unwrap(); + } + let rasters: ArrayRef = Arc::new(b.finish().unwrap()); + let schema = Arc::new(Schema::new(vec![Field::new( + "rast", + rasters.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![rasters]).unwrap(); + (schema, batch) + } + + fn async_expr( + udf: Arc, + name: &str, + schema: &Schema, + ) -> Arc { + let udf = Arc::new(AsyncScalarUDF::new(udf).into_scalar_udf()); + let func = Arc::new( + ScalarFunctionExpr::try_new( + udf, + vec![col("rast", schema).unwrap()], + schema, + Arc::new(ConfigOptions::default()), + ) + .unwrap(), + ); + Arc::new(AsyncFuncExpr::try_new(name, func, schema).unwrap()) + } + + fn task_context(max_batch_bytes: usize) -> Arc { + let config = SessionConfig::new().with_option_extension(SedonaOptions { + raster: RasterOptions { max_batch_bytes }, + ..Default::default() + }); + SessionStateBuilder::new() + .with_config(config) + .build() + .task_ctx() + } + + /// Build `AsyncFuncExec(rs_ensureloaded(rast))` over `batch`, run the + /// rule, execute with `max_batch_bytes`, and return the output batches + /// plus the row counts the mock saw. + async fn run( + sides: &[Option], + max_batch_bytes: usize, + ) -> (Vec, Vec, SchemaRef) { + let (schema, batch) = raster_batch(sides); + let input = + MemorySourceConfig::try_new_exec(&[vec![batch]], Arc::clone(&schema), None).unwrap(); + let calls = Arc::new(Mutex::new(Vec::new())); + let expr = async_expr( + Arc::new(MockEnsureLoaded::new(Arc::clone(&calls))), + "__async_fn_0", + &schema, + ); + let async_exec = Arc::new(AsyncFuncExec::try_new(vec![expr], input).unwrap()); + let expected_schema = async_exec.schema(); + + let plan = RasterBatchBudgetRule + .optimize(async_exec, &ConfigOptions::default()) + .unwrap(); + assert!( + plan.downcast_ref::().is_some(), + "rule must replace AsyncFuncExec, got {}", + plan.name() + ); + assert_eq!( + plan.schema(), + expected_schema, + "schema must match AsyncFuncExec's" + ); + + let stream = plan.execute(0, task_context(max_batch_bytes)).unwrap(); + let batches = collect(stream).await.unwrap(); + let calls = calls.lock().unwrap().clone(); + (batches, calls, expected_schema) + } + + fn row_counts(batches: &[RecordBatch]) -> Vec { + batches.iter().map(RecordBatch::num_rows).collect() + } + + #[test] + fn slice_ranges_packs_rows_up_to_the_budget() { + assert_eq!( + slice_ranges(&[4, 4, 4, 4, 4], 8), + vec![(0, 2), (2, 4), (4, 5)] + ); + assert_eq!(slice_ranges(&[1, 1, 1], 100), vec![(0, 3)]); + // A row over budget goes alone, without dragging its neighbours in. + assert_eq!( + slice_ranges(&[1, 50, 1, 1], 8), + vec![(0, 1), (1, 2), (2, 4)] + ); + assert_eq!(slice_ranges(&[50], 8), vec![(0, 1)]); + assert_eq!(slice_ranges(&[], 8), vec![(0, 0)]); + // Zero-byte rows (nulls) never force a split. + assert_eq!(slice_ranges(&[0, 0, 8, 0], 8), vec![(0, 4)]); + } + + #[tokio::test] + async fn uniform_rasters_are_sliced_to_the_budget() { + // Six 256-byte rasters under a 600-byte budget → 2 per slice. + let (batches, calls, _) = run(&[Some(16); 6], 600).await; + assert_eq!(row_counts(&batches), vec![2, 2, 2]); + assert_eq!(calls, vec![2, 2, 2]); + } + + #[tokio::test] + async fn oversized_rows_go_alone_and_nulls_cost_nothing() { + // 100, 1024, null, 100, 100 bytes under a 256-byte budget. + let (batches, calls, _) = run(&[Some(10), Some(32), None, Some(10), Some(10)], 256).await; + assert_eq!(row_counts(&batches), vec![1, 1, 3]); + assert_eq!(calls, vec![1, 1, 3]); + // Output rows are the input rows in order, and the null survives. + let out: Vec = batches.iter().map(|b| b.column(1).null_count()).collect(); + assert_eq!(out, vec![0, 0, 1]); + } + + #[tokio::test] + async fn zero_budget_disables_slicing() { + let (batches, calls, _) = run(&[Some(16); 6], 0).await; + assert_eq!(row_counts(&batches), vec![6]); + assert_eq!(calls, vec![6]); + } + + #[tokio::test] + async fn output_appends_the_async_column_after_the_input_columns() { + let (batches, _, schema) = run(&[Some(4), Some(4)], 1024).await; + assert_eq!(schema.fields().len(), 2); + assert_eq!(schema.field(1).name(), "__async_fn_0"); + let batch = &batches[0]; + assert_eq!(batch.schema(), schema); + // Identity mock: the appended column equals the raster input. + assert_eq!(batch.column(0).as_ref(), batch.column(1).as_ref()); + } + + #[tokio::test] + async fn empty_batch_flows_through() { + let (batches, calls, _) = run(&[], 1024).await; + assert_eq!(row_counts(&batches), vec![0]); + assert_eq!(calls, vec![0]); + } + + #[tokio::test] + async fn rule_ignores_async_funcs_without_ensure_loaded() { + let (schema, batch) = raster_batch(&[Some(4)]); + let input = + MemorySourceConfig::try_new_exec(&[vec![batch]], Arc::clone(&schema), None).unwrap(); + let expr = async_expr( + Arc::new(OtherAsync { + signature: Signature::any(1, Volatility::Stable), + }), + "__async_fn_0", + &schema, + ); + let async_exec: Arc = + Arc::new(AsyncFuncExec::try_new(vec![expr], input).unwrap()); + let plan = RasterBatchBudgetRule + .optimize(Arc::clone(&async_exec), &ConfigOptions::default()) + .unwrap(); + assert!(plan.downcast_ref::().is_some()); + } + + #[tokio::test] + async fn mixed_async_funcs_are_sized_by_the_ensure_loaded_argument_only() { + let (schema, batch) = raster_batch(&[Some(16); 4]); + let input = + MemorySourceConfig::try_new_exec(&[vec![batch]], Arc::clone(&schema), None).unwrap(); + let calls = Arc::new(Mutex::new(Vec::new())); + let exprs = vec![ + async_expr( + Arc::new(OtherAsync { + signature: Signature::any(1, Volatility::Stable), + }), + "__async_fn_0", + &schema, + ), + async_expr( + Arc::new(MockEnsureLoaded::new(Arc::clone(&calls))), + "__async_fn_1", + &schema, + ), + ]; + let async_exec = Arc::new(AsyncFuncExec::try_new(exprs, input).unwrap()); + let plan = RasterBatchBudgetRule + .optimize(async_exec, &ConfigOptions::default()) + .unwrap(); + let batches = collect(plan.execute(0, task_context(512)).unwrap()) + .await + .unwrap(); + assert_eq!(row_counts(&batches), vec![2, 2]); + assert_eq!(*calls.lock().unwrap(), vec![2, 2]); + // The other async column is evaluated per slice too (scalar → 2 rows). + assert_eq!(batches[0].column(1).len(), 2); + } +} diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs index 8e5efc9bce..791ca629a9 100644 --- a/rust/sedona/src/context.rs +++ b/rust/sedona/src/context.rs @@ -78,6 +78,7 @@ use sedona_query_planner::{ register_vendored_optimizer_rules, }, query_planner::SedonaQueryPlanner, + raster_batch_budget::RasterBatchBudgetRule, }; use sedona_raster::raster_loader::{AsyncRasterLoader, RasterLoaderConfig, RasterLoaderRegistry}; @@ -149,11 +150,24 @@ impl SedonaContext { // large spilled batches. const SPILLED_BATCH_THRESHOLD_PERCENT_DIVISOR: usize = 20; // 5% == 1 / 20 const MIN_SPILLED_BATCH_IN_MEMORY_THRESHOLD_BYTES: usize = 10 * 1024 * 1024; // 10MB + + // Likewise scale the byte budget for one batch of materialized rasters + // (`sedona.raster.max_batch_bytes`, applied per partition) to 1/8 of the + // per-partition limit, floored at 16MB and never above the unbounded + // default: 16 GiB on 16 partitions gives a 1 GiB share and a 128 MiB + // budget; 2 GiB on 16 partitions hits the 16 MiB floor. A later + // `SET sedona.raster.max_batch_bytes` overrides this. + const RASTER_BATCH_BUDGET_DIVISOR: usize = 8; + const MIN_RASTER_MAX_BATCH_BYTES: usize = 16 * 1024 * 1024; // 16MB if let MemoryLimit::Finite(memory_limit) = runtime_env.memory_pool.memory_limit() { let per_partition_memory_limit = memory_limit.div_ceil(target_partitions); opts.spatial_join.spilled_batch_in_memory_size_threshold = per_partition_memory_limit .div_ceil(SPILLED_BATCH_THRESHOLD_PERCENT_DIVISOR) .max(MIN_SPILLED_BATCH_IN_MEMORY_THRESHOLD_BYTES); + opts.raster.max_batch_bytes = (per_partition_memory_limit + / RASTER_BATCH_BUDGET_DIVISOR) + .max(MIN_RASTER_MAX_BATCH_BYTES) + .min(opts.raster.max_batch_bytes); } // Register the spatial join planner extension @@ -247,6 +261,9 @@ impl SedonaContext { state_builder = register_vendored_optimizer_rules(state_builder)?; state_builder = register_spatial_join_logical_optimizer(state_builder)?; state_builder = register_ensure_loaded_optimizer(state_builder)?; + // Re-batch raster materialization by estimated bytes rather than + // row count (see `sedona_query_planner::raster_batch_budget`). + state_builder = state_builder.with_physical_optimizer_rule(Arc::new(RasterBatchBudgetRule)); state_builder = state_builder.with_query_planner(Arc::new(planner)); let mut state = state_builder.build(); @@ -1139,6 +1156,124 @@ mod tests { ); } + #[tokio::test] + async fn rs_ensureloaded_batches_are_bounded_by_raster_max_batch_bytes() { + use arrow_array::{ArrayRef, RecordBatch}; + use arrow_buffer::Buffer; + use arrow_schema::{Field, Schema}; + use datafusion::arrow::util::pretty::pretty_format_batches; + use sedona_raster::builder::{RasterBuilder, StartBandArgs}; + use sedona_raster::raster_loader::{ + AsyncRasterLoader, RasterLoadRequest, RasterLoadResult, + }; + use sedona_schema::datatypes::SedonaType; + use sedona_schema::raster::BandDataType; + use std::sync::Mutex; + + /// Records how many requests each `load()` carried and returns + /// correctly sized zero buffers. + #[derive(Debug, Default)] + struct CountingLoader { + calls: Mutex>, + } + #[async_trait] + impl AsyncRasterLoader for CountingLoader { + fn name(&self) -> &str { + "mock" + } + fn supports_format(&self, format: Option<&str>) -> bool { + format == Some("mock") + } + async fn load( + &self, + reqs: &[&RasterLoadRequest], + ) -> std::result::Result, arrow_schema::ArrowError> { + self.calls.lock().unwrap().push(reqs.len()); + Ok(reqs + .iter() + .map(|req| { + let elements: i64 = req.source_shape.iter().product(); + let len = elements as usize * req.data_type.byte_size(); + RasterLoadResult::unresolved(Buffer::from_vec(vec![0u8; len]), req) + }) + .collect()) + } + } + + // The interactive constructor is the one that installs the planner + // rules (`new()` only wraps an existing SessionContext). + let ctx = SedonaContext::new_local_interactive().await.unwrap(); + let loader = Arc::new(CountingLoader::default()); + ctx.register_raster_loader(loader.clone()); + + // Six 16 × 16 UInt8 OutDb rasters: 256 bytes each once loaded. + let mut b = RasterBuilder::new(6); + for _ in 0..6 { + b.start_raster_nd( + &[0.0, 1.0, 0.0, 0.0, 0.0, -1.0], + &["y", "x"], + &[16, 16], + None, + ) + .unwrap(); + b.start_band(StartBandArgs { + outdb_uri: Some("mock://tile"), + outdb_format: Some("mock"), + ..StartBandArgs::new(&["y", "x"], &[16, 16], BandDataType::UInt8) + }) + .unwrap(); + b.band_data_writer().append_value([0u8; 0]); + b.finish_band().unwrap(); + b.finish_raster().unwrap(); + } + let rasters: ArrayRef = Arc::new(b.finish().unwrap()); + let raster_field = SedonaType::Raster.to_storage_field("rast", true).unwrap(); + let field = Field::new("rast", rasters.data_type().clone(), true) + .with_metadata(raster_field.metadata().clone()); + let batch = + RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![rasters]).unwrap(); + ctx.ctx.register_batch("tiles", batch).unwrap(); + + // One partition keeps the observed batch sequence deterministic (the + // projection above the exec preserves batch boundaries). + ctx.sql("SET datafusion.execution.target_partitions = 1") + .await + .unwrap() + .collect() + .await + .unwrap(); + // A 600-byte budget fits two 256-byte rasters per slice. + ctx.sql("SET sedona.raster.max_batch_bytes = 600") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let explain = ctx + .sql("EXPLAIN SELECT RS_EnsureLoaded(rast) AS r FROM tiles") + .await + .unwrap() + .collect() + .await + .unwrap(); + let plan = pretty_format_batches(&explain).unwrap().to_string(); + assert!(plan.contains("EnsureLoadedExec"), "{plan}"); + assert!(!plan.contains("AsyncFuncExec"), "{plan}"); + + let batches = ctx + .sql("SELECT RS_EnsureLoaded(rast) AS r FROM tiles") + .await + .unwrap() + .collect() + .await + .unwrap(); + let rows: Vec = batches.iter().map(RecordBatch::num_rows).collect(); + assert_eq!(rows, vec![2, 2, 2]); + // Each slice reaches the loader as one bundled call carrying its rows. + assert_eq!(*loader.calls.lock().unwrap(), vec![2, 2, 2]); + } + #[tokio::test] async fn basic_sql() -> Result<()> { let ctx = SedonaContext::new();