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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
81pub enum ChangeOp {
82 Insert,
84 Delete,
86 UpdateDelete,
88 UpdateInsert,
90}
91
92impl ChangeOp {
93 #[must_use]
95 pub const fn is_positive(self) -> bool {
96 matches!(self, Self::Insert | Self::UpdateInsert)
97 }
98
99 #[must_use]
101 pub const fn is_retraction(self) -> bool {
102 matches!(self, Self::Delete | Self::UpdateDelete)
103 }
104}
105
106#[derive(Clone)]
111pub struct ChangelogBatch {
112 ops: Arc<[ChangeOp]>,
113 batch: RecordBatch,
114}
115
116impl ChangelogBatch {
117 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 #[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 #[must_use]
138 pub fn ops(&self) -> &[ChangeOp] {
139 &self.ops
140 }
141
142 #[must_use]
144 pub fn batch(&self) -> &RecordBatch {
145 &self.batch
146 }
147
148 #[must_use]
150 pub fn num_rows(&self) -> usize {
151 self.batch.num_rows()
152 }
153
154 #[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#[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 #[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 #[must_use]
222 pub fn streaming_join_metrics(&self) -> StreamingJoinMetrics {
223 self.join_metrics()
224 }
225
226 #[must_use]
228 pub fn join_state_rows(&self) -> u64 {
229 self.join_state_rows.load(Ordering::Relaxed)
230 }
231
232 #[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
247pub struct DatumSqlContext {
250 inner: SessionContext,
251 sinks: Arc<Mutex<BTreeMap<String, RegisteredSqlSink>>>,
252 streaming_join_config: StreamingJoinConfig,
253}
254
255impl DatumSqlContext {
256 #[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 #[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 pub fn set_streaming_join_config(&mut self, config: StreamingJoinConfig) {
279 self.streaming_join_config = config;
280 }
281
282 #[must_use]
284 pub const fn streaming_join_config(&self) -> &StreamingJoinConfig {
285 &self.streaming_join_config
286 }
287
288 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 {}