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-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 } } 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-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}"); + } +} 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();