algeon_datafusion/native_exec/config/
mod.rs

1//! Native execution settings copied into rewritten DataFusion plan nodes.
2//!
3//! The config is cloneable and cheap to pass through planning. Backend
4//! admission supplies the installed runtime, allocator, streams, and worker
5//! executor when execution acquires the native query context.
6
7use std::sync::Arc;
8
9use datafusion::execution::context::TaskContext;
10use query_planning::{
11    capability::{AggregateStrategyExecutionOptions, KeyedHashJoinStrategyOptions},
12    chunk::SourceChunkPlanningOptions,
13    correlation::{QueryId, TraceId},
14    dynamic_filter::DynamicFilterExecutionOptions,
15    execution_contracts::SourceChunkPlanningMode,
16    mode::StatisticalAggregateExecutionMode,
17    policy::PlanningPolicy,
18};
19use query_runtime::{
20    observability::RuntimeObservationConfig,
21    runtime::{
22        AdmittedQueryContext, ExecutionContext, IoLifecycle, LakehouseCacheFillConfig, S3SourceStorageProfile,
23        SourceSchedulerPolicy, SourceStorageProfile,
24    },
25};
26
27pub use self::{
28    chunking::SourceChunking,
29    source_reads::{LakehouseByteCache, LakehouseByteCacheFill},
30};
31use super::{
32    admission::{
33        QueryAdmissionFacts, QueryAdmissionSource, QueryAdmissionTicket, TicketAdmissionContext,
34        invocation_admission_ticket,
35    },
36    metrics::{NativeExecutionMetricsCapture, NativeMetricsCapture},
37};
38#[cfg(feature = "cugraph")]
39use crate::cugraph_sql::GraphBfsSingleflight;
40use crate::{
41    error::{Error, Result},
42    native::{ParquetReadStrategyConfig, S3_STORAGE_PROFILE_ID, ambient_s3_endpoint},
43};
44
45/// Adapter-side native execution configuration.
46///
47/// This type holds the DataFusion-facing settings for native cuDF execution.
48/// It does not eagerly create cuDF runtime state. Production `ExecutionContext`
49/// values are acquired only through backend admission when an accepted native
50/// plan begins execution.
51#[derive(Debug, Clone)]
52pub struct AlgeonNativeExecutionConfig {
53    pub(super) source_chunking: SourceChunking,
54    pub(super) source_scheduler_policy: SourceSchedulerPolicy,
55    pub(super) source_chunk_planning_mode: SourceChunkPlanningMode,
56    pub(super) lakehouse_cache_fill_config: LakehouseCacheFillConfig,
57    pub(super) host_output_memory_pool_accounting: bool,
58    pub(super) io_lifecycle: Option<IoLifecycle>,
59    pub(super) statistical_aggregate_execution_mode: StatisticalAggregateExecutionMode,
60    pub(super) planning_policy: PlanningPolicy,
61    pub(super) aggregate_strategy_execution_options: AggregateStrategyExecutionOptions,
62    pub(super) source_chunk_planning_options: SourceChunkPlanningOptions,
63    pub(super) s3_storage: Option<S3SourceStorageProfile>,
64    pub(super) lakehouse_byte_cache: Option<LakehouseByteCache>,
65    /// The ambient AWS endpoint, read once here, that scopes byte-cache
66    /// entries when no S3 storage profile names an endpoint.
67    pub(super) ambient_s3_endpoint: Option<String>,
68    pub(super) query_admission_source: Option<QueryAdmissionSource>,
69    pub(super) query_admission_ticket: Option<Arc<QueryAdmissionTicket>>,
70    pub(super) query_id: Option<QueryId>,
71    pub(super) trace_id: Option<TraceId>,
72    pub(super) plan_task_mapping_datafusion_shape: Option<String>,
73    pub(super) metrics_capture: Option<NativeMetricsCapture>,
74    pub(super) lakehouse_footer_pruning: bool,
75    /// Coalesces identical BFS requests across the queries of one backend.
76    #[cfg(feature = "cugraph")]
77    pub(super) bfs_singleflight: GraphBfsSingleflight,
78}
79
80/// Native execution context plus any service leases that must outlive execution.
81#[derive(Debug)]
82pub(crate) struct AlgeonNativeExecutionContext {
83    pub(super) context: ExecutionContext,
84    pub(super) admission_facts: Option<QueryAdmissionFacts>,
85    pub(super) record_admission_facts: bool,
86    pub(super) _admission_context: Option<AdmittedQueryContext>,
87    pub(super) ticket_context: Option<TicketAdmissionContext>,
88}
89
90mod admission;
91mod chunking;
92mod context;
93mod source_reads;
94#[cfg(test)]
95mod tests;
96
97impl Default for AlgeonNativeExecutionConfig {
98    fn default() -> Self {
99        Self::new(SourceChunking::default())
100    }
101}
102
103impl AlgeonNativeExecutionConfig {
104    /// Device-memory headroom, in bytes, reserved by automatic per-device
105    /// backend profiles for libcudf and driver overhead.
106    pub const AUTO_SIZED_DEVICE_BUDGET_HEADROOM_BYTES: u64 = 500 * 1024 * 1024;
107
108    /// A configuration with `source_chunking` and default settings otherwise.
109    /// It carries no device capacity: admitted execution always runs under the
110    /// selected device's immutable grant.
111    pub fn new(source_chunking: SourceChunking) -> Self {
112        Self {
113            source_chunking,
114            source_scheduler_policy: SourceSchedulerPolicy::default(),
115            source_chunk_planning_mode: SourceChunkPlanningMode::default(),
116            lakehouse_cache_fill_config: LakehouseCacheFillConfig::default(),
117            host_output_memory_pool_accounting: false,
118            io_lifecycle: None,
119            statistical_aggregate_execution_mode: StatisticalAggregateExecutionMode::default(),
120            planning_policy: PlanningPolicy::default(),
121            aggregate_strategy_execution_options: AggregateStrategyExecutionOptions::default(),
122            source_chunk_planning_options: SourceChunkPlanningOptions::default(),
123            s3_storage: None,
124            lakehouse_byte_cache: None,
125            ambient_s3_endpoint: ambient_s3_endpoint(),
126            query_admission_source: None,
127            query_admission_ticket: None,
128            query_id: None,
129            trace_id: None,
130            plan_task_mapping_datafusion_shape: None,
131            metrics_capture: Some(NativeMetricsCapture::standalone()),
132            lakehouse_footer_pruning: false,
133            #[cfg(feature = "cugraph")]
134            bfs_singleflight: GraphBfsSingleflight::default(),
135        }
136    }
137
138    /// Set bounded source scheduling controls.
139    pub fn with_source_scheduler_policy(mut self, source_scheduler_policy: SourceSchedulerPolicy) -> Self {
140        self.source_scheduler_policy = source_scheduler_policy;
141        self
142    }
143
144    /// Set bounded source chunk planning mode.
145    pub fn with_source_chunk_planning_mode(mut self, mode: SourceChunkPlanningMode) -> Self {
146        self.source_chunk_planning_mode = mode;
147        self
148    }
149
150    /// Set the lakehouse cache-fill resource policy projected into backend execution.
151    pub fn with_lakehouse_cache_fill_config(mut self, config: LakehouseCacheFillConfig) -> Self {
152        self.lakehouse_cache_fill_config = config;
153        self
154    }
155
156    /// Account native host Arrow output batches against DataFusion's host
157    /// [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool).
158    ///
159    /// This is intentionally host-only. Device memory admission and grants stay
160    /// in the native query service/governor ledger.
161    pub fn with_host_output_memory_pool_accounting(mut self, enabled: bool) -> Self {
162        self.host_output_memory_pool_accounting = enabled;
163        self
164    }
165
166    /// Attaches the server generation coordinator used by file-backed native I/O.
167    #[cfg(feature = "server")]
168    pub(crate) fn with_io_lifecycle(mut self, lifecycle: IoLifecycle) -> Self {
169        self.io_lifecycle = Some(lifecycle);
170        self
171    }
172
173    /// Attach a diagnostic query id to native execution metrics.
174    pub fn with_query_id(mut self, query_id: impl Into<QueryId>) -> Self {
175        self.query_id = Some(query_id.into());
176        self
177    }
178
179    /// Attach a query id only when the caller has not supplied one.
180    pub(crate) fn with_query_id_if_absent(mut self, query_id: impl Into<QueryId>) -> Self {
181        if self.query_id.is_none() {
182            self.query_id = Some(query_id.into());
183        }
184        self
185    }
186
187    /// Attach a parent trace id to native diagnostics.
188    pub fn with_trace_id(mut self, trace_id: impl Into<TraceId>) -> Self {
189        self.trace_id = Some(trace_id.into());
190        self
191    }
192
193    /// Attach the DataFusion candidate shape used to annotate plan/task mapping rows.
194    pub(crate) fn with_plan_task_mapping_datafusion_shape(mut self, shape: impl Into<String>) -> Self {
195        let shape = shape.into();
196        if !shape.is_empty() {
197            self.plan_task_mapping_datafusion_shape = Some(shape);
198        }
199        self
200    }
201
202    /// Borrows the optional diagnostic query correlation id carried by this execution template.
203    pub fn query_id(&self) -> Option<&QueryId> {
204        self.query_id.as_ref()
205    }
206
207    /// Borrows the optional parent trace correlation id carried by this execution template.
208    pub fn trace_id(&self) -> Option<&TraceId> {
209        self.trace_id.as_ref()
210    }
211
212    pub(crate) fn plan_task_mapping_datafusion_shape(&self) -> Option<&str> {
213        self.plan_task_mapping_datafusion_shape.as_deref()
214    }
215
216    pub(crate) fn new_metrics_capture(&self) -> Option<NativeExecutionMetricsCapture> {
217        self.query_admission_ticket
218            .as_ref()
219            .map(|ticket| ticket.metrics_capture())
220            .or_else(|| self.metrics_capture.as_ref().map(NativeMetricsCapture::new_execution))
221    }
222
223    /// Returns the bounded runtime observation retention policy copied into native execution.
224    pub fn runtime_observation_config(&self) -> RuntimeObservationConfig {
225        self.metrics_capture
226            .as_ref()
227            .map_or_else(RuntimeObservationConfig::native_execution, NativeMetricsCapture::runtime_observation_config)
228    }
229
230    /// Prunes `s3://` Parquet row groups from their footers before native
231    /// execution admits the plan. Off by default.
232    pub fn with_lakehouse_footer_pruning(mut self, enabled: bool) -> Self {
233        self.lakehouse_footer_pruning = enabled;
234        self
235    }
236
237    /// Whether footer pruning runs before admission.
238    pub fn lakehouse_footer_pruning(&self) -> bool {
239        self.lakehouse_footer_pruning
240    }
241
242    /// Select the execution contract for floating statistical aggregates.
243    ///
244    /// Defaults to [`StatisticalAggregateExecutionMode::GpuNativeTolerant`],
245    /// whose results can differ from DataFusion in floating-point bits. Select
246    /// [`StatisticalAggregateExecutionMode::DataFusionExactHostWelford`] when
247    /// results must be identical to DataFusion's. In that mode a candidate whose
248    /// statistical aggregate input DataFusion reads through more than one
249    /// partition stays on DataFusion, whose own result then depends on the
250    /// order partitions deliver rows.
251    pub fn with_statistical_aggregate_execution_mode(mut self, mode: StatisticalAggregateExecutionMode) -> Self {
252        self.statistical_aggregate_execution_mode = mode;
253        self
254    }
255
256    /// Set native dynamic-filter execution policy.
257    ///
258    /// Route discovery, admission, and runtime payload handling read one
259    /// threshold set through [`Self::planning_policy`]. It is intentionally
260    /// not exposed as a session/server setting until the expert surface is reviewed.
261    pub fn with_dynamic_filter_execution_options(mut self, options: DynamicFilterExecutionOptions) -> Self {
262        self.planning_policy.dynamic_filter = options;
263        self
264    }
265
266    /// Set native keyed-hash join strategy policy.
267    ///
268    /// Capability admission, admission reports, and runtime observed-build
269    /// strategy selection read one threshold set through
270    /// [`Self::planning_policy`]. It is intentionally not exposed as a
271    /// session/server setting until the expert surface is reviewed.
272    pub fn with_keyed_hash_join_strategy_options(mut self, options: KeyedHashJoinStrategyOptions) -> Self {
273        self.planning_policy.keyed_hash_join = options;
274        self
275    }
276
277    /// Set native grouped aggregate strategy execution policy.
278    ///
279    /// This policy is copied into native execution contexts so runtime
280    /// aggregate strategy selection, exact-cardinality probes, and strategy
281    /// metric rows share one threshold set. It is intentionally not exposed as
282    /// a session/server setting until the expert surface is reviewed.
283    pub fn with_aggregate_strategy_execution_options(mut self, options: AggregateStrategyExecutionOptions) -> Self {
284        self.aggregate_strategy_execution_options = options;
285        self
286    }
287
288    /// Set expert local-throughput source chunk planning policy.
289    ///
290    /// This policy is intentionally Rust-only; public source chunk settings remain
291    /// the bounded memory policy and source chunk planning mode.
292    pub fn with_source_chunk_planning_options(mut self, options: SourceChunkPlanningOptions) -> Self {
293        self.source_chunk_planning_options = options;
294        self
295    }
296
297    /// The S3 storage every `s3://` source reads through, when configured.
298    pub fn s3_storage(&self) -> Option<&S3SourceStorageProfile> {
299        self.s3_storage.as_ref()
300    }
301
302    /// The runtime storage profiles native remote reads and cache fills
303    /// resolve: the S3 storage under [`S3_STORAGE_PROFILE_ID`], when set.
304    pub(crate) fn source_storage_profiles(&self) -> Vec<SourceStorageProfile> {
305        self.s3_storage.iter().map(|profile| SourceStorageProfile::s3(S3_STORAGE_PROFILE_ID, profile.clone())).collect()
306    }
307
308    /// How lowering resolves an `s3://` Parquet file.
309    pub(crate) fn source_read_strategy(&self) -> ParquetReadStrategyConfig {
310        ParquetReadStrategyConfig::new(
311            self.lakehouse_byte_cache.as_ref(),
312            self.s3_storage.as_ref(),
313            self.ambient_s3_endpoint.as_deref(),
314        )
315    }
316
317    #[cfg(feature = "cugraph")]
318    pub(crate) fn bfs_singleflight(&self) -> &GraphBfsSingleflight {
319        &self.bfs_singleflight
320    }
321
322    /// Reads every `s3://` source through `profile`: native reads, cache
323    /// fills, footer pruning, and the DataFusion object store that sessions
324    /// installed with this config resolve for `s3://` URLs.
325    pub fn with_s3_storage(mut self, profile: S3SourceStorageProfile) -> Self {
326        self.s3_storage = Some(profile);
327        self
328    }
329
330    /// Keeps a local byte cache of `s3://` Parquet objects.
331    pub fn with_lakehouse_byte_cache(mut self, cache: LakehouseByteCache) -> Self {
332        self.lakehouse_byte_cache = Some(cache);
333        self
334    }
335
336    /// Source chunking of native reads.
337    pub fn source_chunking(&self) -> SourceChunking {
338        self.source_chunking
339    }
340
341    /// Maximum active native source reads allowed per bounded query.
342    pub fn max_source_reads_in_flight(&self) -> usize {
343        self.source_scheduler_policy.max_source_reads_in_flight()
344    }
345
346    /// Source scheduling policy used by bounded native source reads.
347    pub fn source_scheduler_policy(&self) -> SourceSchedulerPolicy {
348        self.source_scheduler_policy
349    }
350
351    /// Source chunk planning mode used by bounded native source reads.
352    pub fn source_chunk_planning_mode(&self) -> SourceChunkPlanningMode {
353        self.source_chunk_planning_mode
354    }
355
356    /// Lakehouse cache-fill resource policy projected into backend execution.
357    pub fn lakehouse_cache_fill_config(&self) -> &LakehouseCacheFillConfig {
358        &self.lakehouse_cache_fill_config
359    }
360
361    /// Whether host Arrow output from native execution is charged to
362    /// DataFusion's host-side memory pool.
363    pub fn host_output_memory_pool_accounting(&self) -> bool {
364        self.host_output_memory_pool_accounting
365    }
366
367    /// Execution contract for floating statistical aggregates.
368    pub fn statistical_aggregate_execution_mode(&self) -> StatisticalAggregateExecutionMode {
369        self.statistical_aggregate_execution_mode
370    }
371
372    /// Join and dynamic-filter thresholds native plans are optimized,
373    /// admitted, and executed under.
374    pub fn planning_policy(&self) -> PlanningPolicy {
375        self.planning_policy
376    }
377
378    /// Grouped aggregate strategy execution policy copied into native contexts.
379    pub fn aggregate_strategy_execution_options(&self) -> AggregateStrategyExecutionOptions {
380        self.aggregate_strategy_execution_options
381    }
382
383    /// Expert local-throughput source chunk planning policy copied into native contexts.
384    pub fn source_chunk_planning_options(&self) -> SourceChunkPlanningOptions {
385        self.source_chunk_planning_options
386    }
387
388    /// Validate internal expert query policies before planning or execution.
389    pub(crate) fn validate_query_policies(&self) -> Result<()> {
390        self.planning_policy.dynamic_filter.validate().map_err(|error| {
391            Error::from_native(
392                "algeon_datafusion.native_exec.config.validate",
393                "invalid native dynamic-filter execution policy",
394                self.stamp_native_error_correlation(error),
395            )
396        })?;
397        self.planning_policy.keyed_hash_join.validate().map_err(|error| {
398            Error::from_native(
399                "algeon_datafusion.native_exec.config.validate",
400                "invalid native keyed hash-join strategy policy",
401                self.stamp_native_error_correlation(error),
402            )
403        })?;
404        self.aggregate_strategy_execution_options.validate().map_err(|error| {
405            Error::from_native(
406                "algeon_datafusion.native_exec.config.validate",
407                "invalid native aggregate strategy execution policy",
408                self.stamp_native_error_correlation(error),
409            )
410        })?;
411        self.source_chunk_planning_options.validate().map_err(|error| {
412            Error::from_native(
413                "algeon_datafusion.native_exec.config.validate",
414                "invalid native source chunk planning policy",
415                self.stamp_native_error_correlation(error),
416            )
417        })?;
418        Ok(())
419    }
420}
algeon_datafusion/native_exec/config/mod.rs

algeon_datafusion/native_exec/config/
mod.rs
