delta_funnel/query_engine/datafusion/
session.rs

1//! DataFusion session options for query execution.
2//!
3//! These options are query-engine settings. They are intentionally separate
4//! from provider file-read scheduling and from future MSSQL/TDS writer options.
5
6use datafusion::prelude::{SessionConfig, SessionContext};
7
8use crate::DeltaFunnelError;
9use crate::pipeline::{BatchPipelinePhase, validate_nonzero_usize_option};
10
11/// Query execution options applied before DataFusion planning and execution.
12#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
13pub struct QueryOptions {
14    /// Optional DataFusion execution partition target.
15    ///
16    /// When set, this is passed to DataFusion as the session target partition
17    /// count before query planning. When unset, DataFusion's session default
18    /// remains in effect. Provider scans treat the DataFusion target as a cap
19    /// while deriving their own conservative physical scan partition target;
20    /// automatic scan-wide read capacity is resolved from that derived target.
21    pub target_partitions: Option<usize>,
22
23    /// Optional DataFusion execution batch size.
24    ///
25    /// This is an upstream query-engine batch-size request. Arbitrary
26    /// DataFusion operators may still produce smaller or differently shaped
27    /// final output batches.
28    pub output_batch_size: Option<usize>,
29}
30
31impl QueryOptions {
32    /// Validates query options before any source reads or target side effects.
33    ///
34    /// # Errors
35    ///
36    /// Returns [`DeltaFunnelError::BatchPipeline`] when a configured numeric
37    /// option is zero.
38    pub fn validate(&self) -> Result<(), DeltaFunnelError> {
39        if let Some(target_partitions) = self.target_partitions {
40            validate_nonzero_usize_option(
41                BatchPipelinePhase::Configuration,
42                "target_partitions",
43                target_partitions,
44            )?;
45        }
46
47        if let Some(output_batch_size) = self.output_batch_size {
48            validate_nonzero_usize_option(
49                BatchPipelinePhase::Configuration,
50                "output_batch_size",
51                output_batch_size,
52            )?;
53        }
54
55        Ok(())
56    }
57
58    /// Applies validated query options to a DataFusion session config.
59    ///
60    /// # Errors
61    ///
62    /// Returns [`DeltaFunnelError::BatchPipeline`] when validation fails.
63    pub fn apply_to_session_config(
64        &self,
65        mut config: SessionConfig,
66    ) -> Result<SessionConfig, DeltaFunnelError> {
67        self.validate()?;
68
69        if let Some(target_partitions) = self.target_partitions {
70            config = config.with_target_partitions(target_partitions);
71        }
72        if let Some(output_batch_size) = self.output_batch_size {
73            config = config.with_batch_size(output_batch_size);
74        }
75
76        Ok(config)
77    }
78}
79
80/// Builds a DataFusion session config from DeltaFunnel query options.
81///
82/// # Errors
83///
84/// Returns [`DeltaFunnelError::BatchPipeline`] when validation fails.
85pub fn datafusion_session_config(options: QueryOptions) -> Result<SessionConfig, DeltaFunnelError> {
86    options.apply_to_session_config(SessionConfig::new())
87}
88
89/// Builds a DataFusion session context from DeltaFunnel query options.
90///
91/// # Errors
92///
93/// Returns [`DeltaFunnelError::BatchPipeline`] when validation fails.
94pub fn datafusion_session_context(
95    options: QueryOptions,
96) -> Result<SessionContext, DeltaFunnelError> {
97    Ok(SessionContext::new_with_config(datafusion_session_config(
98        options,
99    )?))
100}
101
102#[cfg(test)]
103mod tests {
104    use datafusion::prelude::SessionConfig;
105
106    use super::{QueryOptions, datafusion_session_config, datafusion_session_context};
107    use crate::{BatchPipelinePhase, DeltaFunnelError};
108
109    #[test]
110    fn omitted_query_options_preserve_datafusion_defaults() -> Result<(), Box<dyn std::error::Error>>
111    {
112        let default_config = SessionConfig::new();
113        let configured = datafusion_session_config(QueryOptions::default())?;
114
115        assert_eq!(configured.batch_size(), default_config.batch_size());
116        assert_eq!(
117            configured.target_partitions(),
118            default_config.target_partitions()
119        );
120
121        Ok(())
122    }
123
124    #[test]
125    fn output_batch_size_is_applied_to_session_config() -> Result<(), Box<dyn std::error::Error>> {
126        let configured = datafusion_session_config(QueryOptions {
127            target_partitions: None,
128            output_batch_size: Some(17),
129        })?;
130
131        assert_eq!(configured.batch_size(), 17);
132        assert_eq!(
133            configured.target_partitions(),
134            SessionConfig::new().target_partitions()
135        );
136
137        Ok(())
138    }
139
140    #[test]
141    fn target_partitions_is_applied_to_session_config() -> Result<(), Box<dyn std::error::Error>> {
142        let configured = datafusion_session_config(QueryOptions {
143            target_partitions: Some(3),
144            output_batch_size: None,
145        })?;
146
147        assert_eq!(configured.target_partitions(), 3);
148        assert_eq!(configured.batch_size(), SessionConfig::new().batch_size());
149
150        Ok(())
151    }
152
153    #[test]
154    fn target_partitions_and_output_batch_size_are_independent()
155    -> Result<(), Box<dyn std::error::Error>> {
156        let configured = datafusion_session_config(QueryOptions {
157            target_partitions: Some(4),
158            output_batch_size: Some(11),
159        })?;
160
161        assert_eq!(configured.target_partitions(), 4);
162        assert_eq!(configured.batch_size(), 11);
163
164        Ok(())
165    }
166
167    #[test]
168    fn session_context_uses_configured_query_options() -> Result<(), Box<dyn std::error::Error>> {
169        let ctx = datafusion_session_context(QueryOptions {
170            target_partitions: Some(2),
171            output_batch_size: Some(9),
172        })?;
173
174        assert_eq!(ctx.state().config().target_partitions(), 2);
175        assert_eq!(ctx.state().config().batch_size(), 9);
176
177        Ok(())
178    }
179
180    #[test]
181    fn validation_rejects_zero_output_batch_size() {
182        let error = QueryOptions {
183            target_partitions: None,
184            output_batch_size: Some(0),
185        }
186        .validate();
187
188        assert!(matches!(
189            error,
190            Err(DeltaFunnelError::BatchPipeline {
191                phase: BatchPipelinePhase::Configuration,
192                option: "output_batch_size",
193                ..
194            })
195        ));
196    }
197
198    #[test]
199    fn validation_rejects_zero_target_partitions() {
200        let error = QueryOptions {
201            target_partitions: Some(0),
202            output_batch_size: None,
203        }
204        .validate();
205
206        assert!(matches!(
207            error,
208            Err(DeltaFunnelError::BatchPipeline {
209                phase: BatchPipelinePhase::Configuration,
210                option: "target_partitions",
211                ..
212            })
213        ));
214    }
215}
delta_funnel/query_engine/datafusion/session.rs

delta_funnel/query_engine/datafusion/
session.rs
