laminar_core/time/event_time/
mod.rs

1//! Event-time extraction from Arrow `RecordBatch` columns.
2//!
3//! Columns must be `Timestamp(_)` at any precision; non-millisecond
4//! precisions are cast to `Timestamp(Millisecond)` via the Arrow kernel.
5
6use arrow::array::{Array, TimestampMillisecondArray};
7use arrow::datatypes::{DataType, Schema};
8use arrow::record_batch::RecordBatch;
9
10use super::cast::cast_to_millis_array;
11
12/// Column identifier for timestamp field.
13#[derive(Debug, Clone)]
14pub enum TimestampField {
15    /// Column name (resolved on first use, cached).
16    Name(String),
17    /// Column index (most efficient).
18    Index(usize),
19}
20
21/// Multi-row extraction strategy.
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
23pub enum ExtractionMode {
24    /// Extract from first row — O(1), default.
25    #[default]
26    First,
27    /// Extract from last row — O(1).
28    Last,
29    /// Extract maximum timestamp — O(n).
30    Max,
31    /// Extract minimum timestamp — O(n).
32    Min,
33}
34
35/// Errors that can occur during event time extraction.
36#[derive(Debug, thiserror::Error)]
37pub enum EventTimeError {
38    /// Column not found in schema.
39    #[error("Column not found: {0}")]
40    ColumnNotFound(String),
41
42    /// Column index out of bounds.
43    #[error("Column index {index} out of bounds (batch has {num_columns} columns)")]
44    IndexOutOfBounds {
45        /// Requested index.
46        index: usize,
47        /// Number of columns in the batch.
48        num_columns: usize,
49    },
50
51    /// Column is not a `Timestamp(_)` type.
52    #[error("event-time column must be Timestamp(_), found {found}")]
53    IncompatibleType {
54        /// Actual type found.
55        found: String,
56    },
57
58    /// Null timestamp encountered at a non-null-tolerant position.
59    #[error("Null timestamp at row {row}")]
60    NullTimestamp {
61        /// Row index with the null value.
62        row: usize,
63    },
64
65    /// Empty batch provided.
66    #[error("Cannot extract timestamp from empty batch")]
67    EmptyBatch,
68
69    /// Arrow `cast` kernel failed.
70    #[error("Arrow cast to Timestamp(Millisecond) failed: {0}")]
71    CastFailed(String),
72}
73
74/// Extracts event timestamps from Arrow `RecordBatch` columns.
75#[derive(Debug)]
76pub struct EventTimeExtractor {
77    field: TimestampField,
78    mode: ExtractionMode,
79    cached_index: Option<usize>,
80}
81
82impl EventTimeExtractor {
83    /// Creates an extractor that looks up a column by name. The column
84    /// index is cached after the first extraction.
85    #[must_use]
86    pub fn from_column(name: &str) -> Self {
87        Self {
88            field: TimestampField::Name(name.to_string()),
89            mode: ExtractionMode::default(),
90            cached_index: None,
91        }
92    }
93
94    /// Creates an extractor that uses a column by index. Skips the name
95    /// lookup entirely.
96    #[must_use]
97    pub fn from_index(index: usize) -> Self {
98        Self {
99            field: TimestampField::Index(index),
100            mode: ExtractionMode::default(),
101            cached_index: Some(index),
102        }
103    }
104
105    /// Sets the extraction mode for multi-row batches.
106    #[must_use]
107    pub fn with_mode(mut self, mode: ExtractionMode) -> Self {
108        self.mode = mode;
109        self
110    }
111
112    /// Gets the configured extraction mode.
113    #[must_use]
114    pub fn mode(&self) -> ExtractionMode {
115        self.mode
116    }
117
118    /// Validates that the schema contains a `Timestamp(_)` column at the
119    /// configured position.
120    ///
121    /// # Errors
122    ///
123    /// Returns an error if the column is missing or not a `Timestamp(_)`.
124    pub fn validate_schema(&self, schema: &Schema) -> Result<(), EventTimeError> {
125        let (_, data_type) = self.resolve_column(schema)?;
126        if !matches!(data_type, DataType::Timestamp(_, _)) {
127            return Err(EventTimeError::IncompatibleType {
128                found: format!("{data_type:?}"),
129            });
130        }
131        Ok(())
132    }
133
134    /// Extracts the event timestamp from a batch, in epoch milliseconds.
135    ///
136    /// # Errors
137    ///
138    /// Returns an error if the batch is empty, the column is missing,
139    /// the column is not a `Timestamp(_)`, the value at the selected row
140    /// is null, or Arrow's cast kernel fails.
141    pub fn extract(&mut self, batch: &RecordBatch) -> Result<i64, EventTimeError> {
142        let milliseconds = self.extract_millis_array(batch)?;
143        self.extract_from_millis(&milliseconds)
144    }
145
146    /// Returns the configured event-time column converted to epoch milliseconds.
147    ///
148    /// # Errors
149    ///
150    /// Returns an error if the batch is empty, the column is missing, the
151    /// column is not a `Timestamp(_)`, or Arrow's cast kernel fails.
152    pub fn extract_millis_array(
153        &mut self,
154        batch: &RecordBatch,
155    ) -> Result<TimestampMillisecondArray, EventTimeError> {
156        if batch.num_rows() == 0 {
157            return Err(EventTimeError::EmptyBatch);
158        }
159        let index = self.get_column_index(batch.schema().as_ref())?;
160        let column = batch.column(index);
161        cast_to_millis_array(column.as_ref()).map_err(|e| {
162            if matches!(column.data_type(), DataType::Timestamp(_, _)) {
163                EventTimeError::CastFailed(e.0)
164            } else {
165                EventTimeError::IncompatibleType { found: e.0 }
166            }
167        })
168    }
169
170    fn get_column_index(&mut self, schema: &Schema) -> Result<usize, EventTimeError> {
171        if let Some(idx) = self.cached_index {
172            if idx < schema.fields().len() {
173                return Ok(idx);
174            }
175        }
176        let (index, _) = self.resolve_column(schema)?;
177        self.cached_index = Some(index);
178        Ok(index)
179    }
180
181    fn resolve_column<'a>(
182        &self,
183        schema: &'a Schema,
184    ) -> Result<(usize, &'a DataType), EventTimeError> {
185        match &self.field {
186            TimestampField::Name(name) => {
187                let index = schema
188                    .index_of(name)
189                    .map_err(|_| EventTimeError::ColumnNotFound(name.clone()))?;
190                Ok((index, schema.field(index).data_type()))
191            }
192            TimestampField::Index(index) => {
193                if *index >= schema.fields().len() {
194                    return Err(EventTimeError::IndexOutOfBounds {
195                        index: *index,
196                        num_columns: schema.fields().len(),
197                    });
198                }
199                Ok((*index, schema.field(*index).data_type()))
200            }
201        }
202    }
203
204    fn extract_from_millis(&self, ms: &TimestampMillisecondArray) -> Result<i64, EventTimeError> {
205        match self.mode {
206            ExtractionMode::First => read_indexed(ms, 0),
207            ExtractionMode::Last => read_indexed(ms, ms.len() - 1),
208            ExtractionMode::Max => fold_non_null(ms, i64::MIN, i64::max),
209            ExtractionMode::Min => fold_non_null(ms, i64::MAX, i64::min),
210        }
211    }
212}
213
214fn read_indexed(arr: &TimestampMillisecondArray, idx: usize) -> Result<i64, EventTimeError> {
215    if arr.is_null(idx) {
216        Err(EventTimeError::NullTimestamp { row: idx })
217    } else {
218        Ok(arr.value(idx))
219    }
220}
221
222fn fold_non_null<F>(arr: &TimestampMillisecondArray, init: i64, f: F) -> Result<i64, EventTimeError>
223where
224    F: Fn(i64, i64) -> i64,
225{
226    let mut out = init;
227    let mut found = false;
228    for i in 0..arr.len() {
229        if !arr.is_null(i) {
230            found = true;
231            out = f(out, arr.value(i));
232        }
233    }
234    if found {
235        Ok(out)
236    } else {
237        Err(EventTimeError::NullTimestamp { row: 0 })
238    }
239}
240
241#[cfg(test)]
242mod tests;
laminar_core/time/event_time/mod.rs

laminar_core/time/event_time/
mod.rs
