laminar_sql/datafusion/
mod.rs

1//! DataFusion integration for SQL processing.
2
3/// Marker UDFs for the `ai_*` SQL functions (rewritten by the AI operator).
4pub mod ai_udf;
5mod bridge;
6mod channel_source;
7/// Lambda higher-order functions for arrays and maps (F-SCHEMA-015 Tier 3)
8pub mod complex_type_lambda;
9/// Array, Struct, and Map scalar UDFs (F-SCHEMA-015)
10pub mod complex_type_udf;
11mod exec;
12/// End-to-end streaming SQL execution
13pub mod execute;
14/// Format bridge UDFs for inline format conversion
15pub mod format_bridge_udf;
16/// LaminarDB streaming JSON extension UDFs (F-SCHEMA-013)
17pub mod json_extensions;
18/// SQL/JSON path query compiler and scalar UDFs
19pub mod json_path;
20/// JSON table-valued functions (array/object expansion)
21pub mod json_tvf;
22/// JSONB binary format types for JSON UDF evaluation
23pub mod json_types;
24/// PostgreSQL-compatible JSON aggregate UDAFs
25pub mod json_udaf;
26/// PostgreSQL-compatible JSON scalar UDFs
27pub mod json_udf;
28/// Live source provider for streaming execution with plan caching
29pub mod live_source;
30/// Lookup join plan node for DataFusion.
31pub mod lookup_join;
32/// Physical execution plan and extension planner for lookup joins.
33pub mod lookup_join_exec;
34/// Processing-time UDF for `PROCTIME()` support
35pub mod proctime_udf;
36mod source;
37mod table_provider;
38/// Dynamic watermark filter for scan-level late-data pruning
39/// Watermark UDF for current watermark access
40pub mod watermark_udf;
41/// Window function UDFs (TUMBLE, HOP, SESSION, CUMULATE)
42pub mod window_udf;
43
44pub use ai_udf::{ai_function_markers, AiFunctionMarker};
45pub use bridge::{BridgeSendError, BridgeSender, BridgeStream, BridgeTrySendError, StreamBridge};
46pub use channel_source::ChannelStreamSource;
47pub use complex_type_lambda::{
48    register_lambda_functions, ArrayFilter, ArrayReduce, ArrayTransform, MapFilter,
49    MapTransformValues,
50};
51pub use complex_type_udf::{
52    register_complex_type_functions, MapContainsKey, MapFromArrays, MapKeys, MapValues, StructDrop,
53    StructExtract, StructMerge, StructRename, StructSet,
54};
55pub use exec::StreamingScanExec;
56pub use execute::{execute_streaming_sql, DdlResult, QueryResult, StreamingSqlResult};
57pub use format_bridge_udf::{FromJsonUdf, ParseEpochUdf, ParseTimestampUdf, ToJsonUdf};
58pub use json_extensions::{
59    register_json_extensions, JsonInferSchema, JsonToColumns, JsonbDeepMerge, JsonbExcept,
60    JsonbFlatten, JsonbMerge, JsonbPick, JsonbRenameKeys, JsonbStripNulls, JsonbUnflatten,
61};
62pub use json_path::{CompiledJsonPath, JsonPathStep, JsonbPathExistsUdf, JsonbPathMatchUdf};
63pub use json_tvf::{
64    register_json_table_functions, JsonbArrayElementsTextTvf, JsonbArrayElementsTvf,
65    JsonbEachTextTvf, JsonbEachTvf, JsonbObjectKeysTvf,
66};
67pub use json_udaf::{JsonAgg, JsonObjectAgg};
68pub use json_udf::{
69    JsonBuildArray, JsonBuildObject, JsonTypeof, JsonbContainedBy, JsonbContains, JsonbExists,
70    JsonbExistsAll, JsonbExistsAny, JsonbGet, JsonbGetIdx, JsonbGetPath, JsonbGetPathText,
71    JsonbGetText, JsonbGetTextIdx, ToJsonb,
72};
73pub use live_source::{LiveSourceHandle, LiveSourceProvider};
74pub use lookup_join_exec::{
75    LookupJoinExec, LookupJoinExtensionPlanner, LookupSnapshot, LookupTableRegistry,
76    PartialLookupJoinExec, PartialLookupState, RegisteredLookup,
77};
78pub use proctime_udf::ProcTimeUdf;
79pub use source::{SortColumn, StreamSource, StreamSourceRef};
80pub use table_provider::StreamingTableProvider;
81pub use watermark_udf::WatermarkUdf;
82pub use window_udf::{
83    CumulateWindowEnd, CumulateWindowStart, HopWindowEnd, HopWindowStart, SessionWindowStart,
84    TumbleWindowEnd, TumbleWindowStart,
85};
86
87use std::sync::atomic::AtomicI64;
88use std::sync::Arc;
89
90use datafusion::execution::SessionStateBuilder;
91use datafusion::prelude::*;
92use datafusion_expr::ScalarUDF;
93
94use crate::planner::streaming_optimizer::{StreamingPhysicalValidator, StreamingValidatorMode};
95
96/// Returns a base `SessionConfig` with identifier normalization disabled.
97///
98/// DataFusion's default behaviour lowercases all unquoted SQL identifiers
99/// (per the SQL standard). LaminarDB disables this so that mixed-case
100/// column names from external sources (Kafka, CDC, WebSocket) can be
101/// referenced without double-quoting.
102#[must_use]
103pub fn base_session_config() -> SessionConfig {
104    let mut config = SessionConfig::new();
105    config.options_mut().sql_parser.enable_ident_normalization = false;
106    // Single partition for streaming micro-batch execution. Multi-partition
107    // plans contain stateful operators (RepartitionExec) that cannot be
108    // reused across cycles, causing panics on cached physical plans.
109    config = config.with_target_partitions(1);
110    config
111}
112
113/// Creates a `DataFusion` session context with identifier normalization
114/// disabled.
115///
116/// Suitable for ad-hoc / non-streaming queries (filters, lookups).
117/// For streaming workloads prefer [`create_streaming_context`].
118#[must_use]
119pub fn create_session_context() -> SessionContext {
120    SessionContext::new_with_config(base_session_config())
121}
122
123/// Creates a `DataFusion` session context configured for streaming queries.
124///
125/// The context is configured with:
126/// - Batch size of 8192 (balanced for streaming throughput)
127/// - Single partition (streaming sources are typically not partitioned)
128/// - Identifier normalization disabled (mixed-case columns work unquoted)
129/// - All streaming UDFs registered (TUMBLE, HOP, SESSION, WATERMARK)
130/// - `StreamingPhysicalValidator` in `Reject` mode (blocks unsafe plans)
131///
132/// The watermark UDF is initialized with no watermark set (returns NULL).
133/// Use [`register_streaming_functions_with_watermark`] to provide a live
134/// watermark source.
135///
136/// # Example
137///
138/// ```rust,ignore
139/// let ctx = create_streaming_context();
140/// ctx.register_table("events", provider)?;
141/// let df = ctx.sql("SELECT * FROM events").await?;
142/// ```
143#[must_use]
144pub fn create_streaming_context() -> SessionContext {
145    create_streaming_context_with_validator(StreamingValidatorMode::Reject)
146}
147
148/// Creates a streaming context with a configurable validator mode.
149///
150/// Same as [`create_streaming_context`] but allows choosing how the
151/// [`StreamingPhysicalValidator`] handles plan violations.
152///
153/// Use [`StreamingValidatorMode::Off`] to get the previous behaviour
154/// (no plan-time validation).
155#[must_use]
156pub fn create_streaming_context_with_validator(mode: StreamingValidatorMode) -> SessionContext {
157    let config = base_session_config().with_batch_size(8192);
158
159    let ctx = if matches!(mode, StreamingValidatorMode::Off) {
160        SessionContext::new_with_config(config)
161    } else {
162        // Build a default state to get the standard optimizer rules, then
163        // prepend our streaming validator so it fires before DataFusion's
164        // built-in SanityCheckPlan (which produces generic error messages).
165        let default_state = SessionStateBuilder::new()
166            .with_config(config.clone())
167            .with_default_features()
168            .build();
169        let mut rules: Vec<
170            Arc<dyn datafusion::physical_optimizer::PhysicalOptimizerRule + Send + Sync>,
171        > = vec![Arc::new(StreamingPhysicalValidator::new(mode))];
172        rules.extend(default_state.physical_optimizers().iter().cloned());
173
174        let state = SessionStateBuilder::new()
175            .with_config(config)
176            .with_default_features()
177            .with_physical_optimizer_rules(rules)
178            .build();
179        SessionContext::new_with_state(state)
180    };
181
182    register_streaming_functions(&ctx);
183    ctx
184}
185
186/// Window-time, JSON, complex-type, lambda, and `proctime()` UDFs —
187/// every streaming UDF except `watermark()`. Pulled out of the public
188/// `register_streaming_functions*` entry points so they share a single
189/// list and stay in sync.
190fn register_non_watermark_udfs(ctx: &SessionContext) {
191    ctx.register_udf(ScalarUDF::new_from_impl(TumbleWindowStart::new()));
192    ctx.register_udf(ScalarUDF::new_from_impl(TumbleWindowEnd::new()));
193    ctx.register_udf(ScalarUDF::new_from_impl(HopWindowStart::new()));
194    ctx.register_udf(ScalarUDF::new_from_impl(HopWindowEnd::new()));
195    ctx.register_udf(ScalarUDF::new_from_impl(SessionWindowStart::new()));
196    ctx.register_udf(ScalarUDF::new_from_impl(CumulateWindowStart::new()));
197    ctx.register_udf(ScalarUDF::new_from_impl(CumulateWindowEnd::new()));
198    ctx.register_udf(ScalarUDF::new_from_impl(ProcTimeUdf::new()));
199    for marker in ai_function_markers() {
200        ctx.register_udf(marker);
201    }
202    register_json_functions(ctx);
203    register_json_extensions(ctx);
204    register_complex_type_functions(ctx);
205    register_lambda_functions(ctx);
206}
207
208/// Registers `LaminarDB` streaming UDFs with a session context. The
209/// `watermark()` UDF is registered in unset mode (always returns NULL);
210/// use [`register_streaming_functions_with_watermark`] to provide a
211/// live watermark source from Ring 0.
212pub fn register_streaming_functions(ctx: &SessionContext) {
213    register_non_watermark_udfs(ctx);
214    ctx.register_udf(ScalarUDF::new_from_impl(WatermarkUdf::unset()));
215}
216
217/// Registers streaming UDFs with a live watermark source — same as
218/// [`register_streaming_functions`] but `watermark()` reads
219/// `watermark_ms` (in milliseconds since epoch; values < 0 mean "no
220/// watermark", returning NULL).
221pub fn register_streaming_functions_with_watermark(
222    ctx: &SessionContext,
223    watermark_ms: Arc<AtomicI64>,
224) {
225    register_non_watermark_udfs(ctx);
226    ctx.register_udf(ScalarUDF::new_from_impl(WatermarkUdf::new(watermark_ms)));
227}
228
229/// Registers all PostgreSQL-compatible JSON UDFs and UDAFs
230/// with the given `SessionContext`.
231pub fn register_json_functions(ctx: &SessionContext) {
232    // Extraction operators
233    ctx.register_udf(ScalarUDF::new_from_impl(JsonbGet::new()));
234    ctx.register_udf(ScalarUDF::new_from_impl(JsonbGetIdx::new()));
235    ctx.register_udf(ScalarUDF::new_from_impl(JsonbGetText::new()));
236    ctx.register_udf(ScalarUDF::new_from_impl(JsonbGetTextIdx::new()));
237    ctx.register_udf(ScalarUDF::new_from_impl(JsonbGetPath::new()));
238    ctx.register_udf(ScalarUDF::new_from_impl(JsonbGetPathText::new()));
239
240    // Existence operators
241    ctx.register_udf(ScalarUDF::new_from_impl(JsonbExists::new()));
242    ctx.register_udf(ScalarUDF::new_from_impl(JsonbExistsAny::new()));
243    ctx.register_udf(ScalarUDF::new_from_impl(JsonbExistsAll::new()));
244
245    // Containment operators
246    ctx.register_udf(ScalarUDF::new_from_impl(JsonbContains::new()));
247    ctx.register_udf(ScalarUDF::new_from_impl(JsonbContainedBy::new()));
248
249    // Interrogation / construction
250    ctx.register_udf(ScalarUDF::new_from_impl(JsonTypeof::new()));
251    ctx.register_udf(ScalarUDF::new_from_impl(JsonBuildObject::new()));
252    ctx.register_udf(ScalarUDF::new_from_impl(JsonBuildArray::new()));
253    ctx.register_udf(ScalarUDF::new_from_impl(ToJsonb::new()));
254
255    // Aggregates
256    ctx.register_udaf(datafusion_expr::AggregateUDF::new_from_impl(JsonAgg::new()));
257    ctx.register_udaf(datafusion_expr::AggregateUDF::new_from_impl(
258        JsonObjectAgg::new(),
259    ));
260
261    // Format bridge functions
262    ctx.register_udf(ScalarUDF::new_from_impl(ParseEpochUdf::new()));
263    ctx.register_udf(ScalarUDF::new_from_impl(ParseTimestampUdf::new()));
264    ctx.register_udf(ScalarUDF::new_from_impl(ToJsonUdf::new()));
265    ctx.register_udf(ScalarUDF::new_from_impl(FromJsonUdf::new()));
266
267    // JSON path query functions (scalar)
268    ctx.register_udf(ScalarUDF::new_from_impl(JsonbPathExistsUdf::new()));
269    ctx.register_udf(ScalarUDF::new_from_impl(JsonbPathMatchUdf::new()));
270
271    // JSON table-valued functions
272    register_json_table_functions(ctx);
273}
274
275#[cfg(test)]
276mod tests;
laminar_sql/datafusion/mod.rs

laminar_sql/datafusion/
mod.rs
