datum_sql/
join.rs

1use std::collections::HashMap;
2use std::fmt;
3use std::hash::{Hash, Hasher};
4use std::sync::Arc;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::time::Duration;
7
8use arrow::array::{
9    Array, ArrayRef, RecordBatch, TimestampMicrosecondArray, TimestampMillisecondArray,
10    TimestampNanosecondArray, TimestampSecondArray,
11};
12use arrow::datatypes::{DataType, SchemaRef, TimeUnit};
13use datafusion::common::cast::as_boolean_array;
14use datafusion::common::{DataFusionError, JoinSide, JoinType, NullEquality, Result, ScalarValue};
15use datafusion::physical_expr::PhysicalExpr;
16use datafusion::physical_plan::ExecutionPlan;
17use datafusion::physical_plan::joins::HashJoinExec;
18use datafusion::physical_plan::joins::utils::JoinFilter;
19use datum::{Source, StreamResult};
20
21use crate::{SqlEvent, Watermark, stream_error};
22
23/// Streaming join lowering policy.
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct StreamingJoinConfig {
26    mode: StreamingJoinMode,
27}
28
29impl StreamingJoinConfig {
30    /// Creates a join config from an explicit mode.
31    #[must_use]
32    pub const fn new(mode: StreamingJoinMode) -> Self {
33        Self { mode }
34    }
35
36    /// Creates an event-time windowed join config.
37    #[must_use]
38    pub fn windowed(window: StreamingJoinWindow) -> Self {
39        Self {
40            mode: StreamingJoinMode::Windowed { window },
41        }
42    }
43
44    /// Creates an unbounded equi-join config with hard in-memory limits.
45    #[must_use]
46    pub const fn bounded_state(limits: StreamingJoinStateLimits) -> Self {
47        Self {
48            mode: StreamingJoinMode::BoundedState { limits },
49        }
50    }
51
52    /// Returns the selected join state policy.
53    #[must_use]
54    pub const fn mode(&self) -> &StreamingJoinMode {
55        &self.mode
56    }
57}
58
59impl Default for StreamingJoinConfig {
60    fn default() -> Self {
61        Self::windowed(StreamingJoinWindow::default())
62    }
63}
64
65/// Supported streaming join state classes.
66#[derive(Debug, Clone, PartialEq, Eq)]
67pub enum StreamingJoinMode {
68    /// Event-time interval join. Rows are retained until both input watermarks
69    /// prove they can no longer match a future row within `window`.
70    Windowed {
71        /// Event-time retention policy.
72        window: StreamingJoinWindow,
73    },
74    /// Unbounded SQL equi-join semantics with explicit hard in-memory limits.
75    BoundedState {
76        /// Hard in-memory state limits.
77        limits: StreamingJoinStateLimits,
78    },
79}
80
81/// Event-time interval join settings.
82#[derive(Debug, Clone, PartialEq, Eq)]
83pub struct StreamingJoinWindow {
84    time_column: Arc<str>,
85    max_time_difference: Duration,
86}
87
88impl StreamingJoinWindow {
89    /// Creates a windowed join policy.
90    ///
91    /// `time_column` must exist on both join inputs after scan projection.
92    #[must_use]
93    pub fn new(time_column: impl Into<Arc<str>>, max_time_difference: Duration) -> Self {
94        Self {
95            time_column: time_column.into(),
96            max_time_difference,
97        }
98    }
99
100    /// Returns the input timestamp column used for retention and late-row checks.
101    #[must_use]
102    pub fn time_column(&self) -> &str {
103        &self.time_column
104    }
105
106    /// Returns the maximum allowed timestamp distance between joined rows.
107    #[must_use]
108    pub const fn max_time_difference(&self) -> Duration {
109        self.max_time_difference
110    }
111}
112
113impl Default for StreamingJoinWindow {
114    fn default() -> Self {
115        Self::new("date_time", Duration::from_secs(10))
116    }
117}
118
119/// Hard limits for unbounded streaming equi-joins.
120#[derive(Debug, Clone, Copy, PartialEq, Eq)]
121pub struct StreamingJoinStateLimits {
122    max_rows_per_side: usize,
123    max_total_rows: usize,
124}
125
126impl StreamingJoinStateLimits {
127    /// Creates hard state limits for a bounded-state join.
128    #[must_use]
129    pub const fn new(max_rows_per_side: usize, max_total_rows: usize) -> Self {
130        Self {
131            max_rows_per_side,
132            max_total_rows,
133        }
134    }
135
136    /// Maximum rows retained on either input side.
137    #[must_use]
138    pub const fn max_rows_per_side(self) -> usize {
139        self.max_rows_per_side
140    }
141
142    /// Maximum rows retained across both input sides.
143    #[must_use]
144    pub const fn max_total_rows(self) -> usize {
145        self.max_total_rows
146    }
147}
148
149/// Counters maintained by streaming join stages.
150#[derive(Clone, Default)]
151pub struct StreamingJoinMetrics {
152    state_rows: Arc<AtomicU64>,
153    evicted_rows: Arc<AtomicU64>,
154    late_dropped_rows: Arc<AtomicU64>,
155}
156
157impl StreamingJoinMetrics {
158    pub(crate) fn new(
159        state_rows: Arc<AtomicU64>,
160        evicted_rows: Arc<AtomicU64>,
161        late_dropped_rows: Arc<AtomicU64>,
162    ) -> Self {
163        Self {
164            state_rows,
165            evicted_rows,
166            late_dropped_rows,
167        }
168    }
169
170    /// Rows currently retained in join state.
171    #[must_use]
172    pub fn state_rows(&self) -> u64 {
173        self.state_rows.load(Ordering::Relaxed)
174    }
175
176    /// Rows evicted from join state after watermark progress.
177    #[must_use]
178    pub fn evicted_rows(&self) -> u64 {
179        self.evicted_rows.load(Ordering::Relaxed)
180    }
181
182    /// Rows dropped because they arrived at or behind their input watermark.
183    #[must_use]
184    pub fn late_dropped_rows(&self) -> u64 {
185        self.late_dropped_rows.load(Ordering::Relaxed)
186    }
187
188    fn set_state_rows(&self, rows: usize) {
189        self.state_rows.store(rows as u64, Ordering::Relaxed);
190    }
191
192    fn record_evictions(&self, rows: usize) {
193        self.evicted_rows.fetch_add(rows as u64, Ordering::Relaxed);
194    }
195
196    fn record_late_row(&self) {
197        self.late_dropped_rows.fetch_add(1, Ordering::Relaxed);
198    }
199}
200
201impl fmt::Debug for StreamingJoinMetrics {
202    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
203        f.debug_struct("StreamingJoinMetrics")
204            .field("state_rows", &self.state_rows())
205            .field("evicted_rows", &self.evicted_rows())
206            .field("late_dropped_rows", &self.late_dropped_rows())
207            .finish()
208    }
209}
210
211pub(crate) fn streaming_hash_join_source(
212    left: Source<SqlEvent<RecordBatch>>,
213    right: Source<SqlEvent<RecordBatch>>,
214    join: &HashJoinExec,
215    config: &StreamingJoinConfig,
216    metrics: StreamingJoinMetrics,
217) -> Result<Source<SqlEvent<RecordBatch>>> {
218    let stage = StreamingJoinStage::try_new(join, config, metrics)?;
219    let tagged_left = left.map(TaggedJoinEvent::Left);
220    let tagged_right = right.map(TaggedJoinEvent::Right);
221    Ok(tagged_left
222        .merge_all([tagged_right], false)
223        .try_stateful_map_concat(stage, |stage, event| stage.apply(event)))
224}
225
226#[derive(Debug, Clone, Copy, PartialEq, Eq)]
227enum JoinInputSide {
228    Left,
229    Right,
230}
231
232enum TaggedJoinEvent {
233    Left(SqlEvent<RecordBatch>),
234    Right(SqlEvent<RecordBatch>),
235}
236
237#[derive(Clone)]
238struct StreamingJoinStage {
239    plan: Arc<StreamingJoinPlan>,
240    metrics: StreamingJoinMetrics,
241    latest_left_watermark_ns: Option<i64>,
242    latest_right_watermark_ns: Option<i64>,
243    last_emitted_watermark_ns: Option<i64>,
244    left: JoinSideState,
245    right: JoinSideState,
246}
247
248impl StreamingJoinStage {
249    fn try_new(
250        join: &HashJoinExec,
251        config: &StreamingJoinConfig,
252        metrics: StreamingJoinMetrics,
253    ) -> Result<Self> {
254        Ok(Self {
255            plan: Arc::new(StreamingJoinPlan::try_new(join, config)?),
256            metrics,
257            latest_left_watermark_ns: None,
258            latest_right_watermark_ns: None,
259            last_emitted_watermark_ns: None,
260            left: JoinSideState::default(),
261            right: JoinSideState::default(),
262        })
263    }
264
265    fn apply(&mut self, event: TaggedJoinEvent) -> StreamResult<Vec<SqlEvent<RecordBatch>>> {
266        let result = match event {
267            TaggedJoinEvent::Left(SqlEvent::Data(batch)) => {
268                self.apply_batch(JoinInputSide::Left, batch)
269            }
270            TaggedJoinEvent::Right(SqlEvent::Data(batch)) => {
271                self.apply_batch(JoinInputSide::Right, batch)
272            }
273            TaggedJoinEvent::Left(SqlEvent::Watermark(watermark)) => {
274                self.apply_watermark(JoinInputSide::Left, watermark)
275            }
276            TaggedJoinEvent::Right(SqlEvent::Watermark(watermark)) => {
277                self.apply_watermark(JoinInputSide::Right, watermark)
278            }
279            TaggedJoinEvent::Left(SqlEvent::Barrier(barrier))
280            | TaggedJoinEvent::Right(SqlEvent::Barrier(barrier)) => {
281                Ok(vec![SqlEvent::Barrier(barrier)])
282            }
283        };
284        result.map_err(stream_error)
285    }
286
287    fn apply_batch(
288        &mut self,
289        side: JoinInputSide,
290        batch: RecordBatch,
291    ) -> Result<Vec<SqlEvent<RecordBatch>>> {
292        if batch.num_rows() == 0 {
293            return Ok(Vec::new());
294        }
295
296        let prepared = self.plan.prepare_batch(side, &batch)?;
297        let mut output_rows = Vec::new();
298        for row in 0..batch.num_rows() {
299            let event_time_ns = self.plan.event_time_ns(side, &prepared, row)?;
300            if self.row_is_late(side, event_time_ns) {
301                self.metrics.record_late_row();
302                continue;
303            }
304            let Some(key) = self.plan.key_at(side, &prepared, row)? else {
305                continue;
306            };
307            let state_row = RowState::from_batch(&batch, row, event_time_ns)?;
308
309            match side {
310                JoinInputSide::Left => {
311                    if let Some(candidates) = self.right.rows.get(&key) {
312                        for right_row in candidates {
313                            if self.plan.rows_can_join(&state_row, right_row)?
314                                && self.plan.pair_passes_filter(&state_row, right_row)?
315                            {
316                                output_rows.push(self.plan.output_row(&state_row, right_row));
317                            }
318                        }
319                    }
320                    self.check_state_limit(JoinInputSide::Left)?;
321                    self.left.insert(key, state_row);
322                }
323                JoinInputSide::Right => {
324                    if let Some(candidates) = self.left.rows.get(&key) {
325                        for left_row in candidates {
326                            if self.plan.rows_can_join(left_row, &state_row)?
327                                && self.plan.pair_passes_filter(left_row, &state_row)?
328                            {
329                                output_rows.push(self.plan.output_row(left_row, &state_row));
330                            }
331                        }
332                    }
333                    self.check_state_limit(JoinInputSide::Right)?;
334                    self.right.insert(key, state_row);
335                }
336            }
337        }
338
339        self.update_state_rows_metric();
340        if output_rows.is_empty() {
341            return Ok(Vec::new());
342        }
343        Ok(vec![SqlEvent::Data(
344            self.plan.build_output_batch(output_rows)?,
345        )])
346    }
347
348    fn apply_watermark(
349        &mut self,
350        side: JoinInputSide,
351        watermark: Watermark,
352    ) -> Result<Vec<SqlEvent<RecordBatch>>> {
353        let watermark_ns = watermark.timestamp_ns();
354        match side {
355            JoinInputSide::Left => {
356                self.latest_left_watermark_ns = Some(
357                    self.latest_left_watermark_ns
358                        .map_or(watermark_ns, |current| current.max(watermark_ns)),
359                );
360            }
361            JoinInputSide::Right => {
362                self.latest_right_watermark_ns = Some(
363                    self.latest_right_watermark_ns
364                        .map_or(watermark_ns, |current| current.max(watermark_ns)),
365                );
366            }
367        }
368
369        let Some(min_watermark_ns) = self.min_input_watermark() else {
370            return Ok(Vec::new());
371        };
372        if self
373            .last_emitted_watermark_ns
374            .is_some_and(|last| min_watermark_ns <= last)
375        {
376            return Ok(Vec::new());
377        }
378
379        self.evict_for_watermark(min_watermark_ns)?;
380        self.last_emitted_watermark_ns = Some(min_watermark_ns);
381        Ok(vec![SqlEvent::Watermark(Watermark::new(min_watermark_ns))])
382    }
383
384    fn row_is_late(&self, side: JoinInputSide, event_time_ns: Option<i64>) -> bool {
385        let Some(event_time_ns) = event_time_ns else {
386            return false;
387        };
388        match side {
389            JoinInputSide::Left => self
390                .latest_left_watermark_ns
391                .is_some_and(|watermark_ns| event_time_ns <= watermark_ns),
392            JoinInputSide::Right => self
393                .latest_right_watermark_ns
394                .is_some_and(|watermark_ns| event_time_ns <= watermark_ns),
395        }
396    }
397
398    fn min_input_watermark(&self) -> Option<i64> {
399        Some(
400            self.latest_left_watermark_ns?
401                .min(self.latest_right_watermark_ns?),
402        )
403    }
404
405    fn evict_for_watermark(&mut self, watermark_ns: i64) -> Result<()> {
406        let Some(max_difference_ns) = self.plan.max_time_difference_ns() else {
407            return Ok(());
408        };
409        let left_evicted = self.left.evict_windowed(watermark_ns, max_difference_ns)?;
410        let right_evicted = self.right.evict_windowed(watermark_ns, max_difference_ns)?;
411        let evicted = left_evicted + right_evicted;
412        if evicted > 0 {
413            self.metrics.record_evictions(evicted);
414            self.update_state_rows_metric();
415        }
416        Ok(())
417    }
418
419    fn check_state_limit(&self, side: JoinInputSide) -> Result<()> {
420        let Some(limits) = self.plan.state_limits() else {
421            return Ok(());
422        };
423        let next_side_rows = match side {
424            JoinInputSide::Left => self.left.row_count + 1,
425            JoinInputSide::Right => self.right.row_count + 1,
426        };
427        let next_total_rows = self.left.row_count + self.right.row_count + 1;
428        if next_side_rows > limits.max_rows_per_side() {
429            return Err(DataFusionError::Plan(format!(
430                "streaming join state limit exceeded: {side:?} side would hold {next_side_rows} rows, limit is {}",
431                limits.max_rows_per_side()
432            )));
433        }
434        if next_total_rows > limits.max_total_rows() {
435            return Err(DataFusionError::Plan(format!(
436                "streaming join state limit exceeded: total state would hold {next_total_rows} rows, limit is {}",
437                limits.max_total_rows()
438            )));
439        }
440        Ok(())
441    }
442
443    fn update_state_rows_metric(&self) {
444        self.metrics
445            .set_state_rows(self.left.row_count + self.right.row_count);
446    }
447}
448
449#[derive(Clone)]
450struct StreamingJoinPlan {
451    join_schema: SchemaRef,
452    output_schema: SchemaRef,
453    left_key_exprs: Vec<Arc<dyn PhysicalExpr>>,
454    right_key_exprs: Vec<Arc<dyn PhysicalExpr>>,
455    filter: Option<JoinFilter>,
456    projection: Option<Arc<[usize]>>,
457    null_equality: NullEquality,
458    mode: StreamingJoinPlanMode,
459}
460
461impl StreamingJoinPlan {
462    fn try_new(join: &HashJoinExec, config: &StreamingJoinConfig) -> Result<Self> {
463        if *join.join_type() != JoinType::Inner {
464            return Err(DataFusionError::NotImplemented(format!(
465                "datum-sql streaming joins support INNER equi-joins only for now, found {:?}",
466                join.join_type()
467            )));
468        }
469        if join.on().is_empty() {
470            return Err(DataFusionError::NotImplemented(
471                "datum-sql streaming joins require at least one equi-join key".into(),
472            ));
473        }
474        if join.null_aware {
475            return Err(DataFusionError::NotImplemented(
476                "datum-sql streaming joins do not support null-aware anti joins".into(),
477            ));
478        }
479
480        let left_schema = join.left().schema();
481        let right_schema = join.right().schema();
482        let mode = StreamingJoinPlanMode::try_new(config, &left_schema, &right_schema)?;
483        let output_schema = join.schema();
484        Ok(Self {
485            join_schema: Arc::clone(join.join_schema()),
486            output_schema,
487            left_key_exprs: join
488                .on()
489                .iter()
490                .map(|(left, _right)| Arc::clone(left))
491                .collect(),
492            right_key_exprs: join
493                .on()
494                .iter()
495                .map(|(_left, right)| Arc::clone(right))
496                .collect(),
497            filter: join.filter().cloned(),
498            projection: join.projection.clone(),
499            null_equality: join.null_equality(),
500            mode,
501        })
502    }
503
504    fn prepare_batch(&self, side: JoinInputSide, batch: &RecordBatch) -> Result<PreparedJoinBatch> {
505        let key_exprs = match side {
506            JoinInputSide::Left => &self.left_key_exprs,
507            JoinInputSide::Right => &self.right_key_exprs,
508        };
509        let key_values = key_exprs
510            .iter()
511            .map(|expr| expr.evaluate(batch)?.into_array(batch.num_rows()))
512            .collect::<Result<Vec<_>>>()?;
513        let event_times = match (&self.mode, side) {
514            (
515                StreamingJoinPlanMode::Windowed {
516                    left_time_index, ..
517                },
518                JoinInputSide::Left,
519            ) => Some(Arc::clone(batch.column(*left_time_index))),
520            (
521                StreamingJoinPlanMode::Windowed {
522                    right_time_index, ..
523                },
524                JoinInputSide::Right,
525            ) => Some(Arc::clone(batch.column(*right_time_index))),
526            (StreamingJoinPlanMode::BoundedState { .. }, _) => None,
527        };
528        Ok(PreparedJoinBatch {
529            key_values,
530            event_times,
531        })
532    }
533
534    fn key_at(
535        &self,
536        _side: JoinInputSide,
537        prepared: &PreparedJoinBatch,
538        row: usize,
539    ) -> Result<Option<JoinKey>> {
540        let mut key = Vec::with_capacity(prepared.key_values.len());
541        let mut has_null = false;
542        for array in &prepared.key_values {
543            let value = ScalarValue::try_from_array(array.as_ref(), row)?;
544            has_null |= value.is_null();
545            key.push(value);
546        }
547        if has_null && self.null_equality == NullEquality::NullEqualsNothing {
548            Ok(None)
549        } else {
550            Ok(Some(JoinKey(key)))
551        }
552    }
553
554    fn event_time_ns(
555        &self,
556        side: JoinInputSide,
557        prepared: &PreparedJoinBatch,
558        row: usize,
559    ) -> Result<Option<i64>> {
560        let index = match (&self.mode, side) {
561            (
562                StreamingJoinPlanMode::Windowed {
563                    left_time_index, ..
564                },
565                JoinInputSide::Left,
566            ) => *left_time_index,
567            (
568                StreamingJoinPlanMode::Windowed {
569                    right_time_index, ..
570                },
571                JoinInputSide::Right,
572            ) => *right_time_index,
573            (StreamingJoinPlanMode::BoundedState { .. }, _) => return Ok(None),
574        };
575        let event_times = prepared.event_times.as_ref().ok_or_else(|| {
576            DataFusionError::Internal("windowed join prepared batch is missing event time".into())
577        })?;
578        timestamp_ns_from_array(event_times, row).map(Some).map_err(|error| {
579            DataFusionError::Plan(format!(
580                "streaming join event-time column at index {index} on {side:?} input failed: {error}"
581            ))
582        })
583    }
584
585    fn rows_can_join(&self, left: &RowState, right: &RowState) -> Result<bool> {
586        let Some(max_difference_ns) = self.max_time_difference_ns() else {
587            return Ok(true);
588        };
589        let left_time_ns = left.event_time_ns.ok_or_else(|| {
590            DataFusionError::Internal("windowed join left row is missing event time".into())
591        })?;
592        let right_time_ns = right.event_time_ns.ok_or_else(|| {
593            DataFusionError::Internal("windowed join right row is missing event time".into())
594        })?;
595        Ok(left_time_ns.abs_diff(right_time_ns) <= max_difference_ns as u64)
596    }
597
598    fn pair_passes_filter(&self, left: &RowState, right: &RowState) -> Result<bool> {
599        let Some(filter) = &self.filter else {
600            return Ok(true);
601        };
602        let values = filter
603            .column_indices()
604            .iter()
605            .map(|column| match column.side {
606                JoinSide::Left => Ok(left.values[column.index].clone()),
607                JoinSide::Right => Ok(right.values[column.index].clone()),
608                JoinSide::None => Err(DataFusionError::NotImplemented(
609                    "datum-sql streaming join filters do not support unqualified join-side columns"
610                        .into(),
611                )),
612            })
613            .collect::<Result<Vec<_>>>()?;
614        let arrays = values
615            .into_iter()
616            .map(|value| ScalarValue::iter_to_array([value]))
617            .collect::<Result<Vec<_>>>()?;
618        let batch = RecordBatch::try_new(Arc::clone(filter.schema()), arrays)?;
619        let predicate = filter.expression().evaluate(&batch)?.into_array(1)?;
620        let predicate = as_boolean_array(&predicate)?;
621        Ok(predicate.is_valid(0) && predicate.value(0))
622    }
623
624    fn output_row(&self, left: &RowState, right: &RowState) -> Vec<ScalarValue> {
625        let mut row = Vec::with_capacity(left.values.len() + right.values.len());
626        row.extend(left.values.iter().cloned());
627        row.extend(right.values.iter().cloned());
628        if let Some(projection) = &self.projection {
629            projection.iter().map(|index| row[*index].clone()).collect()
630        } else {
631            row
632        }
633    }
634
635    fn build_output_batch(&self, rows: Vec<Vec<ScalarValue>>) -> Result<RecordBatch> {
636        let schema = if self.projection.is_some() {
637            &self.output_schema
638        } else {
639            &self.join_schema
640        };
641        let column_count = schema.fields().len();
642        let mut columns = vec![Vec::with_capacity(rows.len()); column_count];
643        for row in rows {
644            if row.len() != column_count {
645                return Err(DataFusionError::Internal(format!(
646                    "streaming join produced {} values for {column_count} output columns",
647                    row.len()
648                )));
649            }
650            for (index, value) in row.into_iter().enumerate() {
651                columns[index].push(value);
652            }
653        }
654        let arrays = columns
655            .into_iter()
656            .map(ScalarValue::iter_to_array)
657            .collect::<Result<Vec<_>>>()?;
658        RecordBatch::try_new(Arc::clone(schema), arrays).map_err(DataFusionError::from)
659    }
660
661    fn max_time_difference_ns(&self) -> Option<i64> {
662        match &self.mode {
663            StreamingJoinPlanMode::Windowed {
664                max_time_difference_ns,
665                ..
666            } => Some(*max_time_difference_ns),
667            StreamingJoinPlanMode::BoundedState { .. } => None,
668        }
669    }
670
671    fn state_limits(&self) -> Option<StreamingJoinStateLimits> {
672        match self.mode {
673            StreamingJoinPlanMode::Windowed { .. } => None,
674            StreamingJoinPlanMode::BoundedState { limits } => Some(limits),
675        }
676    }
677}
678
679#[derive(Clone)]
680enum StreamingJoinPlanMode {
681    Windowed {
682        left_time_index: usize,
683        right_time_index: usize,
684        max_time_difference_ns: i64,
685    },
686    BoundedState {
687        limits: StreamingJoinStateLimits,
688    },
689}
690
691impl StreamingJoinPlanMode {
692    fn try_new(
693        config: &StreamingJoinConfig,
694        left_schema: &SchemaRef,
695        right_schema: &SchemaRef,
696    ) -> Result<Self> {
697        match config.mode() {
698            StreamingJoinMode::Windowed { window } => {
699                let max_time_difference_ns = duration_to_ns(window.max_time_difference())?;
700                if max_time_difference_ns <= 0 {
701                    return Err(DataFusionError::Plan(format!(
702                        "streaming join window must be positive, found {:?}",
703                        window.max_time_difference()
704                    )));
705                }
706                Ok(Self::Windowed {
707                    left_time_index: resolve_timestamp_column(left_schema, window.time_column())?,
708                    right_time_index: resolve_timestamp_column(right_schema, window.time_column())?,
709                    max_time_difference_ns,
710                })
711            }
712            StreamingJoinMode::BoundedState { limits } => {
713                if limits.max_rows_per_side() == 0 || limits.max_total_rows() == 0 {
714                    return Err(DataFusionError::Plan(
715                        "streaming join bounded-state limits must be greater than zero".into(),
716                    ));
717                }
718                if limits.max_total_rows() < limits.max_rows_per_side() {
719                    return Err(DataFusionError::Plan(format!(
720                        "streaming join max_total_rows {} must be at least max_rows_per_side {}",
721                        limits.max_total_rows(),
722                        limits.max_rows_per_side()
723                    )));
724                }
725                Ok(Self::BoundedState { limits: *limits })
726            }
727        }
728    }
729}
730
731struct PreparedJoinBatch {
732    key_values: Vec<ArrayRef>,
733    event_times: Option<ArrayRef>,
734}
735
736#[derive(Clone)]
737struct RowState {
738    values: Vec<ScalarValue>,
739    event_time_ns: Option<i64>,
740}
741
742impl RowState {
743    fn from_batch(batch: &RecordBatch, row: usize, event_time_ns: Option<i64>) -> Result<Self> {
744        let values = batch
745            .columns()
746            .iter()
747            .map(|array| ScalarValue::try_from_array(array.as_ref(), row))
748            .collect::<Result<Vec<_>>>()?;
749        Ok(Self {
750            values,
751            event_time_ns,
752        })
753    }
754}
755
756#[derive(Clone, Default)]
757struct JoinSideState {
758    rows: HashMap<JoinKey, Vec<RowState>>,
759    row_count: usize,
760}
761
762impl JoinSideState {
763    fn insert(&mut self, key: JoinKey, row: RowState) {
764        self.rows.entry(key).or_default().push(row);
765        self.row_count += 1;
766    }
767
768    fn evict_windowed(&mut self, watermark_ns: i64, max_difference_ns: i64) -> Result<usize> {
769        let mut evicted = 0;
770        self.rows.retain(|_key, rows| {
771            let before = rows.len();
772            rows.retain(|row| {
773                row.event_time_ns
774                    .and_then(|event_time_ns| event_time_ns.checked_add(max_difference_ns))
775                    .is_none_or(|expires_ns| expires_ns > watermark_ns)
776            });
777            evicted += before - rows.len();
778            !rows.is_empty()
779        });
780        self.row_count = self.row_count.checked_sub(evicted).ok_or_else(|| {
781            DataFusionError::Internal("streaming join state row count underflowed".into())
782        })?;
783        Ok(evicted)
784    }
785}
786
787#[derive(Clone, Eq)]
788struct JoinKey(Vec<ScalarValue>);
789
790impl PartialEq for JoinKey {
791    fn eq(&self, other: &Self) -> bool {
792        self.0 == other.0
793    }
794}
795
796impl Hash for JoinKey {
797    fn hash<H: Hasher>(&self, state: &mut H) {
798        self.0.hash(state);
799    }
800}
801
802fn resolve_timestamp_column(schema: &SchemaRef, name: &str) -> Result<usize> {
803    let index = schema.index_of(name).map_err(|_| {
804        DataFusionError::Plan(format!(
805            "streaming windowed join requires timestamp column '{name}' on both inputs"
806        ))
807    })?;
808    match schema.field(index).data_type() {
809        DataType::Timestamp(_, _) => Ok(index),
810        other => Err(DataFusionError::Plan(format!(
811            "streaming windowed join column '{name}' must be Timestamp, found {other:?}"
812        ))),
813    }
814}
815
816fn timestamp_ns_from_array(array: &ArrayRef, row: usize) -> Result<i64> {
817    match array.data_type() {
818        DataType::Timestamp(TimeUnit::Second, _) => timestamp_value_ns::<TimestampSecondArray>(
819            array.as_any().downcast_ref::<TimestampSecondArray>(),
820            row,
821            1_000_000_000,
822        ),
823        DataType::Timestamp(TimeUnit::Millisecond, _) => {
824            timestamp_value_ns::<TimestampMillisecondArray>(
825                array.as_any().downcast_ref::<TimestampMillisecondArray>(),
826                row,
827                1_000_000,
828            )
829        }
830        DataType::Timestamp(TimeUnit::Microsecond, _) => {
831            timestamp_value_ns::<TimestampMicrosecondArray>(
832                array.as_any().downcast_ref::<TimestampMicrosecondArray>(),
833                row,
834                1_000,
835            )
836        }
837        DataType::Timestamp(TimeUnit::Nanosecond, _) => {
838            timestamp_value_ns::<TimestampNanosecondArray>(
839                array.as_any().downcast_ref::<TimestampNanosecondArray>(),
840                row,
841                1,
842            )
843        }
844        other => Err(DataFusionError::Plan(format!(
845            "streaming join event-time expression must evaluate to Timestamp, found {other:?}"
846        ))),
847    }
848}
849
850fn timestamp_value_ns<T>(array: Option<&T>, row: usize, multiplier: i64) -> Result<i64>
851where
852    T: Array + TimestampArrayValue,
853{
854    let array =
855        array.ok_or_else(|| DataFusionError::Internal("timestamp array type mismatch".into()))?;
856    if array.is_null(row) {
857        return Err(DataFusionError::Plan(format!(
858            "streaming join event-time column contains null at row {row}"
859        )));
860    }
861    array
862        .timestamp_value(row)
863        .checked_mul(multiplier)
864        .ok_or_else(|| DataFusionError::Plan("timestamp overflowed i64 nanoseconds".into()))
865}
866
867trait TimestampArrayValue {
868    fn timestamp_value(&self, row: usize) -> i64;
869}
870
871impl TimestampArrayValue for TimestampSecondArray {
872    fn timestamp_value(&self, row: usize) -> i64 {
873        self.value(row)
874    }
875}
876
877impl TimestampArrayValue for TimestampMillisecondArray {
878    fn timestamp_value(&self, row: usize) -> i64 {
879        self.value(row)
880    }
881}
882
883impl TimestampArrayValue for TimestampMicrosecondArray {
884    fn timestamp_value(&self, row: usize) -> i64 {
885        self.value(row)
886    }
887}
888
889impl TimestampArrayValue for TimestampNanosecondArray {
890    fn timestamp_value(&self, row: usize) -> i64 {
891        self.value(row)
892    }
893}
894
895fn duration_to_ns(duration: Duration) -> Result<i64> {
896    i64::try_from(duration.as_nanos()).map_err(|_| {
897        DataFusionError::Plan(format!(
898            "streaming join window {duration:?} exceeds i64 nanoseconds"
899        ))
900    })
901}
datum_sql/join.rs

datum_sql/
join.rs
