datum-sql 0.10.3

use std::collections::HashMap;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

use arrow::array::{
    Array, ArrayRef, RecordBatch, TimestampMicrosecondArray, TimestampMillisecondArray,
    TimestampNanosecondArray, TimestampSecondArray,
};
use arrow::datatypes::{DataType, SchemaRef, TimeUnit};
use datafusion::common::cast::as_boolean_array;
use datafusion::common::{DataFusionError, JoinSide, JoinType, NullEquality, Result, ScalarValue};
use datafusion::physical_expr::PhysicalExpr;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::joins::HashJoinExec;
use datafusion::physical_plan::joins::utils::JoinFilter;
use datum::{Source, StreamResult};

use crate::{SqlEvent, Watermark, stream_error};

/// Streaming join lowering policy.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StreamingJoinConfig {
    mode: StreamingJoinMode,
}

impl StreamingJoinConfig {
    /// Creates a join config from an explicit mode.
    #[must_use]
    pub const fn new(mode: StreamingJoinMode) -> Self {
        Self { mode }
    }

    /// Creates an event-time windowed join config.
    #[must_use]
    pub fn windowed(window: StreamingJoinWindow) -> Self {
        Self {
            mode: StreamingJoinMode::Windowed { window },
        }
    }

    /// Creates an unbounded equi-join config with hard in-memory limits.
    #[must_use]
    pub const fn bounded_state(limits: StreamingJoinStateLimits) -> Self {
        Self {
            mode: StreamingJoinMode::BoundedState { limits },
        }
    }

    /// Returns the selected join state policy.
    #[must_use]
    pub const fn mode(&self) -> &StreamingJoinMode {
        &self.mode
    }
}

impl Default for StreamingJoinConfig {
    fn default() -> Self {
        Self::windowed(StreamingJoinWindow::default())
    }
}

/// Supported streaming join state classes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StreamingJoinMode {
    /// Event-time interval join. Rows are retained until both input watermarks
    /// prove they can no longer match a future row within `window`.
    Windowed {
        /// Event-time retention policy.
        window: StreamingJoinWindow,
    },
    /// Unbounded SQL equi-join semantics with explicit hard in-memory limits.
    BoundedState {
        /// Hard in-memory state limits.
        limits: StreamingJoinStateLimits,
    },
}

/// Event-time interval join settings.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StreamingJoinWindow {
    time_column: Arc<str>,
    max_time_difference: Duration,
}

impl StreamingJoinWindow {
    /// Creates a windowed join policy.
    ///
    /// `time_column` must exist on both join inputs after scan projection.
    #[must_use]
    pub fn new(time_column: impl Into<Arc<str>>, max_time_difference: Duration) -> Self {
        Self {
            time_column: time_column.into(),
            max_time_difference,
        }
    }

    /// Returns the input timestamp column used for retention and late-row checks.
    #[must_use]
    pub fn time_column(&self) -> &str {
        &self.time_column
    }

    /// Returns the maximum allowed timestamp distance between joined rows.
    #[must_use]
    pub const fn max_time_difference(&self) -> Duration {
        self.max_time_difference
    }
}

impl Default for StreamingJoinWindow {
    fn default() -> Self {
        Self::new("date_time", Duration::from_secs(10))
    }
}

/// Hard limits for unbounded streaming equi-joins.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StreamingJoinStateLimits {
    max_rows_per_side: usize,
    max_total_rows: usize,
}

impl StreamingJoinStateLimits {
    /// Creates hard state limits for a bounded-state join.
    #[must_use]
    pub const fn new(max_rows_per_side: usize, max_total_rows: usize) -> Self {
        Self {
            max_rows_per_side,
            max_total_rows,
        }
    }

    /// Maximum rows retained on either input side.
    #[must_use]
    pub const fn max_rows_per_side(self) -> usize {
        self.max_rows_per_side
    }

    /// Maximum rows retained across both input sides.
    #[must_use]
    pub const fn max_total_rows(self) -> usize {
        self.max_total_rows
    }
}

/// Counters maintained by streaming join stages.
#[derive(Clone, Default)]
pub struct StreamingJoinMetrics {
    state_rows: Arc<AtomicU64>,
    evicted_rows: Arc<AtomicU64>,
    late_dropped_rows: Arc<AtomicU64>,
}

impl StreamingJoinMetrics {
    pub(crate) fn new(
        state_rows: Arc<AtomicU64>,
        evicted_rows: Arc<AtomicU64>,
        late_dropped_rows: Arc<AtomicU64>,
    ) -> Self {
        Self {
            state_rows,
            evicted_rows,
            late_dropped_rows,
        }
    }

    /// Rows currently retained in join state.
    #[must_use]
    pub fn state_rows(&self) -> u64 {
        self.state_rows.load(Ordering::Relaxed)
    }

    /// Rows evicted from join state after watermark progress.
    #[must_use]
    pub fn evicted_rows(&self) -> u64 {
        self.evicted_rows.load(Ordering::Relaxed)
    }

    /// Rows dropped because they arrived at or behind their input watermark.
    #[must_use]
    pub fn late_dropped_rows(&self) -> u64 {
        self.late_dropped_rows.load(Ordering::Relaxed)
    }

    fn set_state_rows(&self, rows: usize) {
        self.state_rows.store(rows as u64, Ordering::Relaxed);
    }

    fn record_evictions(&self, rows: usize) {
        self.evicted_rows.fetch_add(rows as u64, Ordering::Relaxed);
    }

    fn record_late_row(&self) {
        self.late_dropped_rows.fetch_add(1, Ordering::Relaxed);
    }
}

impl fmt::Debug for StreamingJoinMetrics {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("StreamingJoinMetrics")
            .field("state_rows", &self.state_rows())
            .field("evicted_rows", &self.evicted_rows())
            .field("late_dropped_rows", &self.late_dropped_rows())
            .finish()
    }
}

pub(crate) fn streaming_hash_join_source(
    left: Source<SqlEvent<RecordBatch>>,
    right: Source<SqlEvent<RecordBatch>>,
    join: &HashJoinExec,
    config: &StreamingJoinConfig,
    metrics: StreamingJoinMetrics,
) -> Result<Source<SqlEvent<RecordBatch>>> {
    let stage = StreamingJoinStage::try_new(join, config, metrics)?;
    let tagged_left = left.map(TaggedJoinEvent::Left);
    let tagged_right = right.map(TaggedJoinEvent::Right);
    Ok(tagged_left
        .merge_all([tagged_right], false)
        .try_stateful_map_concat(stage, |stage, event| stage.apply(event)))
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum JoinInputSide {
    Left,
    Right,
}

enum TaggedJoinEvent {
    Left(SqlEvent<RecordBatch>),
    Right(SqlEvent<RecordBatch>),
}

#[derive(Clone)]
struct StreamingJoinStage {
    plan: Arc<StreamingJoinPlan>,
    metrics: StreamingJoinMetrics,
    latest_left_watermark_ns: Option<i64>,
    latest_right_watermark_ns: Option<i64>,
    last_emitted_watermark_ns: Option<i64>,
    left: JoinSideState,
    right: JoinSideState,
}

impl StreamingJoinStage {
    fn try_new(
        join: &HashJoinExec,
        config: &StreamingJoinConfig,
        metrics: StreamingJoinMetrics,
    ) -> Result<Self> {
        Ok(Self {
            plan: Arc::new(StreamingJoinPlan::try_new(join, config)?),
            metrics,
            latest_left_watermark_ns: None,
            latest_right_watermark_ns: None,
            last_emitted_watermark_ns: None,
            left: JoinSideState::default(),
            right: JoinSideState::default(),
        })
    }

    fn apply(&mut self, event: TaggedJoinEvent) -> StreamResult<Vec<SqlEvent<RecordBatch>>> {
        let result = match event {
            TaggedJoinEvent::Left(SqlEvent::Data(batch)) => {
                self.apply_batch(JoinInputSide::Left, batch)
            }
            TaggedJoinEvent::Right(SqlEvent::Data(batch)) => {
                self.apply_batch(JoinInputSide::Right, batch)
            }
            TaggedJoinEvent::Left(SqlEvent::Watermark(watermark)) => {
                self.apply_watermark(JoinInputSide::Left, watermark)
            }
            TaggedJoinEvent::Right(SqlEvent::Watermark(watermark)) => {
                self.apply_watermark(JoinInputSide::Right, watermark)
            }
            TaggedJoinEvent::Left(SqlEvent::Barrier(barrier))
            | TaggedJoinEvent::Right(SqlEvent::Barrier(barrier)) => {
                Ok(vec![SqlEvent::Barrier(barrier)])
            }
        };
        result.map_err(stream_error)
    }

    fn apply_batch(
        &mut self,
        side: JoinInputSide,
        batch: RecordBatch,
    ) -> Result<Vec<SqlEvent<RecordBatch>>> {
        if batch.num_rows() == 0 {
            return Ok(Vec::new());
        }

        let prepared = self.plan.prepare_batch(side, &batch)?;
        let mut output_rows = Vec::new();
        for row in 0..batch.num_rows() {
            let event_time_ns = self.plan.event_time_ns(side, &prepared, row)?;
            if self.row_is_late(side, event_time_ns) {
                self.metrics.record_late_row();
                continue;
            }
            let Some(key) = self.plan.key_at(side, &prepared, row)? else {
                continue;
            };
            let state_row = RowState::from_batch(&batch, row, event_time_ns)?;

            match side {
                JoinInputSide::Left => {
                    if let Some(candidates) = self.right.rows.get(&key) {
                        for right_row in candidates {
                            if self.plan.rows_can_join(&state_row, right_row)?
                                && self.plan.pair_passes_filter(&state_row, right_row)?
                            {
                                output_rows.push(self.plan.output_row(&state_row, right_row));
                            }
                        }
                    }
                    self.check_state_limit(JoinInputSide::Left)?;
                    self.left.insert(key, state_row);
                }
                JoinInputSide::Right => {
                    if let Some(candidates) = self.left.rows.get(&key) {
                        for left_row in candidates {
                            if self.plan.rows_can_join(left_row, &state_row)?
                                && self.plan.pair_passes_filter(left_row, &state_row)?
                            {
                                output_rows.push(self.plan.output_row(left_row, &state_row));
                            }
                        }
                    }
                    self.check_state_limit(JoinInputSide::Right)?;
                    self.right.insert(key, state_row);
                }
            }
        }

        self.update_state_rows_metric();
        if output_rows.is_empty() {
            return Ok(Vec::new());
        }
        Ok(vec![SqlEvent::Data(
            self.plan.build_output_batch(output_rows)?,
        )])
    }

    fn apply_watermark(
        &mut self,
        side: JoinInputSide,
        watermark: Watermark,
    ) -> Result<Vec<SqlEvent<RecordBatch>>> {
        let watermark_ns = watermark.timestamp_ns();
        match side {
            JoinInputSide::Left => {
                self.latest_left_watermark_ns = Some(
                    self.latest_left_watermark_ns
                        .map_or(watermark_ns, |current| current.max(watermark_ns)),
                );
            }
            JoinInputSide::Right => {
                self.latest_right_watermark_ns = Some(
                    self.latest_right_watermark_ns
                        .map_or(watermark_ns, |current| current.max(watermark_ns)),
                );
            }
        }

        let Some(min_watermark_ns) = self.min_input_watermark() else {
            return Ok(Vec::new());
        };
        if self
            .last_emitted_watermark_ns
            .is_some_and(|last| min_watermark_ns <= last)
        {
            return Ok(Vec::new());
        }

        self.evict_for_watermark(min_watermark_ns)?;
        self.last_emitted_watermark_ns = Some(min_watermark_ns);
        Ok(vec![SqlEvent::Watermark(Watermark::new(min_watermark_ns))])
    }

    fn row_is_late(&self, side: JoinInputSide, event_time_ns: Option<i64>) -> bool {
        let Some(event_time_ns) = event_time_ns else {
            return false;
        };
        match side {
            JoinInputSide::Left => self
                .latest_left_watermark_ns
                .is_some_and(|watermark_ns| event_time_ns <= watermark_ns),
            JoinInputSide::Right => self
                .latest_right_watermark_ns
                .is_some_and(|watermark_ns| event_time_ns <= watermark_ns),
        }
    }

    fn min_input_watermark(&self) -> Option<i64> {
        Some(
            self.latest_left_watermark_ns?
                .min(self.latest_right_watermark_ns?),
        )
    }

    fn evict_for_watermark(&mut self, watermark_ns: i64) -> Result<()> {
        let Some(max_difference_ns) = self.plan.max_time_difference_ns() else {
            return Ok(());
        };
        let left_evicted = self.left.evict_windowed(watermark_ns, max_difference_ns)?;
        let right_evicted = self.right.evict_windowed(watermark_ns, max_difference_ns)?;
        let evicted = left_evicted + right_evicted;
        if evicted > 0 {
            self.metrics.record_evictions(evicted);
            self.update_state_rows_metric();
        }
        Ok(())
    }

    fn check_state_limit(&self, side: JoinInputSide) -> Result<()> {
        let Some(limits) = self.plan.state_limits() else {
            return Ok(());
        };
        let next_side_rows = match side {
            JoinInputSide::Left => self.left.row_count + 1,
            JoinInputSide::Right => self.right.row_count + 1,
        };
        let next_total_rows = self.left.row_count + self.right.row_count + 1;
        if next_side_rows > limits.max_rows_per_side() {
            return Err(DataFusionError::Plan(format!(
                "streaming join state limit exceeded: {side:?} side would hold {next_side_rows} rows, limit is {}",
                limits.max_rows_per_side()
            )));
        }
        if next_total_rows > limits.max_total_rows() {
            return Err(DataFusionError::Plan(format!(
                "streaming join state limit exceeded: total state would hold {next_total_rows} rows, limit is {}",
                limits.max_total_rows()
            )));
        }
        Ok(())
    }

    fn update_state_rows_metric(&self) {
        self.metrics
            .set_state_rows(self.left.row_count + self.right.row_count);
    }
}

#[derive(Clone)]
struct StreamingJoinPlan {
    join_schema: SchemaRef,
    output_schema: SchemaRef,
    left_key_exprs: Vec<Arc<dyn PhysicalExpr>>,
    right_key_exprs: Vec<Arc<dyn PhysicalExpr>>,
    filter: Option<JoinFilter>,
    projection: Option<Arc<[usize]>>,
    null_equality: NullEquality,
    mode: StreamingJoinPlanMode,
}

impl StreamingJoinPlan {
    fn try_new(join: &HashJoinExec, config: &StreamingJoinConfig) -> Result<Self> {
        if *join.join_type() != JoinType::Inner {
            return Err(DataFusionError::NotImplemented(format!(
                "datum-sql streaming joins support INNER equi-joins only for now, found {:?}",
                join.join_type()
            )));
        }
        if join.on().is_empty() {
            return Err(DataFusionError::NotImplemented(
                "datum-sql streaming joins require at least one equi-join key".into(),
            ));
        }
        if join.null_aware {
            return Err(DataFusionError::NotImplemented(
                "datum-sql streaming joins do not support null-aware anti joins".into(),
            ));
        }

        let left_schema = join.left().schema();
        let right_schema = join.right().schema();
        let mode = StreamingJoinPlanMode::try_new(config, &left_schema, &right_schema)?;
        let output_schema = join.schema();
        Ok(Self {
            join_schema: Arc::clone(join.join_schema()),
            output_schema,
            left_key_exprs: join
                .on()
                .iter()
                .map(|(left, _right)| Arc::clone(left))
                .collect(),
            right_key_exprs: join
                .on()
                .iter()
                .map(|(_left, right)| Arc::clone(right))
                .collect(),
            filter: join.filter().cloned(),
            projection: join.projection.clone(),
            null_equality: join.null_equality(),
            mode,
        })
    }

    fn prepare_batch(&self, side: JoinInputSide, batch: &RecordBatch) -> Result<PreparedJoinBatch> {
        let key_exprs = match side {
            JoinInputSide::Left => &self.left_key_exprs,
            JoinInputSide::Right => &self.right_key_exprs,
        };
        let key_values = key_exprs
            .iter()
            .map(|expr| expr.evaluate(batch)?.into_array(batch.num_rows()))
            .collect::<Result<Vec<_>>>()?;
        let event_times = match (&self.mode, side) {
            (
                StreamingJoinPlanMode::Windowed {
                    left_time_index, ..
                },
                JoinInputSide::Left,
            ) => Some(Arc::clone(batch.column(*left_time_index))),
            (
                StreamingJoinPlanMode::Windowed {
                    right_time_index, ..
                },
                JoinInputSide::Right,
            ) => Some(Arc::clone(batch.column(*right_time_index))),
            (StreamingJoinPlanMode::BoundedState { .. }, _) => None,
        };
        Ok(PreparedJoinBatch {
            key_values,
            event_times,
        })
    }

    fn key_at(
        &self,
        _side: JoinInputSide,
        prepared: &PreparedJoinBatch,
        row: usize,
    ) -> Result<Option<JoinKey>> {
        let mut key = Vec::with_capacity(prepared.key_values.len());
        let mut has_null = false;
        for array in &prepared.key_values {
            let value = ScalarValue::try_from_array(array.as_ref(), row)?;
            has_null |= value.is_null();
            key.push(value);
        }
        if has_null && self.null_equality == NullEquality::NullEqualsNothing {
            Ok(None)
        } else {
            Ok(Some(JoinKey(key)))
        }
    }

    fn event_time_ns(
        &self,
        side: JoinInputSide,
        prepared: &PreparedJoinBatch,
        row: usize,
    ) -> Result<Option<i64>> {
        let index = match (&self.mode, side) {
            (
                StreamingJoinPlanMode::Windowed {
                    left_time_index, ..
                },
                JoinInputSide::Left,
            ) => *left_time_index,
            (
                StreamingJoinPlanMode::Windowed {
                    right_time_index, ..
                },
                JoinInputSide::Right,
            ) => *right_time_index,
            (StreamingJoinPlanMode::BoundedState { .. }, _) => return Ok(None),
        };
        let event_times = prepared.event_times.as_ref().ok_or_else(|| {
            DataFusionError::Internal("windowed join prepared batch is missing event time".into())
        })?;
        timestamp_ns_from_array(event_times, row).map(Some).map_err(|error| {
            DataFusionError::Plan(format!(
                "streaming join event-time column at index {index} on {side:?} input failed: {error}"
            ))
        })
    }

    fn rows_can_join(&self, left: &RowState, right: &RowState) -> Result<bool> {
        let Some(max_difference_ns) = self.max_time_difference_ns() else {
            return Ok(true);
        };
        let left_time_ns = left.event_time_ns.ok_or_else(|| {
            DataFusionError::Internal("windowed join left row is missing event time".into())
        })?;
        let right_time_ns = right.event_time_ns.ok_or_else(|| {
            DataFusionError::Internal("windowed join right row is missing event time".into())
        })?;
        Ok(left_time_ns.abs_diff(right_time_ns) <= max_difference_ns as u64)
    }

    fn pair_passes_filter(&self, left: &RowState, right: &RowState) -> Result<bool> {
        let Some(filter) = &self.filter else {
            return Ok(true);
        };
        let values = filter
            .column_indices()
            .iter()
            .map(|column| match column.side {
                JoinSide::Left => Ok(left.values[column.index].clone()),
                JoinSide::Right => Ok(right.values[column.index].clone()),
                JoinSide::None => Err(DataFusionError::NotImplemented(
                    "datum-sql streaming join filters do not support unqualified join-side columns"
                        .into(),
                )),
            })
            .collect::<Result<Vec<_>>>()?;
        let arrays = values
            .into_iter()
            .map(|value| ScalarValue::iter_to_array([value]))
            .collect::<Result<Vec<_>>>()?;
        let batch = RecordBatch::try_new(Arc::clone(filter.schema()), arrays)?;
        let predicate = filter.expression().evaluate(&batch)?.into_array(1)?;
        let predicate = as_boolean_array(&predicate)?;
        Ok(predicate.is_valid(0) && predicate.value(0))
    }

    fn output_row(&self, left: &RowState, right: &RowState) -> Vec<ScalarValue> {
        let mut row = Vec::with_capacity(left.values.len() + right.values.len());
        row.extend(left.values.iter().cloned());
        row.extend(right.values.iter().cloned());
        if let Some(projection) = &self.projection {
            projection.iter().map(|index| row[*index].clone()).collect()
        } else {
            row
        }
    }

    fn build_output_batch(&self, rows: Vec<Vec<ScalarValue>>) -> Result<RecordBatch> {
        let schema = if self.projection.is_some() {
            &self.output_schema
        } else {
            &self.join_schema
        };
        let column_count = schema.fields().len();
        let mut columns = vec![Vec::with_capacity(rows.len()); column_count];
        for row in rows {
            if row.len() != column_count {
                return Err(DataFusionError::Internal(format!(
                    "streaming join produced {} values for {column_count} output columns",
                    row.len()
                )));
            }
            for (index, value) in row.into_iter().enumerate() {
                columns[index].push(value);
            }
        }
        let arrays = columns
            .into_iter()
            .map(ScalarValue::iter_to_array)
            .collect::<Result<Vec<_>>>()?;
        RecordBatch::try_new(Arc::clone(schema), arrays).map_err(DataFusionError::from)
    }

    fn max_time_difference_ns(&self) -> Option<i64> {
        match &self.mode {
            StreamingJoinPlanMode::Windowed {
                max_time_difference_ns,
                ..
            } => Some(*max_time_difference_ns),
            StreamingJoinPlanMode::BoundedState { .. } => None,
        }
    }

    fn state_limits(&self) -> Option<StreamingJoinStateLimits> {
        match self.mode {
            StreamingJoinPlanMode::Windowed { .. } => None,
            StreamingJoinPlanMode::BoundedState { limits } => Some(limits),
        }
    }
}

#[derive(Clone)]
enum StreamingJoinPlanMode {
    Windowed {
        left_time_index: usize,
        right_time_index: usize,
        max_time_difference_ns: i64,
    },
    BoundedState {
        limits: StreamingJoinStateLimits,
    },
}

impl StreamingJoinPlanMode {
    fn try_new(
        config: &StreamingJoinConfig,
        left_schema: &SchemaRef,
        right_schema: &SchemaRef,
    ) -> Result<Self> {
        match config.mode() {
            StreamingJoinMode::Windowed { window } => {
                let max_time_difference_ns = duration_to_ns(window.max_time_difference())?;
                if max_time_difference_ns <= 0 {
                    return Err(DataFusionError::Plan(format!(
                        "streaming join window must be positive, found {:?}",
                        window.max_time_difference()
                    )));
                }
                Ok(Self::Windowed {
                    left_time_index: resolve_timestamp_column(left_schema, window.time_column())?,
                    right_time_index: resolve_timestamp_column(right_schema, window.time_column())?,
                    max_time_difference_ns,
                })
            }
            StreamingJoinMode::BoundedState { limits } => {
                if limits.max_rows_per_side() == 0 || limits.max_total_rows() == 0 {
                    return Err(DataFusionError::Plan(
                        "streaming join bounded-state limits must be greater than zero".into(),
                    ));
                }
                if limits.max_total_rows() < limits.max_rows_per_side() {
                    return Err(DataFusionError::Plan(format!(
                        "streaming join max_total_rows {} must be at least max_rows_per_side {}",
                        limits.max_total_rows(),
                        limits.max_rows_per_side()
                    )));
                }
                Ok(Self::BoundedState { limits: *limits })
            }
        }
    }
}

struct PreparedJoinBatch {
    key_values: Vec<ArrayRef>,
    event_times: Option<ArrayRef>,
}

#[derive(Clone)]
struct RowState {
    values: Vec<ScalarValue>,
    event_time_ns: Option<i64>,
}

impl RowState {
    fn from_batch(batch: &RecordBatch, row: usize, event_time_ns: Option<i64>) -> Result<Self> {
        let values = batch
            .columns()
            .iter()
            .map(|array| ScalarValue::try_from_array(array.as_ref(), row))
            .collect::<Result<Vec<_>>>()?;
        Ok(Self {
            values,
            event_time_ns,
        })
    }
}

#[derive(Clone, Default)]
struct JoinSideState {
    rows: HashMap<JoinKey, Vec<RowState>>,
    row_count: usize,
}

impl JoinSideState {
    fn insert(&mut self, key: JoinKey, row: RowState) {
        self.rows.entry(key).or_default().push(row);
        self.row_count += 1;
    }

    fn evict_windowed(&mut self, watermark_ns: i64, max_difference_ns: i64) -> Result<usize> {
        let mut evicted = 0;
        self.rows.retain(|_key, rows| {
            let before = rows.len();
            rows.retain(|row| {
                row.event_time_ns
                    .and_then(|event_time_ns| event_time_ns.checked_add(max_difference_ns))
                    .is_none_or(|expires_ns| expires_ns > watermark_ns)
            });
            evicted += before - rows.len();
            !rows.is_empty()
        });
        self.row_count = self.row_count.checked_sub(evicted).ok_or_else(|| {
            DataFusionError::Internal("streaming join state row count underflowed".into())
        })?;
        Ok(evicted)
    }
}

#[derive(Clone, Eq)]
struct JoinKey(Vec<ScalarValue>);

impl PartialEq for JoinKey {
    fn eq(&self, other: &Self) -> bool {
        self.0 == other.0
    }
}

impl Hash for JoinKey {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.hash(state);
    }
}

fn resolve_timestamp_column(schema: &SchemaRef, name: &str) -> Result<usize> {
    let index = schema.index_of(name).map_err(|_| {
        DataFusionError::Plan(format!(
            "streaming windowed join requires timestamp column '{name}' on both inputs"
        ))
    })?;
    match schema.field(index).data_type() {
        DataType::Timestamp(_, _) => Ok(index),
        other => Err(DataFusionError::Plan(format!(
            "streaming windowed join column '{name}' must be Timestamp, found {other:?}"
        ))),
    }
}

fn timestamp_ns_from_array(array: &ArrayRef, row: usize) -> Result<i64> {
    match array.data_type() {
        DataType::Timestamp(TimeUnit::Second, _) => timestamp_value_ns::<TimestampSecondArray>(
            array.as_any().downcast_ref::<TimestampSecondArray>(),
            row,
            1_000_000_000,
        ),
        DataType::Timestamp(TimeUnit::Millisecond, _) => {
            timestamp_value_ns::<TimestampMillisecondArray>(
                array.as_any().downcast_ref::<TimestampMillisecondArray>(),
                row,
                1_000_000,
            )
        }
        DataType::Timestamp(TimeUnit::Microsecond, _) => {
            timestamp_value_ns::<TimestampMicrosecondArray>(
                array.as_any().downcast_ref::<TimestampMicrosecondArray>(),
                row,
                1_000,
            )
        }
        DataType::Timestamp(TimeUnit::Nanosecond, _) => {
            timestamp_value_ns::<TimestampNanosecondArray>(
                array.as_any().downcast_ref::<TimestampNanosecondArray>(),
                row,
                1,
            )
        }
        other => Err(DataFusionError::Plan(format!(
            "streaming join event-time expression must evaluate to Timestamp, found {other:?}"
        ))),
    }
}

fn timestamp_value_ns<T>(array: Option<&T>, row: usize, multiplier: i64) -> Result<i64>
where
    T: Array + TimestampArrayValue,
{
    let array =
        array.ok_or_else(|| DataFusionError::Internal("timestamp array type mismatch".into()))?;
    if array.is_null(row) {
        return Err(DataFusionError::Plan(format!(
            "streaming join event-time column contains null at row {row}"
        )));
    }
    array
        .timestamp_value(row)
        .checked_mul(multiplier)
        .ok_or_else(|| DataFusionError::Plan("timestamp overflowed i64 nanoseconds".into()))
}

trait TimestampArrayValue {
    fn timestamp_value(&self, row: usize) -> i64;
}

impl TimestampArrayValue for TimestampSecondArray {
    fn timestamp_value(&self, row: usize) -> i64 {
        self.value(row)
    }
}

impl TimestampArrayValue for TimestampMillisecondArray {
    fn timestamp_value(&self, row: usize) -> i64 {
        self.value(row)
    }
}

impl TimestampArrayValue for TimestampMicrosecondArray {
    fn timestamp_value(&self, row: usize) -> i64 {
        self.value(row)
    }
}

impl TimestampArrayValue for TimestampNanosecondArray {
    fn timestamp_value(&self, row: usize) -> i64 {
        self.value(row)
    }
}

fn duration_to_ns(duration: Duration) -> Result<i64> {
    i64::try_from(duration.as_nanos()).map_err(|_| {
        DataFusionError::Plan(format!(
            "streaming join window {duration:?} exceeds i64 nanoseconds"
        ))
    })
}