laminar_connectors/kafka/source/
record_metadata.rs

1//! Typed Kafka record metadata: the selectable columns appended after the decoded value columns,
2//! and the per-row values captured for them.
3
4use std::sync::Arc;
5
6use arrow_array::builder::{BinaryBuilder, ListBuilder, StringBuilder, StructBuilder};
7use arrow_array::{
8    ArrayRef, BinaryArray, Int32Array, Int64Array, StringArray, TimestampMillisecondArray,
9};
10use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef, TimeUnit};
11
12use crate::error::ConnectorError;
13
14/// One selectable metadata column, in canonical output order.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
16pub enum KafkaMetadataColumn {
17    /// `_key`: the record key bytes, NULL for a record without a key.
18    Key,
19    /// `_topic`: the source topic.
20    Topic,
21    /// `_partition`: the source partition.
22    Partition,
23    /// `_offset`: the record offset.
24    Offset,
25    /// `_timestamp`: the record timestamp in milliseconds, NULL when the broker reports none.
26    Timestamp,
27    /// `_timestamp_type`: `create_time`, `log_append_time` or `not_available`.
28    TimestampType,
29    /// `_headers`: every header in order, duplicates kept, values as nullable bytes.
30    Headers,
31}
32
33impl KafkaMetadataColumn {
34    const ALL: [Self; 7] = [
35        Self::Key,
36        Self::Topic,
37        Self::Partition,
38        Self::Offset,
39        Self::Timestamp,
40        Self::TimestampType,
41        Self::Headers,
42    ];
43
44    /// Output column name.
45    #[must_use]
46    pub const fn column_name(self) -> &'static str {
47        match self {
48            Self::Key => "_key",
49            Self::Topic => "_topic",
50            Self::Partition => "_partition",
51            Self::Offset => "_offset",
52            Self::Timestamp => "_timestamp",
53            Self::TimestampType => "_timestamp_type",
54            Self::Headers => "_headers",
55        }
56    }
57
58    const fn option_name(self) -> &'static str {
59        match self {
60            Self::Key => "key",
61            Self::Topic => "topic",
62            Self::Partition => "partition",
63            Self::Offset => "offset",
64            Self::Timestamp => "timestamp",
65            Self::TimestampType => "timestamp_type",
66            Self::Headers => "headers",
67        }
68    }
69
70    fn field(self) -> Field {
71        match self {
72            Self::Key => Field::new(self.column_name(), DataType::Binary, true),
73            Self::Topic | Self::TimestampType => {
74                Field::new(self.column_name(), DataType::Utf8, false)
75            }
76            Self::Partition => Field::new(self.column_name(), DataType::Int32, false),
77            Self::Offset => Field::new(self.column_name(), DataType::Int64, false),
78            Self::Timestamp => Field::new(
79                self.column_name(),
80                DataType::Timestamp(TimeUnit::Millisecond, None),
81                true,
82            ),
83            Self::Headers => Field::new(
84                self.column_name(),
85                DataType::List(Arc::new(Field::new_struct("item", header_fields(), false))),
86                false,
87            ),
88        }
89    }
90}
91
92fn header_fields() -> Fields {
93    Fields::from(vec![
94        Field::new("key", DataType::Utf8, false),
95        Field::new("value", DataType::Binary, true),
96    ])
97}
98
99/// The metadata columns a source appends, in canonical order without duplicates.
100#[derive(Debug, Clone, Default, PartialEq, Eq)]
101pub struct KafkaMetadataColumns(Vec<KafkaMetadataColumn>);
102
103impl KafkaMetadataColumns {
104    /// Parse a comma-separated `metadata.columns` list such as `key,offset,headers`.
105    ///
106    /// # Errors
107    /// Returns a configuration error for an unknown name.
108    pub fn parse(list: &str) -> Result<Self, ConnectorError> {
109        let mut columns = Vec::new();
110        for name in list
111            .split(',')
112            .map(str::trim)
113            .filter(|name| !name.is_empty())
114        {
115            let column = KafkaMetadataColumn::ALL
116                .into_iter()
117                .find(|column| column.option_name().eq_ignore_ascii_case(name))
118                .ok_or_else(|| {
119                    ConnectorError::ConfigurationError(format!(
120                        "unknown Kafka metadata column '{name}' (expected key, topic, partition, \
121                         offset, timestamp, timestamp_type or headers)"
122                    ))
123                })?;
124            columns.push(column);
125        }
126        Ok(Self::from_columns(columns))
127    }
128
129    /// Select `columns`, normalized to canonical order.
130    #[must_use]
131    pub fn from_columns(mut columns: Vec<KafkaMetadataColumn>) -> Self {
132        columns.sort_unstable();
133        columns.dedup();
134        Self(columns)
135    }
136
137    /// Whether `column` is selected.
138    #[must_use]
139    pub fn contains(&self, column: KafkaMetadataColumn) -> bool {
140        self.0.contains(&column)
141    }
142
143    /// Whether no metadata column is selected.
144    #[must_use]
145    pub fn is_empty(&self) -> bool {
146        self.0.is_empty()
147    }
148
149    /// Number of selected columns.
150    #[must_use]
151    pub fn len(&self) -> usize {
152        self.0.len()
153    }
154
155    /// The canonical `metadata.columns` option value.
156    #[must_use]
157    pub fn option_value(&self) -> String {
158        self.0
159            .iter()
160            .map(|column| column.option_name())
161            .collect::<Vec<_>>()
162            .join(",")
163    }
164
165    /// Reject a payload field that collides, case-insensitively, with a selected column.
166    ///
167    /// # Errors
168    /// Returns a schema mismatch naming the colliding field.
169    pub fn validate_payload(&self, payload: &Schema) -> Result<(), ConnectorError> {
170        let collision = payload.fields().iter().find(|field| {
171            self.0
172                .iter()
173                .any(|column| field.name().eq_ignore_ascii_case(column.column_name()))
174        });
175        match collision {
176            Some(field) => Err(ConnectorError::SchemaMismatch(format!(
177                "Kafka payload schema contains configured connector metadata column '{}'",
178                field.name()
179            ))),
180            None => Ok(()),
181        }
182    }
183
184    /// The payload schema followed by the selected metadata fields.
185    #[must_use]
186    pub fn output_schema(&self, payload: &SchemaRef) -> SchemaRef {
187        if self.0.is_empty() {
188            return Arc::clone(payload);
189        }
190        let mut fields = payload.fields().to_vec();
191        fields.extend(self.0.iter().map(|column| Arc::new(column.field())));
192        Arc::new(Schema::new_with_metadata(
193            fields,
194            payload.metadata().clone(),
195        ))
196    }
197
198    /// Build the selected metadata arrays for `rows`.
199    pub(super) fn arrays(
200        &self,
201        rows: &[KafkaRowMetadata],
202    ) -> Result<Vec<ArrayRef>, ConnectorError> {
203        self.0
204            .iter()
205            .map(|column| column_array(*column, rows))
206            .collect()
207    }
208}
209
210/// How the broker assigned a record's timestamp.
211#[derive(Debug, Clone, Copy, PartialEq, Eq)]
212pub(super) enum KafkaTimestampType {
213    NotAvailable,
214    CreateTime,
215    LogAppendTime,
216}
217
218impl KafkaTimestampType {
219    const fn as_str(self) -> &'static str {
220        match self {
221            Self::NotAvailable => "not_available",
222            Self::CreateTime => "create_time",
223            Self::LogAppendTime => "log_append_time",
224        }
225    }
226}
227
228/// One record header: a UTF-8 name and nullable binary value.
229#[derive(Debug, Clone, PartialEq, Eq)]
230pub(super) struct KafkaHeader {
231    pub(super) key: String,
232    pub(super) value: Option<Vec<u8>>,
233}
234
235/// Metadata captured for one emitted row. `key` and `headers` are populated only when selected.
236#[derive(Debug, Clone)]
237pub(super) struct KafkaRowMetadata {
238    pub(super) key: Option<Vec<u8>>,
239    pub(super) topic: Arc<str>,
240    pub(super) partition: i32,
241    pub(super) offset: i64,
242    pub(super) timestamp_ms: Option<i64>,
243    pub(super) timestamp_type: KafkaTimestampType,
244    pub(super) headers: Vec<KafkaHeader>,
245}
246
247fn column_array(
248    column: KafkaMetadataColumn,
249    rows: &[KafkaRowMetadata],
250) -> Result<ArrayRef, ConnectorError> {
251    let array: ArrayRef = match column {
252        KafkaMetadataColumn::Key => Arc::new(
253            rows.iter()
254                .map(|row| row.key.as_deref())
255                .collect::<BinaryArray>(),
256        ),
257        KafkaMetadataColumn::Topic => Arc::new(
258            rows.iter()
259                .map(|row| Some(row.topic.as_ref()))
260                .collect::<StringArray>(),
261        ),
262        KafkaMetadataColumn::Partition => {
263            Arc::new(rows.iter().map(|row| row.partition).collect::<Int32Array>())
264        }
265        KafkaMetadataColumn::Offset => {
266            Arc::new(rows.iter().map(|row| row.offset).collect::<Int64Array>())
267        }
268        KafkaMetadataColumn::Timestamp => Arc::new(
269            rows.iter()
270                .map(|row| row.timestamp_ms)
271                .collect::<TimestampMillisecondArray>(),
272        ),
273        KafkaMetadataColumn::TimestampType => Arc::new(
274            rows.iter()
275                .map(|row| Some(row.timestamp_type.as_str()))
276                .collect::<StringArray>(),
277        ),
278        KafkaMetadataColumn::Headers => headers_array(rows)?,
279    };
280    Ok(array)
281}
282
283fn headers_array(rows: &[KafkaRowMetadata]) -> Result<ArrayRef, ConnectorError> {
284    let entries = StructBuilder::new(
285        header_fields(),
286        vec![
287            Box::new(StringBuilder::new()),
288            Box::new(BinaryBuilder::new()),
289        ],
290    );
291    let item = Field::new_struct("item", header_fields(), false);
292    let mut list = ListBuilder::new(entries).with_field(Arc::new(item));
293    for row in rows {
294        let entries = list.values();
295        for header in &row.headers {
296            entries
297                .field_builder::<StringBuilder>(0)
298                .ok_or_else(|| ConnectorError::Internal("header key builder".into()))?
299                .append_value(&header.key);
300            entries
301                .field_builder::<BinaryBuilder>(1)
302                .ok_or_else(|| ConnectorError::Internal("header value builder".into()))?
303                .append_option(header.value.as_deref());
304            entries.append(true);
305        }
306        list.append(true);
307    }
308    Ok(Arc::new(list.finish()))
309}
310
311#[cfg(test)]
312mod tests {
313    use arrow_array::cast::AsArray;
314    use arrow_array::{Array, RecordBatch};
315
316    use super::*;
317
318    fn row(offset: i64, key: Option<&[u8]>, headers: Vec<KafkaHeader>) -> KafkaRowMetadata {
319        KafkaRowMetadata {
320            key: key.map(<[u8]>::to_vec),
321            topic: Arc::from("events"),
322            partition: 3,
323            offset,
324            timestamp_ms: Some(1_000 + offset),
325            timestamp_type: KafkaTimestampType::LogAppendTime,
326            headers,
327        }
328    }
329
330    #[test]
331    fn parse_normalizes_order_and_rejects_unknown_names() {
332        let columns = KafkaMetadataColumns::parse(" headers ,KEY,offset,key").unwrap();
333        assert_eq!(columns.option_value(), "key,offset,headers");
334        assert!(KafkaMetadataColumns::parse("partition,epoch").is_err());
335        assert!(KafkaMetadataColumns::parse("").unwrap().is_empty());
336    }
337
338    #[test]
339    fn collisions_are_case_insensitive() {
340        let columns = KafkaMetadataColumns::parse("topic").unwrap();
341        let payload = Schema::new(vec![Field::new("_TOPIC", DataType::Utf8, true)]);
342        assert!(columns.validate_payload(&payload).is_err());
343        let unselected = KafkaMetadataColumns::parse("offset").unwrap();
344        unselected.validate_payload(&payload).unwrap();
345    }
346
347    #[test]
348    fn arrays_preserve_null_and_empty_keys_and_binary_duplicate_headers() {
349        let columns = KafkaMetadataColumns::parse("key,topic,timestamp_type,headers").unwrap();
350        let rows = vec![
351            row(
352                10,
353                None,
354                vec![
355                    KafkaHeader {
356                        key: "h".into(),
357                        value: Some(vec![0xff, 0x00]),
358                    },
359                    KafkaHeader {
360                        key: "h".into(),
361                        value: None,
362                    },
363                    KafkaHeader {
364                        key: "e".into(),
365                        value: Some(Vec::new()),
366                    },
367                ],
368            ),
369            row(11, Some(b""), Vec::new()),
370        ];
371        let payload: SchemaRef = Arc::new(Schema::empty());
372        let batch = RecordBatch::try_new_with_options(
373            columns.output_schema(&payload),
374            columns.arrays(&rows).unwrap(),
375            &arrow_array::RecordBatchOptions::new().with_row_count(Some(2)),
376        )
377        .unwrap();
378
379        let keys = batch.column(0).as_binary::<i32>();
380        assert!(keys.is_null(0));
381        assert_eq!(keys.value(1), b"");
382        assert_eq!(
383            batch.column(2).as_string::<i32>().value(0),
384            "log_append_time"
385        );
386
387        let headers = batch.column(3).as_list::<i32>();
388        let first = headers.value(0);
389        let first = first.as_struct();
390        let names = first.column(0).as_string::<i32>();
391        let values = first.column(1).as_binary::<i32>();
392        assert_eq!((names.value(0), values.value(0)), ("h", &[0xff, 0x00][..]));
393        assert_eq!(names.value(1), "h");
394        assert!(values.is_null(1));
395        assert_eq!((names.value(2), values.value(2)), ("e", &b""[..]));
396        assert_eq!(headers.value(1).len(), 0);
397        assert!(!headers.is_null(1));
398    }
399}
laminar_connectors/kafka/source/record_metadata.rs

laminar_connectors/kafka/source/
record_metadata.rs
