Skip to main content

datum_sql/
lib.rs

1#![forbid(unsafe_code)]
2#![doc = include_str!("../README.md")]
3
4mod join;
5mod sink;
6mod time;
7mod window;
8
9use std::collections::BTreeMap;
10use std::fmt;
11use std::sync::atomic::{AtomicU64, Ordering};
12use std::sync::{Arc, Mutex};
13
14use arrow::array::{Array, RecordBatch};
15use arrow::compute::filter_record_batch;
16use arrow::datatypes::{DataType, Field, FieldRef, SchemaRef};
17use async_trait::async_trait;
18use datafusion::catalog::{Session, TableProvider};
19use datafusion::common::ColumnStatistics;
20use datafusion::common::ScalarValue;
21use datafusion::common::assert_or_internal_err;
22use datafusion::common::cast::as_boolean_array;
23use datafusion::common::stats::Precision;
24use datafusion::common::{DataFusionError, Result, project_schema};
25use datafusion::execution::TaskContext;
26use datafusion::execution::context::SessionContext;
27use datafusion::logical_expr::{
28    ColumnarValue, Expr, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
29    TableProviderFilterPushDown, TableType, Volatility,
30};
31use datafusion::physical_expr::EquivalenceProperties;
32use datafusion::physical_expr::projection::ProjectionExpr;
33use datafusion::physical_plan::PhysicalExpr;
34use datafusion::physical_plan::Statistics;
35use datafusion::physical_plan::aggregates::AggregateExec;
36#[allow(deprecated)]
37use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
38use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
39use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType, SchedulingType};
40use datafusion::physical_plan::filter::FilterExec;
41use datafusion::physical_plan::joins::HashJoinExec;
42use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
43use datafusion::physical_plan::memory::MemoryStream;
44use datafusion::physical_plan::projection::ProjectionExec;
45use datafusion::physical_plan::{
46    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
47    SendableRecordBatchStream,
48};
49use datafusion::prelude::SessionConfig;
50use datafusion::sql::parser::{DFParser, Statement as DfStatement};
51use datafusion::sql::sqlparser::ast::{Statement as SqlStatement, TableObject as SqlTableObject};
52use datum::{Keep, KillSwitches, Sink, Source, StreamError, StreamResult};
53
54pub use join::{
55    StreamingJoinConfig, StreamingJoinMetrics, StreamingJoinMode, StreamingJoinStateLimits,
56    StreamingJoinWindow,
57};
58use sink::{AppendSinkSource, RegisteredSqlSink, sink_registry_error};
59pub use sink::{CommittableRecordBatch, SourceCommit, SqlSourcePosition};
60pub use time::{
61    ContinuousQueryHandle, EventTimeConfig, SqlBarrier, SqlEvent, SqlEventPayload, Watermark,
62    WatermarkStrategy, assign_event_time_watermarks, data_events, map_sql_event_data,
63};
64pub use window::WindowedAggregationMetrics;
65
66pub mod connect;
67#[cfg(feature = "cdc")]
68pub use connect::cdc::{CdcChangelogBatch, CdcSourcePosition, CdcTableAdapter};
69#[cfg(feature = "mq")]
70pub use connect::mq::{
71    JsonRowFormat, KafkaPartitionOffset, KafkaRecordBatch, KafkaSourcePosition, MqPayloadFormat,
72};
73pub use connect::{BatchEnvelope, EnvelopedBatch, EnvelopedRecordBatch};
74
75type SourceFactory = Arc<dyn Fn() -> Source<RecordBatch> + Send + Sync>;
76type CommittableSourceFactory = Arc<dyn Fn() -> Source<CommittableRecordBatch> + Send + Sync>;
77type ChangelogSourceFactory = Arc<dyn Fn() -> Source<ChangelogBatch> + Send + Sync>;
78
79/// Row-level operation for SQL streams that can revise prior output.
80#[derive(Debug, Clone, Copy, PartialEq, Eq)]
81pub enum ChangeOp {
82    /// Adds a row that was not previously present.
83    Insert,
84    /// Retracts a row without a replacement.
85    Delete,
86    /// First half of an adjacent update pair, carrying the old row.
87    UpdateDelete,
88    /// Second half of an adjacent update pair, carrying the new row.
89    UpdateInsert,
90}
91
92impl ChangeOp {
93    /// Returns true for rows that add state to an updating stream.
94    #[must_use]
95    pub const fn is_positive(self) -> bool {
96        matches!(self, Self::Insert | Self::UpdateInsert)
97    }
98
99    /// Returns true for rows that retract state from an updating stream.
100    #[must_use]
101    pub const fn is_retraction(self) -> bool {
102        matches!(self, Self::Delete | Self::UpdateDelete)
103    }
104}
105
106/// A typed changelog envelope over a plain Arrow `RecordBatch`.
107///
108/// Append-only SQL stays as `Source<RecordBatch>`. This type is only used once a
109/// source or operator can emit retractions.
110#[derive(Clone)]
111pub struct ChangelogBatch {
112    ops: Arc<[ChangeOp]>,
113    batch: RecordBatch,
114}
115
116impl ChangelogBatch {
117    /// Builds a changelog batch and validates that row operations match the
118    /// payload row count and update-pair adjacency rules.
119    pub fn try_new(ops: Vec<ChangeOp>, batch: RecordBatch) -> Result<Self> {
120        validate_change_ops(&ops, batch.num_rows())?;
121        Ok(Self {
122            ops: Arc::from(ops),
123            batch,
124        })
125    }
126
127    /// Wraps an append-only Arrow batch as insert-only changelog data.
128    #[must_use]
129    pub fn insert_only(batch: RecordBatch) -> Self {
130        Self {
131            ops: vec![ChangeOp::Insert; batch.num_rows()].into(),
132            batch,
133        }
134    }
135
136    /// Returns the row-level changelog operations.
137    #[must_use]
138    pub fn ops(&self) -> &[ChangeOp] {
139        &self.ops
140    }
141
142    /// Returns the Arrow payload batch.
143    #[must_use]
144    pub fn batch(&self) -> &RecordBatch {
145        &self.batch
146    }
147
148    /// Returns the number of rows in the payload batch.
149    #[must_use]
150    pub fn num_rows(&self) -> usize {
151        self.batch.num_rows()
152    }
153
154    /// Returns true when every row is an insert.
155    #[must_use]
156    pub fn is_insert_only(&self) -> bool {
157        self.ops.iter().all(|op| *op == ChangeOp::Insert)
158    }
159}
160
161impl fmt::Debug for ChangelogBatch {
162    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
163        f.debug_struct("ChangelogBatch")
164            .field("ops", &self.ops)
165            .field("batch", &self.batch)
166            .finish()
167    }
168}
169
170fn validate_change_ops(ops: &[ChangeOp], rows: usize) -> Result<()> {
171    if ops.len() != rows {
172        return Err(DataFusionError::Plan(format!(
173            "changelog batch has {} ops for {rows} rows",
174            ops.len()
175        )));
176    }
177
178    let mut row = 0;
179    while row < ops.len() {
180        match ops[row] {
181            ChangeOp::UpdateDelete => {
182                if ops.get(row + 1) != Some(&ChangeOp::UpdateInsert) {
183                    return Err(DataFusionError::Plan(format!(
184                        "UpdateDelete at row {row} must be immediately followed by UpdateInsert"
185                    )));
186                }
187                row += 2;
188            }
189            ChangeOp::UpdateInsert => {
190                return Err(DataFusionError::Plan(format!(
191                    "UpdateInsert at row {row} must be immediately preceded by UpdateDelete"
192                )));
193            }
194            ChangeOp::Insert | ChangeOp::Delete => row += 1,
195        }
196    }
197
198    Ok(())
199}
200
201/// Per-query counters populated by SQL lowering and native Datum stages.
202#[derive(Debug, Clone, Default)]
203pub struct SqlExecutionMetrics {
204    late_dropped_rows: Arc<AtomicU64>,
205    join_state_rows: Arc<AtomicU64>,
206    join_evicted_rows: Arc<AtomicU64>,
207}
208
209impl SqlExecutionMetrics {
210    /// Rows dropped after the active watermark in windowed or joined streams.
211    #[must_use]
212    pub fn late_dropped_rows(&self) -> u64 {
213        self.late_dropped_rows.load(Ordering::Relaxed)
214    }
215
216    pub(crate) fn window_metrics(&self) -> WindowedAggregationMetrics {
217        WindowedAggregationMetrics::new(Arc::clone(&self.late_dropped_rows))
218    }
219
220    /// Returns a snapshot handle for streaming join counters.
221    #[must_use]
222    pub fn streaming_join_metrics(&self) -> StreamingJoinMetrics {
223        self.join_metrics()
224    }
225
226    /// Rows currently retained by streaming join state.
227    #[must_use]
228    pub fn join_state_rows(&self) -> u64 {
229        self.join_state_rows.load(Ordering::Relaxed)
230    }
231
232    /// Rows evicted from streaming join state by watermark progress.
233    #[must_use]
234    pub fn join_evicted_rows(&self) -> u64 {
235        self.join_evicted_rows.load(Ordering::Relaxed)
236    }
237
238    pub(crate) fn join_metrics(&self) -> StreamingJoinMetrics {
239        StreamingJoinMetrics::new(
240            Arc::clone(&self.join_state_rows),
241            Arc::clone(&self.join_evicted_rows),
242            Arc::clone(&self.late_dropped_rows),
243        )
244    }
245}
246
247/// SQL planning context that registers Datum sources and executes supported
248/// DataFusion physical plans on Datum.
249pub struct DatumSqlContext {
250    inner: SessionContext,
251    sinks: Arc<Mutex<BTreeMap<String, RegisteredSqlSink>>>,
252    streaming_join_config: StreamingJoinConfig,
253}
254
255impl DatumSqlContext {
256    /// Creates a SQL context with one DataFusion target partition and Datum's
257    /// planner placeholder functions registered.
258    #[must_use]
259    pub fn new() -> Self {
260        let config = SessionConfig::new().with_target_partitions(1);
261        let inner = SessionContext::new_with_config(config);
262        register_window_placeholder_udfs(&inner);
263        Self {
264            inner,
265            sinks: Arc::new(Mutex::new(BTreeMap::new())),
266            streaming_join_config: StreamingJoinConfig::default(),
267        }
268    }
269
270    /// Returns this context with a different streaming join policy.
271    #[must_use]
272    pub fn with_streaming_join_config(mut self, config: StreamingJoinConfig) -> Self {
273        self.streaming_join_config = config;
274        self
275    }
276
277    /// Replaces the streaming join policy used by future query lowering.
278    pub fn set_streaming_join_config(&mut self, config: StreamingJoinConfig) {
279        self.streaming_join_config = config;
280    }
281
282    /// Returns the configured streaming join policy.
283    #[must_use]
284    pub const fn streaming_join_config(&self) -> &StreamingJoinConfig {
285        &self.streaming_join_config
286    }
287
288    /// Registers an append-only Datum source as a SQL table.
289    pub fn register_source(
290        &self,
291        name: &str,
292        schema: SchemaRef,
293        source: Source<RecordBatch>,
294    ) -> Result<()> {
295        self.register_source_factory(name, schema, move || source.clone())
296    }
297
298    /// Registers an append-only Datum source and event-time declaration as a
299    /// SQL table.
300    pub fn register_source_with_event_time(
301        &self,
302        name: &str,
303        schema: SchemaRef,
304        source: Source<RecordBatch>,
305        event_time: EventTimeConfig,
306    ) -> Result<()> {
307        self.register_source_factory_with_event_time(
308            name,
309            schema,
310            move || source.clone(),
311            event_time,
312        )
313    }
314
315    /// Registers a factory for append-only Datum sources as a SQL table.
316    pub fn register_source_factory<F>(
317        &self,
318        name: &str,
319        schema: SchemaRef,
320        source_factory: F,
321    ) -> Result<()>
322    where
323        F: Fn() -> Source<RecordBatch> + Send + Sync + 'static,
324    {
325        let table = DatumTable::new(schema, Arc::new(source_factory), None);
326        self.inner.register_table(name, Arc::new(table))?;
327        Ok(())
328    }
329
330    /// Registers a factory for append-only Datum sources with event time.
331    pub fn register_source_factory_with_event_time<F>(
332        &self,
333        name: &str,
334        schema: SchemaRef,
335        source_factory: F,
336        event_time: EventTimeConfig,
337    ) -> Result<()>
338    where
339        F: Fn() -> Source<RecordBatch> + Send + Sync + 'static,
340    {
341        event_time.resolve(schema.as_ref())?;
342        let table = DatumTable::new(schema, Arc::new(source_factory), Some(event_time));
343        self.inner.register_table(name, Arc::new(table))?;
344        Ok(())
345    }
346
347    /// Registers a committable append-only source as a SQL table.
348    pub fn register_committable_source(
349        &self,
350        name: &str,
351        schema: SchemaRef,
352        source: Source<CommittableRecordBatch>,
353    ) -> Result<()> {
354        self.register_committable_source_factory(name, schema, move || source.clone())
355    }
356
357    /// Registers a factory for committable append-only sources as a SQL table.
358    pub fn register_committable_source_factory<F>(
359        &self,
360        name: &str,
361        schema: SchemaRef,
362        source_factory: F,
363    ) -> Result<()>
364    where
365        F: Fn() -> Source<CommittableRecordBatch> + Send + Sync + 'static,
366    {
367        let table = DatumCommittableTable::new(schema, Arc::new(source_factory), None);
368        self.inner.register_table(name, Arc::new(table))?;
369        Ok(())
370    }
371
372    /// Registers a committable append-only source and event-time declaration as
373    /// a SQL table.
374    pub fn register_committable_source_with_event_time(
375        &self,
376        name: &str,
377        schema: SchemaRef,
378        source: Source<CommittableRecordBatch>,
379        event_time: EventTimeConfig,
380    ) -> Result<()> {
381        self.register_committable_source_factory_with_event_time(
382            name,
383            schema,
384            move || source.clone(),
385            event_time,
386        )
387    }
388
389    /// Registers a factory for committable append-only sources with event time.
390    pub fn register_committable_source_factory_with_event_time<F>(
391        &self,
392        name: &str,
393        schema: SchemaRef,
394        source_factory: F,
395        event_time: EventTimeConfig,
396    ) -> Result<()>
397    where
398        F: Fn() -> Source<CommittableRecordBatch> + Send + Sync + 'static,
399    {
400        event_time.resolve(schema.as_ref())?;
401        let table = DatumCommittableTable::new(schema, Arc::new(source_factory), Some(event_time));
402        self.inner.register_table(name, Arc::new(table))?;
403        Ok(())
404    }
405
406    /// Registers an updating source as a SQL table.
407    pub fn register_changelog_source(
408        &self,
409        name: &str,
410        schema: SchemaRef,
411        source: Source<ChangelogBatch>,
412    ) -> Result<()> {
413        self.register_changelog_source_factory(name, schema, move || source.clone())
414    }
415
416    /// Registers an updating source and event-time declaration as a SQL table.
417    pub fn register_changelog_source_with_event_time(
418        &self,
419        name: &str,
420        schema: SchemaRef,
421        source: Source<ChangelogBatch>,
422        event_time: EventTimeConfig,
423    ) -> Result<()> {
424        self.register_changelog_source_factory_with_event_time(
425            name,
426            schema,
427            move || source.clone(),
428            event_time,
429        )
430    }
431
432    /// Registers a factory for updating sources as a SQL table.
433    pub fn register_changelog_source_factory<F>(
434        &self,
435        name: &str,
436        schema: SchemaRef,
437        source_factory: F,
438    ) -> Result<()>
439    where
440        F: Fn() -> Source<ChangelogBatch> + Send + Sync + 'static,
441    {
442        let table = DatumChangelogTable::new(schema, Arc::new(source_factory), None);
443        self.inner.register_table(name, Arc::new(table))?;
444        Ok(())
445    }
446
447    /// Registers a factory for updating sources with event time.
448    pub fn register_changelog_source_factory_with_event_time<F>(
449        &self,
450        name: &str,
451        schema: SchemaRef,
452        source_factory: F,
453        event_time: EventTimeConfig,
454    ) -> Result<()>
455    where
456        F: Fn() -> Source<ChangelogBatch> + Send + Sync + 'static,
457    {
458        event_time.resolve(schema.as_ref())?;
459        let table = DatumChangelogTable::new(schema, Arc::new(source_factory), Some(event_time));
460        self.inner.register_table(name, Arc::new(table))?;
461        Ok(())
462    }
463
464    #[cfg(feature = "mq")]
465    /// Registers a Kafka JSON topic as a committable SQL table.
466    pub fn register_mq_topic(
467        &self,
468        name: &str,
469        settings: datum_mq::KafkaConsumerSettings,
470        topic: impl Into<String>,
471        format: connect::mq::JsonRowFormat,
472    ) -> Result<()> {
473        let topic = topic.into();
474        self.register_committable_source_factory(name, format.schema(), move || {
475            connect::mq::kafka_json_committable_source_uncontrolled(
476                settings.clone(),
477                topic.clone(),
478                format.clone(),
479            )
480        })
481    }
482
483    #[cfg(feature = "mq")]
484    /// Registers a Kafka JSON topic and event-time declaration as a
485    /// committable SQL table.
486    pub fn register_mq_topic_with_event_time(
487        &self,
488        name: &str,
489        settings: datum_mq::KafkaConsumerSettings,
490        topic: impl Into<String>,
491        format: connect::mq::JsonRowFormat,
492        event_time: EventTimeConfig,
493    ) -> Result<()> {
494        let topic = topic.into();
495        self.register_committable_source_factory_with_event_time(
496            name,
497            format.schema(),
498            move || {
499                connect::mq::kafka_json_committable_source_uncontrolled(
500                    settings.clone(),
501                    topic.clone(),
502                    format.clone(),
503                )
504            },
505            event_time,
506        )
507    }
508
509    #[cfg(feature = "mq")]
510    /// Registers a Kafka JSON producer as an append-only SQL sink.
511    pub fn register_mq_json_sink(
512        &self,
513        name: &str,
514        settings: datum_mq::KafkaProducerSettings,
515        topic: impl Into<String>,
516        format: connect::mq::JsonRowFormat,
517    ) -> Result<()> {
518        self.register_sql_sink(
519            name,
520            RegisteredSqlSink::kafka_json(settings, topic.into(), format),
521        )
522    }
523
524    #[cfg(feature = "cdc")]
525    /// Registers a CDC source factory as an updating SQL table.
526    pub fn register_cdc_source<F>(
527        &self,
528        name: &str,
529        schema: SchemaRef,
530        source_factory: F,
531    ) -> Result<()>
532    where
533        F: Fn() -> Source<datum_cdc::ChangeEvent, datum_cdc::CdcHandle> + Send + Sync + 'static,
534    {
535        self.register_changelog_source_factory(name, Arc::clone(&schema), move || {
536            connect::cdc::cdc_changelog_batch_source_uncontrolled(
537                source_factory(),
538                Arc::clone(&schema),
539            )
540        })
541    }
542
543    /// Returns DataFusion's physical plan for SQL text without lowering it.
544    pub async fn physical_plan(&self, sql: &str) -> Result<Arc<dyn ExecutionPlan>> {
545        self.inner.sql(sql).await?.create_physical_plan().await
546    }
547
548    /// Executes a bounded append-only query and collects the output batches.
549    pub async fn execute(&self, sql: &str) -> Result<Vec<RecordBatch>> {
550        let plan = self.physical_plan(sql).await?;
551        let source = lower_physical_plan(plan.as_ref())?;
552        source
553            .run_with(Sink::collect())
554            .map_err(datafusion_error)?
555            .wait()
556            .map_err(datafusion_error)
557    }
558
559    /// Lowers a continuous SQL query into a Datum source of data/control events.
560    pub async fn streaming_source(&self, sql: &str) -> Result<Source<SqlEvent<RecordBatch>>> {
561        self.streaming_source_with_metrics(sql)
562            .await
563            .map(|(source, _metrics)| source)
564    }
565
566    /// Lowers a continuous SQL query and returns shared execution metrics.
567    pub async fn streaming_source_with_metrics(
568        &self,
569        sql: &str,
570    ) -> Result<(Source<SqlEvent<RecordBatch>>, SqlExecutionMetrics)> {
571        let plan = self.physical_plan(sql).await?;
572        let metrics = SqlExecutionMetrics::default();
573        let source = lower_physical_plan_streaming_with_metrics_and_config(
574            plan.as_ref(),
575            &metrics,
576            &self.streaming_join_config,
577        )?;
578        Ok((source, metrics))
579    }
580
581    /// Materializes a continuous SQL query, or an `INSERT INTO` sink query.
582    pub async fn execute_streaming(&self, sql: &str) -> Result<ContinuousQueryHandle> {
583        if looks_like_insert_into(sql) {
584            return self.execute_insert_into(sql).await;
585        }
586        let source = self.streaming_source(sql).await?;
587        let (kill_switch, completion) = source
588            .via_mat(KillSwitches::single(), Keep::right)
589            .to_mat(Sink::ignore(), Keep::both)
590            .run()
591            .map_err(datafusion_error)?;
592        Ok(ContinuousQueryHandle::new(kill_switch, completion))
593    }
594
595    /// Parses and materializes `INSERT INTO <sink> SELECT ...`.
596    pub async fn execute_insert_into(&self, sql: &str) -> Result<ContinuousQueryHandle> {
597        let insert = parse_insert_into(sql)?;
598        self.select_into(&insert.select_sql, &insert.sink_name)
599            .await
600    }
601
602    /// Materializes a `SELECT` body into a registered sink.
603    pub async fn select_into(
604        &self,
605        select_sql: &str,
606        sink_name: &str,
607    ) -> Result<ContinuousQueryHandle> {
608        let sink = self.lookup_sink(sink_name)?;
609        let plan = self.physical_plan(select_sql).await?;
610        self.run_plan_into_sink(plan.as_ref(), sink)
611    }
612
613    /// Registers an append-only sink callback.
614    pub fn register_append_sink<F>(&self, name: &str, writer: F) -> Result<()>
615    where
616        F: Fn(RecordBatch) -> StreamResult<()> + Send + Sync + 'static,
617    {
618        self.register_sql_sink(name, RegisteredSqlSink::append_only(writer))
619    }
620
621    /// Registers a sink callback that can consume both append-only and updating
622    /// streams.
623    pub fn register_changelog_sink<F, G>(
624        &self,
625        name: &str,
626        append_writer: F,
627        changelog_writer: G,
628    ) -> Result<()>
629    where
630        F: Fn(RecordBatch) -> StreamResult<()> + Send + Sync + 'static,
631        G: Fn(ChangelogBatch) -> StreamResult<()> + Send + Sync + 'static,
632    {
633        self.register_sql_sink(
634            name,
635            RegisteredSqlSink::changelog(append_writer, changelog_writer),
636        )
637    }
638
639    fn register_sql_sink(&self, name: &str, sink: RegisteredSqlSink) -> Result<()> {
640        self.sinks
641            .lock()
642            .map_err(|_| sink_registry_error("datum-sql sink registry lock poisoned"))?
643            .insert(name.to_owned(), sink);
644        Ok(())
645    }
646
647    fn lookup_sink(&self, name: &str) -> Result<RegisteredSqlSink> {
648        self.sinks
649            .lock()
650            .map_err(|_| sink_registry_error("datum-sql sink registry lock poisoned"))?
651            .get(name)
652            .cloned()
653            .ok_or_else(|| {
654                DataFusionError::Plan(format!("datum-sql sink '{name}' is not registered"))
655            })
656    }
657
658    fn run_plan_into_sink(
659        &self,
660        plan: &dyn ExecutionPlan,
661        sink: RegisteredSqlSink,
662    ) -> Result<ContinuousQueryHandle> {
663        if plan_contains_changelog_scan(plan) {
664            if !sink.accepts_changelog() {
665                return Err(DataFusionError::Plan(
666                    "datum-sql append-only sink cannot consume an updating stream; register a changelog-aware sink to accept retractions".into(),
667                ));
668            }
669            let source = lower_changelog_physical_plan(plan)?;
670            let (kill_switch, completion) = source
671                .via_mat(KillSwitches::single(), Keep::right)
672                .to_mat(
673                    Sink::foreach_result(move |changes| sink.write_changelog(changes)),
674                    Keep::both,
675                )
676                .run()
677                .map_err(datafusion_error)?;
678            return Ok(ContinuousQueryHandle::new(kill_switch, completion));
679        }
680
681        let source = lower_append_sink_plan(plan)?.into_committable();
682        let (kill_switch, completion) = source
683            .via_mat(KillSwitches::single(), Keep::right)
684            .to_mat(
685                Sink::foreach_result(move |batch: CommittableRecordBatch| {
686                    sink.write_append(batch.batch().clone())?;
687                    batch.commit()
688                }),
689                Keep::both,
690            )
691            .run()
692            .map_err(datafusion_error)?;
693        Ok(ContinuousQueryHandle::new(kill_switch, completion))
694    }
695}
696
697impl Default for DatumSqlContext {
698    fn default() -> Self {
699        Self::new()
700    }
701}
702
703impl fmt::Debug for DatumSqlContext {
704    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
705        f.debug_struct("DatumSqlContext").finish_non_exhaustive()
706    }
707}
708
709fn register_window_placeholder_udfs(context: &SessionContext) {
710    context.register_udf(ScalarUDF::from(WindowPlaceholderUdf::new("tumble", 2)));
711    context.register_udf(ScalarUDF::from(WindowPlaceholderUdf::new("hop", 3)));
712}
713
714#[derive(Debug, Clone, PartialEq, Eq, Hash)]
715struct WindowPlaceholderUdf {
716    name: &'static str,
717    arity: usize,
718    signature: Signature,
719}
720
721impl WindowPlaceholderUdf {
722    fn new(name: &'static str, arity: usize) -> Self {
723        Self {
724            name,
725            arity,
726            signature: Signature::variadic_any(Volatility::Immutable),
727        }
728    }
729
730    fn timestamp_arg<'a>(&self, args: &'a [FieldRef]) -> Result<&'a Field> {
731        if args.len() != self.arity {
732            return Err(DataFusionError::Plan(format!(
733                "{} expects {} arguments, found {}",
734                self.name,
735                self.arity,
736                args.len()
737            )));
738        }
739        let timestamp_index = self.arity - 1;
740        let field = args[timestamp_index].as_ref();
741        if !matches!(field.data_type(), DataType::Timestamp(_, _)) {
742            return Err(DataFusionError::Plan(format!(
743                "{} expects its final argument to be a timestamp, found {:?}",
744                self.name,
745                field.data_type()
746            )));
747        }
748        Ok(field)
749    }
750}
751
752impl ScalarUDFImpl for WindowPlaceholderUdf {
753    fn name(&self) -> &str {
754        self.name
755    }
756
757    fn signature(&self) -> &Signature {
758        &self.signature
759    }
760
761    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
762        if arg_types.len() != self.arity {
763            return Err(DataFusionError::Plan(format!(
764                "{} expects {} arguments, found {}",
765                self.name,
766                self.arity,
767                arg_types.len()
768            )));
769        }
770        match &arg_types[self.arity - 1] {
771            DataType::Timestamp(unit, timezone) => Ok(DataType::Timestamp(*unit, timezone.clone())),
772            other => Err(DataFusionError::Plan(format!(
773                "{} expects its final argument to be a timestamp, found {other:?}",
774                self.name
775            ))),
776        }
777    }
778
779    fn return_field_from_args(&self, args: ReturnFieldArgs<'_>) -> Result<FieldRef> {
780        let timestamp = self.timestamp_arg(args.arg_fields)?;
781        Ok(Arc::new(Field::new(
782            self.name,
783            timestamp.data_type().clone(),
784            timestamp.is_nullable(),
785        )))
786    }
787
788    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
789        args.args.get(self.arity - 1).cloned().ok_or_else(|| {
790            DataFusionError::Plan(format!("{} missing timestamp argument", self.name))
791        })
792    }
793}
794
795#[derive(Clone)]
796struct DatumTable {
797    schema: SchemaRef,
798    source_factory: SourceFactory,
799    event_time: Option<EventTimeConfig>,
800}
801
802impl DatumTable {
803    fn new(
804        schema: SchemaRef,
805        source_factory: SourceFactory,
806        event_time: Option<EventTimeConfig>,
807    ) -> Self {
808        Self {
809            schema,
810            source_factory,
811            event_time,
812        }
813    }
814}
815
816impl fmt::Debug for DatumTable {
817    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
818        f.debug_struct("DatumTable")
819            .field("schema", &self.schema)
820            .finish_non_exhaustive()
821    }
822}
823
824#[derive(Clone)]
825struct DatumCommittableTable {
826    schema: SchemaRef,
827    source_factory: CommittableSourceFactory,
828    event_time: Option<EventTimeConfig>,
829}
830
831impl DatumCommittableTable {
832    fn new(
833        schema: SchemaRef,
834        source_factory: CommittableSourceFactory,
835        event_time: Option<EventTimeConfig>,
836    ) -> Self {
837        Self {
838            schema,
839            source_factory,
840            event_time,
841        }
842    }
843}
844
845impl fmt::Debug for DatumCommittableTable {
846    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
847        f.debug_struct("DatumCommittableTable")
848            .field("schema", &self.schema)
849            .finish_non_exhaustive()
850    }
851}
852
853#[derive(Clone)]
854struct DatumChangelogTable {
855    schema: SchemaRef,
856    source_factory: ChangelogSourceFactory,
857    event_time: Option<EventTimeConfig>,
858}
859
860impl DatumChangelogTable {
861    fn new(
862        schema: SchemaRef,
863        source_factory: ChangelogSourceFactory,
864        event_time: Option<EventTimeConfig>,
865    ) -> Self {
866        Self {
867            schema,
868            source_factory,
869            event_time,
870        }
871    }
872}
873
874impl fmt::Debug for DatumChangelogTable {
875    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
876        f.debug_struct("DatumChangelogTable")
877            .field("schema", &self.schema)
878            .finish_non_exhaustive()
879    }
880}
881
882#[async_trait]
883impl TableProvider for DatumTable {
884    fn schema(&self) -> SchemaRef {
885        Arc::clone(&self.schema)
886    }
887
888    fn table_type(&self) -> TableType {
889        TableType::Base
890    }
891
892    async fn scan(
893        &self,
894        _state: &dyn Session,
895        projection: Option<&Vec<usize>>,
896        filters: &[Expr],
897        limit: Option<usize>,
898    ) -> Result<Arc<dyn ExecutionPlan>> {
899        if !filters.is_empty() {
900            return Err(DataFusionError::Internal(
901                "datum-sql scan received pushed filters despite reporting them unsupported".into(),
902            ));
903        }
904        Ok(Arc::new(DatumScanExec::try_new(
905            Arc::clone(&self.schema),
906            Arc::clone(&self.source_factory),
907            projection.cloned(),
908            limit,
909            self.event_time.clone(),
910        )?))
911    }
912
913    fn supports_filters_pushdown(
914        &self,
915        filters: &[&Expr],
916    ) -> Result<Vec<TableProviderFilterPushDown>> {
917        Ok(vec![
918            TableProviderFilterPushDown::Unsupported;
919            filters.len()
920        ])
921    }
922}
923
924#[async_trait]
925impl TableProvider for DatumCommittableTable {
926    fn schema(&self) -> SchemaRef {
927        Arc::clone(&self.schema)
928    }
929
930    fn table_type(&self) -> TableType {
931        TableType::Base
932    }
933
934    async fn scan(
935        &self,
936        _state: &dyn Session,
937        projection: Option<&Vec<usize>>,
938        filters: &[Expr],
939        limit: Option<usize>,
940    ) -> Result<Arc<dyn ExecutionPlan>> {
941        if !filters.is_empty() {
942            return Err(DataFusionError::Internal(
943                "datum-sql committable scan received pushed filters despite reporting them unsupported"
944                    .into(),
945            ));
946        }
947        Ok(Arc::new(DatumCommittableScanExec::try_new(
948            Arc::clone(&self.schema),
949            Arc::clone(&self.source_factory),
950            projection.cloned(),
951            limit,
952            self.event_time.clone(),
953        )?))
954    }
955
956    fn supports_filters_pushdown(
957        &self,
958        filters: &[&Expr],
959    ) -> Result<Vec<TableProviderFilterPushDown>> {
960        Ok(vec![
961            TableProviderFilterPushDown::Unsupported;
962            filters.len()
963        ])
964    }
965}
966
967#[derive(Clone)]
968struct DatumCommittableScanExec {
969    input_schema: SchemaRef,
970    schema: SchemaRef,
971    source_factory: CommittableSourceFactory,
972    projection: Option<Vec<usize>>,
973    limit: Option<usize>,
974    event_time: Option<EventTimeConfig>,
975    cache: Arc<PlanProperties>,
976}
977
978impl DatumCommittableScanExec {
979    fn try_new(
980        input_schema: SchemaRef,
981        source_factory: CommittableSourceFactory,
982        projection: Option<Vec<usize>>,
983        limit: Option<usize>,
984        event_time: Option<EventTimeConfig>,
985    ) -> Result<Self> {
986        let schema = project_schema(&input_schema, projection.as_ref())?;
987        if let Some(event_time) = &event_time {
988            event_time.resolve(input_schema.as_ref())?;
989        }
990        let cache = Arc::new(DatumScanExec::compute_properties(Arc::clone(&schema)));
991        Ok(Self {
992            input_schema,
993            schema,
994            source_factory,
995            projection,
996            limit,
997            event_time,
998            cache,
999        })
1000    }
1001
1002    fn source(&self) -> Source<RecordBatch> {
1003        self.committable_source().map(|batch| batch.batch().clone())
1004    }
1005
1006    fn committable_source(&self) -> Source<CommittableRecordBatch> {
1007        let mut source = (self.source_factory)();
1008        if let Some(projection) = self.projection.clone() {
1009            source = source.try_map(move |batch| {
1010                batch
1011                    .try_map_batch(|record_batch| {
1012                        record_batch
1013                            .project(&projection)
1014                            .map_err(DataFusionError::from)
1015                    })
1016                    .map_err(stream_error)
1017            });
1018        }
1019        if let Some(limit) = self.limit {
1020            source = apply_committable_row_window(source, 0, Some(limit));
1021        }
1022        source
1023    }
1024
1025    fn event_source(&self) -> Result<Source<SqlEvent<RecordBatch>>> {
1026        let raw_source = (self.source_factory)();
1027        let source = match &self.event_time {
1028            Some(event_time) => {
1029                let resolved = event_time.resolve(self.input_schema.as_ref())?;
1030                time::assign_resolved_event_time_watermarks(raw_source, resolved)
1031            }
1032            None => data_events(raw_source),
1033        };
1034        let mut source = map_sql_event_data(source, |batch: CommittableRecordBatch| {
1035            Ok(batch.batch().clone())
1036        });
1037        if let Some(projection) = self.projection.clone() {
1038            source = map_sql_event_data(source, move |batch| {
1039                batch.project(&projection).map_err(stream_error)
1040            });
1041        }
1042        if let Some(limit) = self.limit {
1043            source = apply_event_row_window(source, 0, Some(limit));
1044        }
1045        Ok(source)
1046    }
1047}
1048
1049impl fmt::Debug for DatumCommittableScanExec {
1050    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1051        f.debug_struct("DatumCommittableScanExec")
1052            .field("schema", &self.schema)
1053            .field("projection", &self.projection)
1054            .field("limit", &self.limit)
1055            .field("event_time", &self.event_time)
1056            .finish_non_exhaustive()
1057    }
1058}
1059
1060impl DisplayAs for DatumCommittableScanExec {
1061    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1062        match t {
1063            DisplayFormatType::Default | DisplayFormatType::Verbose => write!(
1064                f,
1065                "DatumCommittableScanExec: projection={:?}, limit={:?}",
1066                self.projection, self.limit
1067            ),
1068            DisplayFormatType::TreeRender => write!(f, "DatumCommittableScanExec"),
1069        }
1070    }
1071}
1072
1073impl ExecutionPlan for DatumCommittableScanExec {
1074    fn name(&self) -> &'static str {
1075        "DatumCommittableScanExec"
1076    }
1077
1078    fn properties(&self) -> &Arc<PlanProperties> {
1079        &self.cache
1080    }
1081
1082    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
1083        vec![]
1084    }
1085
1086    fn with_new_children(
1087        self: Arc<Self>,
1088        children: Vec<Arc<dyn ExecutionPlan>>,
1089    ) -> Result<Arc<dyn ExecutionPlan>> {
1090        if children.is_empty() {
1091            Ok(self)
1092        } else {
1093            Err(DataFusionError::Internal(
1094                "DatumCommittableScanExec does not accept child plans".into(),
1095            ))
1096        }
1097    }
1098
1099    fn execute(
1100        &self,
1101        partition: usize,
1102        _context: Arc<TaskContext>,
1103    ) -> Result<SendableRecordBatchStream> {
1104        assert_or_internal_err!(
1105            partition == 0,
1106            "DatumCommittableScanExec invalid partition {partition}; only partition 0 exists"
1107        );
1108        let batches = self.source().run_collect().map_err(datafusion_error)?;
1109        Ok(Box::pin(MemoryStream::try_new(
1110            batches,
1111            Arc::clone(&self.schema),
1112            None,
1113        )?))
1114    }
1115
1116    fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
1117        if let Some(partition) = partition {
1118            assert_or_internal_err!(
1119                partition == 0,
1120                "DatumCommittableScanExec invalid partition {partition}; only partition 0 exists"
1121            );
1122        }
1123
1124        let mut stats = Statistics::default();
1125        for _ in self.schema.fields() {
1126            stats = stats.add_column_statistics(ColumnStatistics {
1127                null_count: Precision::Absent,
1128                distinct_count: Precision::Absent,
1129                min_value: Precision::<ScalarValue>::Absent,
1130                max_value: Precision::<ScalarValue>::Absent,
1131                sum_value: Precision::<ScalarValue>::Absent,
1132                byte_size: Precision::Absent,
1133            });
1134        }
1135        Ok(Arc::new(stats))
1136    }
1137}
1138
1139#[async_trait]
1140impl TableProvider for DatumChangelogTable {
1141    fn schema(&self) -> SchemaRef {
1142        Arc::clone(&self.schema)
1143    }
1144
1145    fn table_type(&self) -> TableType {
1146        TableType::Base
1147    }
1148
1149    async fn scan(
1150        &self,
1151        _state: &dyn Session,
1152        projection: Option<&Vec<usize>>,
1153        filters: &[Expr],
1154        limit: Option<usize>,
1155    ) -> Result<Arc<dyn ExecutionPlan>> {
1156        if !filters.is_empty() {
1157            return Err(DataFusionError::Internal(
1158                "datum-sql changelog scan received pushed filters despite reporting them unsupported"
1159                    .into(),
1160            ));
1161        }
1162        Ok(Arc::new(DatumChangelogScanExec::try_new(
1163            Arc::clone(&self.schema),
1164            Arc::clone(&self.source_factory),
1165            projection.cloned(),
1166            limit,
1167            self.event_time.clone(),
1168        )?))
1169    }
1170
1171    fn supports_filters_pushdown(
1172        &self,
1173        filters: &[&Expr],
1174    ) -> Result<Vec<TableProviderFilterPushDown>> {
1175        Ok(vec![
1176            TableProviderFilterPushDown::Unsupported;
1177            filters.len()
1178        ])
1179    }
1180}
1181
1182#[derive(Clone)]
1183struct DatumScanExec {
1184    input_schema: SchemaRef,
1185    schema: SchemaRef,
1186    source_factory: SourceFactory,
1187    projection: Option<Vec<usize>>,
1188    limit: Option<usize>,
1189    event_time: Option<EventTimeConfig>,
1190    cache: Arc<PlanProperties>,
1191}
1192
1193impl DatumScanExec {
1194    fn try_new(
1195        input_schema: SchemaRef,
1196        source_factory: SourceFactory,
1197        projection: Option<Vec<usize>>,
1198        limit: Option<usize>,
1199        event_time: Option<EventTimeConfig>,
1200    ) -> Result<Self> {
1201        let schema = project_schema(&input_schema, projection.as_ref())?;
1202        if let Some(event_time) = &event_time {
1203            event_time.resolve(input_schema.as_ref())?;
1204        }
1205        let cache = Arc::new(Self::compute_properties(Arc::clone(&schema)));
1206        Ok(Self {
1207            input_schema,
1208            schema,
1209            source_factory,
1210            projection,
1211            limit,
1212            event_time,
1213            cache,
1214        })
1215    }
1216
1217    fn compute_properties(schema: SchemaRef) -> PlanProperties {
1218        PlanProperties::new(
1219            EquivalenceProperties::new(schema),
1220            Partitioning::UnknownPartitioning(1),
1221            EmissionType::Incremental,
1222            Boundedness::Bounded,
1223        )
1224        .with_scheduling_type(SchedulingType::Cooperative)
1225    }
1226
1227    fn source(&self) -> Source<RecordBatch> {
1228        let mut source = (self.source_factory)();
1229        if let Some(projection) = self.projection.clone() {
1230            source = source.try_map(move |batch| batch.project(&projection).map_err(stream_error));
1231        }
1232        if let Some(limit) = self.limit {
1233            source = apply_row_window(source, 0, Some(limit));
1234        }
1235        source
1236    }
1237
1238    fn event_source(&self) -> Result<Source<SqlEvent<RecordBatch>>> {
1239        let raw_source = (self.source_factory)();
1240        let mut source = match &self.event_time {
1241            Some(event_time) => {
1242                let resolved = event_time.resolve(self.input_schema.as_ref())?;
1243                time::assign_resolved_event_time_watermarks(raw_source, resolved)
1244            }
1245            None => data_events(raw_source),
1246        };
1247        if let Some(projection) = self.projection.clone() {
1248            source = map_sql_event_data(source, move |batch| {
1249                batch.project(&projection).map_err(stream_error)
1250            });
1251        }
1252        if let Some(limit) = self.limit {
1253            source = apply_event_row_window(source, 0, Some(limit));
1254        }
1255        Ok(source)
1256    }
1257}
1258
1259#[derive(Clone)]
1260struct DatumChangelogScanExec {
1261    input_schema: SchemaRef,
1262    schema: SchemaRef,
1263    source_factory: ChangelogSourceFactory,
1264    projection: Option<Vec<usize>>,
1265    limit: Option<usize>,
1266    event_time: Option<EventTimeConfig>,
1267    cache: Arc<PlanProperties>,
1268}
1269
1270impl DatumChangelogScanExec {
1271    fn try_new(
1272        input_schema: SchemaRef,
1273        source_factory: ChangelogSourceFactory,
1274        projection: Option<Vec<usize>>,
1275        limit: Option<usize>,
1276        event_time: Option<EventTimeConfig>,
1277    ) -> Result<Self> {
1278        let schema = project_schema(&input_schema, projection.as_ref())?;
1279        if let Some(event_time) = &event_time {
1280            event_time.resolve(input_schema.as_ref())?;
1281        }
1282        let cache = Arc::new(DatumScanExec::compute_properties(Arc::clone(&schema)));
1283        Ok(Self {
1284            input_schema,
1285            schema,
1286            source_factory,
1287            projection,
1288            limit,
1289            event_time,
1290            cache,
1291        })
1292    }
1293
1294    fn source(&self) -> Source<ChangelogBatch> {
1295        let mut source = (self.source_factory)();
1296        if let Some(projection) = self.projection.clone() {
1297            source = source.try_map(move |changes| {
1298                project_changelog_batch(&changes, &projection).map_err(stream_error)
1299            });
1300        }
1301        if let Some(limit) = self.limit {
1302            source = apply_changelog_row_window(source, 0, Some(limit));
1303        }
1304        source
1305    }
1306
1307    fn event_source(&self) -> Result<Source<SqlEvent<ChangelogBatch>>> {
1308        let raw_source = (self.source_factory)();
1309        let mut source = match &self.event_time {
1310            Some(event_time) => {
1311                let resolved = event_time.resolve(self.input_schema.as_ref())?;
1312                time::assign_resolved_event_time_watermarks(raw_source, resolved)
1313            }
1314            None => data_events(raw_source),
1315        };
1316        if let Some(projection) = self.projection.clone() {
1317            source = map_sql_event_data(source, move |changes| {
1318                project_changelog_batch(&changes, &projection).map_err(stream_error)
1319            });
1320        }
1321        if let Some(limit) = self.limit {
1322            source = apply_changelog_event_row_window(source, 0, Some(limit));
1323        }
1324        Ok(source)
1325    }
1326}
1327
1328impl fmt::Debug for DatumChangelogScanExec {
1329    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1330        f.debug_struct("DatumChangelogScanExec")
1331            .field("schema", &self.schema)
1332            .field("projection", &self.projection)
1333            .field("limit", &self.limit)
1334            .field("event_time", &self.event_time)
1335            .finish_non_exhaustive()
1336    }
1337}
1338
1339impl DisplayAs for DatumChangelogScanExec {
1340    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1341        match t {
1342            DisplayFormatType::Default | DisplayFormatType::Verbose => write!(
1343                f,
1344                "DatumChangelogScanExec: projection={:?}, limit={:?}",
1345                self.projection, self.limit
1346            ),
1347            DisplayFormatType::TreeRender => write!(f, "DatumChangelogScanExec"),
1348        }
1349    }
1350}
1351
1352impl ExecutionPlan for DatumChangelogScanExec {
1353    fn name(&self) -> &'static str {
1354        "DatumChangelogScanExec"
1355    }
1356
1357    fn properties(&self) -> &Arc<PlanProperties> {
1358        &self.cache
1359    }
1360
1361    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
1362        vec![]
1363    }
1364
1365    fn with_new_children(
1366        self: Arc<Self>,
1367        children: Vec<Arc<dyn ExecutionPlan>>,
1368    ) -> Result<Arc<dyn ExecutionPlan>> {
1369        if children.is_empty() {
1370            Ok(self)
1371        } else {
1372            Err(DataFusionError::Internal(
1373                "DatumChangelogScanExec does not accept child plans".into(),
1374            ))
1375        }
1376    }
1377
1378    fn execute(
1379        &self,
1380        _partition: usize,
1381        _context: Arc<TaskContext>,
1382    ) -> Result<SendableRecordBatchStream> {
1383        Err(updating_scan_error())
1384    }
1385
1386    fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
1387        if let Some(partition) = partition {
1388            assert_or_internal_err!(
1389                partition == 0,
1390                "DatumChangelogScanExec invalid partition {partition}; only partition 0 exists"
1391            );
1392        }
1393
1394        let mut stats = Statistics::default();
1395        for _ in self.schema.fields() {
1396            stats = stats.add_column_statistics(ColumnStatistics {
1397                null_count: Precision::Absent,
1398                distinct_count: Precision::Absent,
1399                min_value: Precision::<ScalarValue>::Absent,
1400                max_value: Precision::<ScalarValue>::Absent,
1401                sum_value: Precision::<ScalarValue>::Absent,
1402                byte_size: Precision::Absent,
1403            });
1404        }
1405        Ok(Arc::new(stats))
1406    }
1407}
1408
1409impl fmt::Debug for DatumScanExec {
1410    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1411        f.debug_struct("DatumScanExec")
1412            .field("schema", &self.schema)
1413            .field("projection", &self.projection)
1414            .field("limit", &self.limit)
1415            .field("event_time", &self.event_time)
1416            .finish_non_exhaustive()
1417    }
1418}
1419
1420impl DisplayAs for DatumScanExec {
1421    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1422        match t {
1423            DisplayFormatType::Default | DisplayFormatType::Verbose => write!(
1424                f,
1425                "DatumScanExec: projection={:?}, limit={:?}",
1426                self.projection, self.limit
1427            ),
1428            DisplayFormatType::TreeRender => write!(f, "DatumScanExec"),
1429        }
1430    }
1431}
1432
1433impl ExecutionPlan for DatumScanExec {
1434    fn name(&self) -> &'static str {
1435        "DatumScanExec"
1436    }
1437
1438    fn properties(&self) -> &Arc<PlanProperties> {
1439        &self.cache
1440    }
1441
1442    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
1443        vec![]
1444    }
1445
1446    fn with_new_children(
1447        self: Arc<Self>,
1448        children: Vec<Arc<dyn ExecutionPlan>>,
1449    ) -> Result<Arc<dyn ExecutionPlan>> {
1450        if children.is_empty() {
1451            Ok(self)
1452        } else {
1453            Err(DataFusionError::Internal(
1454                "DatumScanExec does not accept child plans".into(),
1455            ))
1456        }
1457    }
1458
1459    fn execute(
1460        &self,
1461        partition: usize,
1462        _context: Arc<TaskContext>,
1463    ) -> Result<SendableRecordBatchStream> {
1464        assert_or_internal_err!(
1465            partition == 0,
1466            "DatumScanExec invalid partition {partition}; only partition 0 exists"
1467        );
1468        let batches = self.source().run_collect().map_err(datafusion_error)?;
1469        Ok(Box::pin(MemoryStream::try_new(
1470            batches,
1471            Arc::clone(&self.schema),
1472            None,
1473        )?))
1474    }
1475
1476    fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
1477        if let Some(partition) = partition {
1478            assert_or_internal_err!(
1479                partition == 0,
1480                "DatumScanExec invalid partition {partition}; only partition 0 exists"
1481            );
1482        }
1483
1484        let mut stats = Statistics::default();
1485        for _ in self.schema.fields() {
1486            stats = stats.add_column_statistics(ColumnStatistics {
1487                null_count: Precision::Absent,
1488                distinct_count: Precision::Absent,
1489                min_value: Precision::<ScalarValue>::Absent,
1490                max_value: Precision::<ScalarValue>::Absent,
1491                sum_value: Precision::<ScalarValue>::Absent,
1492                byte_size: Precision::Absent,
1493            });
1494        }
1495        Ok(Arc::new(stats))
1496    }
1497}
1498
1499fn lower_physical_plan(plan: &dyn ExecutionPlan) -> Result<Source<RecordBatch>> {
1500    if let Some(scan) = plan.downcast_ref::<DatumScanExec>() {
1501        return Ok(scan.source());
1502    }
1503
1504    if let Some(scan) = plan.downcast_ref::<DatumCommittableScanExec>() {
1505        return Ok(scan.source());
1506    }
1507
1508    if plan.downcast_ref::<DatumChangelogScanExec>().is_some() {
1509        return Err(updating_scan_error());
1510    }
1511
1512    if let Some(projection) = plan.downcast_ref::<ProjectionExec>() {
1513        let input = lower_physical_plan(projection.input().as_ref())?;
1514        let exprs = projection.expr().to_vec();
1515        let schema = projection.schema();
1516        return Ok(input.try_map(move |batch| {
1517            project_batch(&batch, &exprs, Arc::clone(&schema)).map_err(stream_error)
1518        }));
1519    }
1520
1521    if let Some(filter) = plan.downcast_ref::<FilterExec>() {
1522        let input = lower_physical_plan(filter.input().as_ref())?;
1523        let predicate = Arc::clone(filter.predicate());
1524        let projection = filter.projection().as_ref().map(|p| p.as_ref().to_vec());
1525        return Ok(input.try_map(move |batch| {
1526            filter_batch(&batch, predicate.as_ref(), projection.as_deref()).map_err(stream_error)
1527        }));
1528    }
1529
1530    #[allow(deprecated)]
1531    if let Some(coalesce) = plan.downcast_ref::<CoalesceBatchesExec>() {
1532        return lower_physical_plan(coalesce.input().as_ref());
1533    }
1534
1535    if let Some(coalesce) = plan.downcast_ref::<CoalescePartitionsExec>() {
1536        return lower_physical_plan(coalesce.input().as_ref());
1537    }
1538
1539    if let Some(limit) = plan.downcast_ref::<GlobalLimitExec>() {
1540        let input = lower_physical_plan(limit.input().as_ref())?;
1541        return Ok(apply_row_window(input, limit.skip(), limit.fetch()));
1542    }
1543
1544    if let Some(limit) = plan.downcast_ref::<LocalLimitExec>() {
1545        let input = lower_physical_plan(limit.input().as_ref())?;
1546        return Ok(apply_row_window(input, 0, Some(limit.fetch())));
1547    }
1548
1549    Err(DataFusionError::NotImplemented(format!(
1550        "datum-sql does not lower DataFusion physical node {} yet",
1551        plan.name()
1552    )))
1553}
1554
1555fn lower_physical_plan_streaming_with_metrics_and_config(
1556    plan: &dyn ExecutionPlan,
1557    metrics: &SqlExecutionMetrics,
1558    join_config: &StreamingJoinConfig,
1559) -> Result<Source<SqlEvent<RecordBatch>>> {
1560    if let Some(scan) = plan.downcast_ref::<DatumScanExec>() {
1561        return scan.event_source();
1562    }
1563
1564    if let Some(scan) = plan.downcast_ref::<DatumCommittableScanExec>() {
1565        return scan.event_source();
1566    }
1567
1568    if let Some(projection) = plan.downcast_ref::<ProjectionExec>() {
1569        let input = lower_physical_plan_streaming_with_metrics_and_config(
1570            projection.input().as_ref(),
1571            metrics,
1572            join_config,
1573        )?;
1574        let exprs = projection.expr().to_vec();
1575        let schema = projection.schema();
1576        return Ok(map_sql_event_data(input, move |batch| {
1577            project_batch(&batch, &exprs, Arc::clone(&schema)).map_err(stream_error)
1578        }));
1579    }
1580
1581    if let Some(filter) = plan.downcast_ref::<FilterExec>() {
1582        let input = lower_physical_plan_streaming_with_metrics_and_config(
1583            filter.input().as_ref(),
1584            metrics,
1585            join_config,
1586        )?;
1587        let predicate = Arc::clone(filter.predicate());
1588        let projection = filter.projection().as_ref().map(|p| p.as_ref().to_vec());
1589        return Ok(map_sql_event_data(input, move |batch| {
1590            filter_batch(&batch, predicate.as_ref(), projection.as_deref()).map_err(stream_error)
1591        }));
1592    }
1593
1594    if let Some(join) = plan.downcast_ref::<HashJoinExec>() {
1595        if plan_contains_changelog_scan(join.left().as_ref())
1596            || plan_contains_changelog_scan(join.right().as_ref())
1597        {
1598            return Err(DataFusionError::NotImplemented(
1599                "datum-sql streaming joins do not support changelog inputs yet; append-only inputs are required"
1600                    .into(),
1601            ));
1602        }
1603        let left = lower_physical_plan_streaming_with_metrics_and_config(
1604            join.left().as_ref(),
1605            metrics,
1606            join_config,
1607        )?;
1608        let right = lower_physical_plan_streaming_with_metrics_and_config(
1609            join.right().as_ref(),
1610            metrics,
1611            join_config,
1612        )?;
1613        return join::streaming_hash_join_source(
1614            left,
1615            right,
1616            join,
1617            join_config,
1618            metrics.join_metrics(),
1619        );
1620    }
1621
1622    if let Some(aggregate) = plan.downcast_ref::<AggregateExec>() {
1623        if plan_contains_changelog_scan(aggregate.input().as_ref()) {
1624            let input = lower_changelog_physical_plan_streaming(aggregate.input().as_ref())?;
1625            return window::windowed_changelog_aggregate_source(
1626                input,
1627                aggregate,
1628                metrics.window_metrics(),
1629            );
1630        }
1631        let input = lower_physical_plan_streaming_with_metrics_and_config(
1632            aggregate.input().as_ref(),
1633            metrics,
1634            join_config,
1635        )?;
1636        return window::windowed_aggregate_source(input, aggregate, metrics.window_metrics());
1637    }
1638
1639    #[allow(deprecated)]
1640    if let Some(coalesce) = plan.downcast_ref::<CoalesceBatchesExec>() {
1641        return lower_physical_plan_streaming_with_metrics_and_config(
1642            coalesce.input().as_ref(),
1643            metrics,
1644            join_config,
1645        );
1646    }
1647
1648    if let Some(coalesce) = plan.downcast_ref::<CoalescePartitionsExec>() {
1649        return lower_physical_plan_streaming_with_metrics_and_config(
1650            coalesce.input().as_ref(),
1651            metrics,
1652            join_config,
1653        );
1654    }
1655
1656    if let Some(limit) = plan.downcast_ref::<GlobalLimitExec>() {
1657        let input = lower_physical_plan_streaming_with_metrics_and_config(
1658            limit.input().as_ref(),
1659            metrics,
1660            join_config,
1661        )?;
1662        return Ok(apply_event_row_window(input, limit.skip(), limit.fetch()));
1663    }
1664
1665    if let Some(limit) = plan.downcast_ref::<LocalLimitExec>() {
1666        let input = lower_physical_plan_streaming_with_metrics_and_config(
1667            limit.input().as_ref(),
1668            metrics,
1669            join_config,
1670        )?;
1671        return Ok(apply_event_row_window(input, 0, Some(limit.fetch())));
1672    }
1673
1674    Err(DataFusionError::NotImplemented(format!(
1675        "datum-sql does not lower DataFusion physical node {} yet",
1676        plan.name()
1677    )))
1678}
1679
1680fn lower_changelog_physical_plan_streaming(
1681    plan: &dyn ExecutionPlan,
1682) -> Result<Source<SqlEvent<ChangelogBatch>>> {
1683    if let Some(scan) = plan.downcast_ref::<DatumChangelogScanExec>() {
1684        return scan.event_source();
1685    }
1686
1687    if let Some(projection) = plan.downcast_ref::<ProjectionExec>() {
1688        let input = lower_changelog_physical_plan_streaming(projection.input().as_ref())?;
1689        let exprs = projection.expr().to_vec();
1690        let schema = projection.schema();
1691        return Ok(map_sql_event_data(input, move |changes| {
1692            let batch = project_batch(changes.batch(), &exprs, Arc::clone(&schema))
1693                .map_err(stream_error)?;
1694            ChangelogBatch::try_new(changes.ops().to_vec(), batch).map_err(stream_error)
1695        }));
1696    }
1697
1698    if let Some(filter) = plan.downcast_ref::<FilterExec>() {
1699        let input = lower_changelog_physical_plan_streaming(filter.input().as_ref())?;
1700        let predicate = Arc::clone(filter.predicate());
1701        let projection = filter.projection().as_ref().map(|p| p.as_ref().to_vec());
1702        return Ok(map_sql_event_data(input, move |changes| {
1703            filter_changelog_batch(&changes, predicate.as_ref(), projection.as_deref())
1704                .map_err(stream_error)
1705        }));
1706    }
1707
1708    #[allow(deprecated)]
1709    if let Some(coalesce) = plan.downcast_ref::<CoalesceBatchesExec>() {
1710        return lower_changelog_physical_plan_streaming(coalesce.input().as_ref());
1711    }
1712
1713    if let Some(coalesce) = plan.downcast_ref::<CoalescePartitionsExec>() {
1714        return lower_changelog_physical_plan_streaming(coalesce.input().as_ref());
1715    }
1716
1717    if let Some(limit) = plan.downcast_ref::<GlobalLimitExec>() {
1718        let input = lower_changelog_physical_plan_streaming(limit.input().as_ref())?;
1719        return Ok(apply_changelog_event_row_window(
1720            input,
1721            limit.skip(),
1722            limit.fetch(),
1723        ));
1724    }
1725
1726    if let Some(limit) = plan.downcast_ref::<LocalLimitExec>() {
1727        let input = lower_changelog_physical_plan_streaming(limit.input().as_ref())?;
1728        return Ok(apply_changelog_event_row_window(
1729            input,
1730            0,
1731            Some(limit.fetch()),
1732        ));
1733    }
1734
1735    Err(DataFusionError::NotImplemented(format!(
1736        "datum-sql does not lower updating DataFusion physical node {} yet",
1737        plan.name()
1738    )))
1739}
1740
1741fn lower_changelog_physical_plan(plan: &dyn ExecutionPlan) -> Result<Source<ChangelogBatch>> {
1742    if let Some(scan) = plan.downcast_ref::<DatumChangelogScanExec>() {
1743        return Ok(scan.source());
1744    }
1745
1746    if let Some(projection) = plan.downcast_ref::<ProjectionExec>() {
1747        let input = lower_changelog_physical_plan(projection.input().as_ref())?;
1748        let exprs = projection.expr().to_vec();
1749        let schema = projection.schema();
1750        return Ok(input.try_map(move |changes| {
1751            let batch = project_batch(changes.batch(), &exprs, Arc::clone(&schema))
1752                .map_err(stream_error)?;
1753            ChangelogBatch::try_new(changes.ops().to_vec(), batch).map_err(stream_error)
1754        }));
1755    }
1756
1757    if let Some(filter) = plan.downcast_ref::<FilterExec>() {
1758        let input = lower_changelog_physical_plan(filter.input().as_ref())?;
1759        let predicate = Arc::clone(filter.predicate());
1760        let projection = filter.projection().as_ref().map(|p| p.as_ref().to_vec());
1761        return Ok(input.try_map(move |changes| {
1762            filter_changelog_batch(&changes, predicate.as_ref(), projection.as_deref())
1763                .map_err(stream_error)
1764        }));
1765    }
1766
1767    #[allow(deprecated)]
1768    if let Some(coalesce) = plan.downcast_ref::<CoalesceBatchesExec>() {
1769        return lower_changelog_physical_plan(coalesce.input().as_ref());
1770    }
1771
1772    if let Some(coalesce) = plan.downcast_ref::<CoalescePartitionsExec>() {
1773        return lower_changelog_physical_plan(coalesce.input().as_ref());
1774    }
1775
1776    if let Some(limit) = plan.downcast_ref::<GlobalLimitExec>() {
1777        let input = lower_changelog_physical_plan(limit.input().as_ref())?;
1778        return Ok(apply_changelog_row_window(
1779            input,
1780            limit.skip(),
1781            limit.fetch(),
1782        ));
1783    }
1784
1785    if let Some(limit) = plan.downcast_ref::<LocalLimitExec>() {
1786        let input = lower_changelog_physical_plan(limit.input().as_ref())?;
1787        return Ok(apply_changelog_row_window(input, 0, Some(limit.fetch())));
1788    }
1789
1790    Err(DataFusionError::NotImplemented(format!(
1791        "datum-sql does not lower updating DataFusion physical node {} yet",
1792        plan.name()
1793    )))
1794}
1795
1796fn lower_append_sink_plan(plan: &dyn ExecutionPlan) -> Result<AppendSinkSource> {
1797    if !plan_contains_committable_scan(plan) {
1798        return lower_physical_plan(plan).map(AppendSinkSource::Plain);
1799    }
1800
1801    if let Some(scan) = plan.downcast_ref::<DatumCommittableScanExec>() {
1802        return Ok(AppendSinkSource::Committable(scan.committable_source()));
1803    }
1804
1805    if let Some(projection) = plan.downcast_ref::<ProjectionExec>() {
1806        let input = lower_append_sink_plan(projection.input().as_ref())?.into_committable();
1807        let exprs = projection.expr().to_vec();
1808        let schema = projection.schema();
1809        return Ok(AppendSinkSource::Committable(input.try_map(move |batch| {
1810            batch
1811                .try_map_batch(|record_batch| {
1812                    project_batch(record_batch, &exprs, Arc::clone(&schema))
1813                })
1814                .map_err(stream_error)
1815        })));
1816    }
1817
1818    if let Some(filter) = plan.downcast_ref::<FilterExec>() {
1819        let input = lower_append_sink_plan(filter.input().as_ref())?.into_committable();
1820        let predicate = Arc::clone(filter.predicate());
1821        let projection = filter.projection().as_ref().map(|p| p.as_ref().to_vec());
1822        return Ok(AppendSinkSource::Committable(input.try_map(move |batch| {
1823            batch
1824                .try_map_batch(|record_batch| {
1825                    filter_batch(record_batch, predicate.as_ref(), projection.as_deref())
1826                })
1827                .map_err(stream_error)
1828        })));
1829    }
1830
1831    #[allow(deprecated)]
1832    if let Some(coalesce) = plan.downcast_ref::<CoalesceBatchesExec>() {
1833        return lower_append_sink_plan(coalesce.input().as_ref());
1834    }
1835
1836    if let Some(coalesce) = plan.downcast_ref::<CoalescePartitionsExec>() {
1837        return lower_append_sink_plan(coalesce.input().as_ref());
1838    }
1839
1840    if let Some(limit) = plan.downcast_ref::<GlobalLimitExec>() {
1841        let input = lower_append_sink_plan(limit.input().as_ref())?.into_committable();
1842        return Ok(AppendSinkSource::Committable(apply_committable_row_window(
1843            input,
1844            limit.skip(),
1845            limit.fetch(),
1846        )));
1847    }
1848
1849    if let Some(limit) = plan.downcast_ref::<LocalLimitExec>() {
1850        let input = lower_append_sink_plan(limit.input().as_ref())?.into_committable();
1851        return Ok(AppendSinkSource::Committable(apply_committable_row_window(
1852            input,
1853            0,
1854            Some(limit.fetch()),
1855        )));
1856    }
1857
1858    Err(DataFusionError::NotImplemented(format!(
1859        "datum-sql does not lower committable DataFusion physical node {} yet",
1860        plan.name()
1861    )))
1862}
1863
1864fn project_batch(
1865    batch: &RecordBatch,
1866    exprs: &[ProjectionExpr],
1867    schema: SchemaRef,
1868) -> Result<RecordBatch> {
1869    let columns = exprs
1870        .iter()
1871        .map(|expr| expr.expr.evaluate(batch)?.into_array(batch.num_rows()))
1872        .collect::<Result<Vec<_>>>()?;
1873    RecordBatch::try_new(schema, columns).map_err(DataFusionError::from)
1874}
1875
1876fn filter_batch(
1877    batch: &RecordBatch,
1878    predicate: &dyn PhysicalExpr,
1879    projection: Option<&[usize]>,
1880) -> Result<RecordBatch> {
1881    let array = predicate.evaluate(batch)?.into_array(batch.num_rows())?;
1882    let filter = as_boolean_array(&array)?;
1883    match projection {
1884        Some(projection) => {
1885            let projected = batch.project(projection)?;
1886            filter_record_batch(&projected, filter).map_err(DataFusionError::from)
1887        }
1888        None => filter_record_batch(batch, filter).map_err(DataFusionError::from),
1889    }
1890}
1891
1892fn project_changelog_batch(
1893    changes: &ChangelogBatch,
1894    projection: &[usize],
1895) -> Result<ChangelogBatch> {
1896    let batch = changes.batch().project(projection)?;
1897    ChangelogBatch::try_new(changes.ops().to_vec(), batch)
1898}
1899
1900fn filter_changelog_batch(
1901    changes: &ChangelogBatch,
1902    predicate: &dyn PhysicalExpr,
1903    projection: Option<&[usize]>,
1904) -> Result<ChangelogBatch> {
1905    let array = predicate
1906        .evaluate(changes.batch())?
1907        .into_array(changes.num_rows())?;
1908    let filter = as_boolean_array(&array)?;
1909    let payload = match projection {
1910        Some(projection) => changes.batch().project(projection)?,
1911        None => changes.batch().clone(),
1912    };
1913    let batch = filter_record_batch(&payload, filter)?;
1914    let ops = (0..changes.num_rows())
1915        .filter(|row| filter.is_valid(*row) && filter.value(*row))
1916        .map(|row| changes.ops()[row])
1917        .collect::<Vec<_>>();
1918    ChangelogBatch::try_new(ops, batch)
1919}
1920
1921fn plan_contains_changelog_scan(plan: &dyn ExecutionPlan) -> bool {
1922    plan.downcast_ref::<DatumChangelogScanExec>().is_some()
1923        || plan
1924            .children()
1925            .into_iter()
1926            .any(|child| plan_contains_changelog_scan(child.as_ref()))
1927}
1928
1929fn plan_contains_committable_scan(plan: &dyn ExecutionPlan) -> bool {
1930    plan.downcast_ref::<DatumCommittableScanExec>().is_some()
1931        || plan
1932            .children()
1933            .into_iter()
1934            .any(|child| plan_contains_committable_scan(child.as_ref()))
1935}
1936
1937#[derive(Clone)]
1938struct RowWindow {
1939    skip: usize,
1940    fetch: Option<usize>,
1941}
1942
1943fn apply_row_window(
1944    source: Source<RecordBatch>,
1945    skip: usize,
1946    fetch: Option<usize>,
1947) -> Source<RecordBatch> {
1948    source.stateful_map_concat(RowWindow { skip, fetch }, |window, batch| {
1949        let mut offset = 0;
1950        if window.skip > 0 {
1951            let skipped = window.skip.min(batch.num_rows());
1952            window.skip -= skipped;
1953            offset = skipped;
1954        }
1955
1956        if offset == batch.num_rows() {
1957            return Vec::new();
1958        }
1959
1960        let mut len = batch.num_rows() - offset;
1961        if let Some(remaining) = &mut window.fetch {
1962            if *remaining == 0 {
1963                return Vec::new();
1964            }
1965            len = len.min(*remaining);
1966            *remaining -= len;
1967        }
1968
1969        vec![batch.slice(offset, len)]
1970    })
1971}
1972
1973fn apply_event_row_window(
1974    source: Source<SqlEvent<RecordBatch>>,
1975    skip: usize,
1976    fetch: Option<usize>,
1977) -> Source<SqlEvent<RecordBatch>> {
1978    source.stateful_map_concat(RowWindow { skip, fetch }, |window, event| match event {
1979        SqlEvent::Data(batch) => {
1980            let mut offset = 0;
1981            if window.skip > 0 {
1982                let skipped = window.skip.min(batch.num_rows());
1983                window.skip -= skipped;
1984                offset = skipped;
1985            }
1986
1987            if offset == batch.num_rows() {
1988                return Vec::new();
1989            }
1990
1991            let mut len = batch.num_rows() - offset;
1992            if let Some(remaining) = &mut window.fetch {
1993                if *remaining == 0 {
1994                    return Vec::new();
1995                }
1996                len = len.min(*remaining);
1997                *remaining -= len;
1998            }
1999
2000            vec![SqlEvent::Data(batch.slice(offset, len))]
2001        }
2002        SqlEvent::Watermark(watermark) => vec![SqlEvent::Watermark(watermark)],
2003        SqlEvent::Barrier(barrier) => vec![SqlEvent::Barrier(barrier)],
2004    })
2005}
2006
2007fn apply_changelog_row_window(
2008    source: Source<ChangelogBatch>,
2009    skip: usize,
2010    fetch: Option<usize>,
2011) -> Source<ChangelogBatch> {
2012    source.try_stateful_map_concat(RowWindow { skip, fetch }, |window, changes| {
2013        let mut offset = 0;
2014        if window.skip > 0 {
2015            let skipped = window.skip.min(changes.num_rows());
2016            window.skip -= skipped;
2017            offset = skipped;
2018        }
2019
2020        if offset == changes.num_rows() {
2021            return Ok(Vec::new());
2022        }
2023
2024        let mut len = changes.num_rows() - offset;
2025        if let Some(remaining) = &mut window.fetch {
2026            if *remaining == 0 {
2027                return Ok(Vec::new());
2028            }
2029            len = len.min(*remaining);
2030            *remaining -= len;
2031        }
2032
2033        let ops = changes.ops()[offset..offset + len].to_vec();
2034        let batch = changes.batch().slice(offset, len);
2035        let changes = ChangelogBatch::try_new(ops, batch).map_err(stream_error)?;
2036        Ok(vec![changes])
2037    })
2038}
2039
2040fn apply_committable_row_window(
2041    source: Source<CommittableRecordBatch>,
2042    skip: usize,
2043    fetch: Option<usize>,
2044) -> Source<CommittableRecordBatch> {
2045    source.try_stateful_map_concat(RowWindow { skip, fetch }, |window, batch| {
2046        apply_committable_row_window_to_batch(window, batch)
2047    })
2048}
2049
2050fn apply_committable_row_window_to_batch(
2051    window: &mut RowWindow,
2052    batch: CommittableRecordBatch,
2053) -> StreamResult<Vec<CommittableRecordBatch>> {
2054    let mut offset = 0;
2055    if window.skip > 0 {
2056        let skipped = window.skip.min(batch.batch().num_rows());
2057        window.skip -= skipped;
2058        offset = skipped;
2059    }
2060
2061    let mut len = batch.batch().num_rows().saturating_sub(offset);
2062    if let Some(remaining) = &mut window.fetch {
2063        if *remaining == 0 {
2064            return empty_committable_batch(batch).map(|batch| vec![batch]);
2065        }
2066        len = len.min(*remaining);
2067        *remaining -= len;
2068    }
2069
2070    if len == 0 {
2071        return empty_committable_batch(batch).map(|batch| vec![batch]);
2072    }
2073
2074    batch
2075        .try_map_batch(|record_batch| Ok(record_batch.slice(offset, len)))
2076        .map(|batch| vec![batch])
2077        .map_err(stream_error)
2078}
2079
2080fn empty_committable_batch(batch: CommittableRecordBatch) -> StreamResult<CommittableRecordBatch> {
2081    batch
2082        .try_map_batch(|record_batch| Ok(record_batch.slice(0, 0)))
2083        .map_err(stream_error)
2084}
2085
2086fn apply_changelog_event_row_window(
2087    source: Source<SqlEvent<ChangelogBatch>>,
2088    skip: usize,
2089    fetch: Option<usize>,
2090) -> Source<SqlEvent<ChangelogBatch>> {
2091    source.try_stateful_map_concat(RowWindow { skip, fetch }, |window, event| match event {
2092        SqlEvent::Data(changes) => apply_changelog_row_window_to_batch(window, changes)
2093            .map(|batches| batches.into_iter().map(SqlEvent::Data).collect::<Vec<_>>()),
2094        SqlEvent::Watermark(watermark) => Ok(vec![SqlEvent::Watermark(watermark)]),
2095        SqlEvent::Barrier(barrier) => Ok(vec![SqlEvent::Barrier(barrier)]),
2096    })
2097}
2098
2099fn apply_changelog_row_window_to_batch(
2100    window: &mut RowWindow,
2101    changes: ChangelogBatch,
2102) -> StreamResult<Vec<ChangelogBatch>> {
2103    let mut offset = 0;
2104    if window.skip > 0 {
2105        let skipped = window.skip.min(changes.num_rows());
2106        window.skip -= skipped;
2107        offset = skipped;
2108    }
2109
2110    if offset == changes.num_rows() {
2111        return Ok(Vec::new());
2112    }
2113
2114    let mut len = changes.num_rows() - offset;
2115    if let Some(remaining) = &mut window.fetch {
2116        if *remaining == 0 {
2117            return Ok(Vec::new());
2118        }
2119        len = len.min(*remaining);
2120        *remaining -= len;
2121    }
2122
2123    let ops = changes.ops()[offset..offset + len].to_vec();
2124    let batch = changes.batch().slice(offset, len);
2125    let changes = ChangelogBatch::try_new(ops, batch).map_err(stream_error)?;
2126    Ok(vec![changes])
2127}
2128
2129#[derive(Debug, Clone, PartialEq, Eq)]
2130struct ParsedInsertInto {
2131    sink_name: String,
2132    select_sql: String,
2133}
2134
2135fn looks_like_insert_into(sql: &str) -> bool {
2136    sql.trim_start()
2137        .get(..11)
2138        .is_some_and(|prefix| prefix.eq_ignore_ascii_case("insert into"))
2139}
2140
2141fn parse_insert_into(sql: &str) -> Result<ParsedInsertInto> {
2142    let mut statements = DFParser::parse_sql(sql)?;
2143    if statements.len() != 1 {
2144        return Err(DataFusionError::Plan(format!(
2145            "datum-sql INSERT INTO sink expects exactly one statement, found {}",
2146            statements.len()
2147        )));
2148    }
2149
2150    let statement = statements
2151        .pop_front()
2152        .expect("statement length was checked above");
2153    let DfStatement::Statement(statement) = statement else {
2154        return Err(DataFusionError::Plan(
2155            "datum-sql INSERT INTO sink supports SQL INSERT statements only".into(),
2156        ));
2157    };
2158    let SqlStatement::Insert(insert) = *statement else {
2159        return Err(DataFusionError::Plan(
2160            "datum-sql INSERT INTO sink supports INSERT INTO <sink> SELECT ... only".into(),
2161        ));
2162    };
2163
2164    if !insert.into {
2165        return Err(DataFusionError::Plan(
2166            "datum-sql sink SQL requires INSERT INTO <sink> SELECT ...".into(),
2167        ));
2168    }
2169    if !insert.optimizer_hints.is_empty()
2170        || insert.or.is_some()
2171        || insert.ignore
2172        || insert.table_alias.is_some()
2173        || !insert.columns.is_empty()
2174        || insert.overwrite
2175        || !insert.assignments.is_empty()
2176        || insert.partitioned.is_some()
2177        || !insert.after_columns.is_empty()
2178        || insert.has_table_keyword
2179        || insert.on.is_some()
2180        || insert.returning.is_some()
2181        || insert.output.is_some()
2182        || insert.replace_into
2183        || insert.priority.is_some()
2184        || insert.insert_alias.is_some()
2185        || insert.settings.is_some()
2186        || insert.format_clause.is_some()
2187        || insert.multi_table_insert_type.is_some()
2188        || !insert.multi_table_into_clauses.is_empty()
2189        || !insert.multi_table_when_clauses.is_empty()
2190        || insert.multi_table_else_clause.is_some()
2191    {
2192        return Err(DataFusionError::Plan(
2193            "datum-sql sink SQL supports only plain INSERT INTO <sink> SELECT ...".into(),
2194        ));
2195    }
2196
2197    let SqlTableObject::TableName(table_name) = insert.table else {
2198        return Err(DataFusionError::Plan(
2199            "datum-sql sink SQL requires a registered sink name after INSERT INTO".into(),
2200        ));
2201    };
2202    let sink_name = object_name_to_registry_key(table_name)?;
2203    let source = insert.source.ok_or_else(|| {
2204        DataFusionError::Plan("datum-sql sink SQL requires INSERT INTO <sink> SELECT ...".into())
2205    })?;
2206
2207    Ok(ParsedInsertInto {
2208        sink_name,
2209        select_sql: source.to_string(),
2210    })
2211}
2212
2213fn object_name_to_registry_key(
2214    name: datafusion::sql::sqlparser::ast::ObjectName,
2215) -> Result<String> {
2216    let mut parts = Vec::with_capacity(name.0.len());
2217    for part in name.0 {
2218        let Some(ident) = part.as_ident() else {
2219            return Err(DataFusionError::Plan(
2220                "datum-sql sink names must be identifiers".into(),
2221            ));
2222        };
2223        parts.push(ident.value.clone());
2224    }
2225    if parts.is_empty() {
2226        return Err(DataFusionError::Plan(
2227            "datum-sql sink name cannot be empty".into(),
2228        ));
2229    }
2230    Ok(parts.join("."))
2231}
2232
2233fn updating_scan_error() -> DataFusionError {
2234    DataFusionError::Plan(
2235        "datum-sql updating tables produce Source<ChangelogBatch>; use changelog-aware lowering or a changelog-aware sink instead of the append-only RecordBatch executor"
2236            .into(),
2237    )
2238}
2239
2240fn datafusion_error(error: impl fmt::Display) -> DataFusionError {
2241    DataFusionError::External(Box::new(StreamBridgeError(error.to_string())))
2242}
2243
2244fn stream_error(error: impl fmt::Display) -> StreamError {
2245    StreamError::Failed(error.to_string())
2246}
2247
2248#[derive(Debug)]
2249struct StreamBridgeError(String);
2250
2251impl fmt::Display for StreamBridgeError {
2252    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2253        f.write_str(&self.0)
2254    }
2255}
2256
2257impl std::error::Error for StreamBridgeError {}