codecs/encoding/format/
arrow.rs

1//! Arrow IPC streaming format codec for batched event encoding
2//!
3//! Provides Apache Arrow IPC stream format encoding with static schema support.
4//! This implements the streaming variant of the Arrow IPC protocol, which writes
5//! a continuous stream of record batches without a file footer.
6
7use arrow::{
8    datatypes::{DataType, Field, Fields, Schema, SchemaRef},
9    error::ArrowError,
10    ipc::writer::StreamWriter,
11    json::reader::ReaderBuilder,
12    record_batch::RecordBatch,
13};
14use async_trait::async_trait;
15use bytes::{BufMut, Bytes, BytesMut};
16use snafu::{ResultExt, Snafu, ensure};
17use vector_config::configurable_component;
18use vector_core::event::Event;
19
20/// Provides Arrow schema for encoding.
21///
22/// Sinks can implement this trait to provide custom schema fetching logic.
23#[async_trait]
24pub trait SchemaProvider: Send + Sync + std::fmt::Debug {
25    /// Fetch the Arrow schema from the data store.
26    ///
27    /// This is called during sink configuration build phase to fetch
28    /// the schema once at startup, rather than at runtime.
29    async fn get_schema(&self) -> Result<Schema, ArrowEncodingError>;
30}
31
32/// Configuration for Arrow IPC stream serialization
33#[configurable_component]
34#[derive(Clone, Default)]
35pub struct ArrowStreamSerializerConfig {
36    /// The Arrow schema to use for encoding
37    #[serde(skip)]
38    #[configurable(derived)]
39    pub schema: Option<arrow::datatypes::Schema>,
40
41    /// Allow null values for non-nullable fields in the schema.
42    ///
43    /// When enabled, missing or incompatible values are encoded as null, even for fields
44    /// marked as non-nullable in the Arrow schema. This is useful when working with downstream
45    /// systems that can handle null values through defaults, computed columns, or other mechanisms.
46    ///
47    /// When disabled (default), missing values for non-nullable fields results in encoding errors. This is to
48    /// help ensure all required data is present before sending it to the sink.
49    #[serde(default)]
50    #[configurable(derived)]
51    pub allow_nullable_fields: bool,
52}
53
54impl std::fmt::Debug for ArrowStreamSerializerConfig {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        f.debug_struct("ArrowStreamSerializerConfig")
57            .field(
58                "schema",
59                &self
60                    .schema
61                    .as_ref()
62                    .map(|s| format!("{} fields", s.fields().len())),
63            )
64            .field("allow_nullable_fields", &self.allow_nullable_fields)
65            .finish()
66    }
67}
68
69impl ArrowStreamSerializerConfig {
70    /// Create a new ArrowStreamSerializerConfig with a schema
71    pub fn new(schema: arrow::datatypes::Schema) -> Self {
72        Self {
73            schema: Some(schema),
74            allow_nullable_fields: false,
75        }
76    }
77
78    /// The data type of events that are accepted by `ArrowStreamEncoder`.
79    pub fn input_type(&self) -> vector_core::config::DataType {
80        vector_core::config::DataType::Log
81    }
82
83    /// The schema required by the serializer.
84    pub fn schema_requirement(&self) -> vector_core::schema::Requirement {
85        vector_core::schema::Requirement::empty()
86    }
87}
88
89/// Arrow IPC stream batch serializer that holds the schema
90#[derive(Clone, Debug)]
91pub struct ArrowStreamSerializer {
92    schema: SchemaRef,
93}
94
95impl ArrowStreamSerializer {
96    /// Encode events into a `RecordBatch` without writing to IPC stream format.
97    pub fn encode_to_record_batch(
98        &self,
99        events: &[Event],
100    ) -> Result<RecordBatch, ArrowEncodingError> {
101        let values = vector_log_events_to_json_values(events).map_err(|e| {
102            ArrowEncodingError::RecordBatchCreation {
103                source: arrow::error::ArrowError::JsonError(e.to_string()),
104            }
105        })?;
106        build_record_batch(self.schema.clone(), &values)
107    }
108
109    /// Create a new ArrowStreamSerializer with the given configuration
110    pub fn new(config: ArrowStreamSerializerConfig) -> Result<Self, ArrowEncodingError> {
111        let schema = config.schema.ok_or(ArrowEncodingError::MissingSchema)?;
112
113        // If allow_nullable_fields is enabled, transform the schema once here
114        // instead of on every batch encoding
115        let schema = if config.allow_nullable_fields {
116            let nullable_fields: Fields = schema
117                .fields()
118                .iter()
119                .map(|f| make_field_nullable(f))
120                .collect::<Result<Vec<_>, _>>()?
121                .into();
122            Schema::new_with_metadata(nullable_fields, schema.metadata().clone())
123        } else {
124            schema
125        };
126
127        Ok(Self {
128            schema: SchemaRef::new(schema),
129        })
130    }
131}
132
133impl tokio_util::codec::Encoder<Vec<Event>> for ArrowStreamSerializer {
134    type Error = ArrowEncodingError;
135
136    fn encode(&mut self, events: Vec<Event>, buffer: &mut BytesMut) -> Result<(), Self::Error> {
137        if events.is_empty() {
138            return Err(ArrowEncodingError::NoEvents);
139        }
140
141        let bytes = encode_events_to_arrow_ipc_stream(&events, self.schema.clone())?;
142
143        buffer.extend_from_slice(&bytes);
144        Ok(())
145    }
146}
147
148/// Errors that can occur during Arrow encoding
149#[derive(Debug, Snafu)]
150pub enum ArrowEncodingError {
151    /// Failed to create Arrow record batch
152    #[snafu(display("Failed to create Arrow record batch: {source}"))]
153    RecordBatchCreation {
154        /// The underlying Arrow error
155        source: arrow::error::ArrowError,
156    },
157
158    /// Failed to write Arrow IPC data
159    #[snafu(display("Failed to write Arrow IPC data: {source}"))]
160    IpcWrite {
161        /// The underlying Arrow error
162        source: arrow::error::ArrowError,
163    },
164
165    /// No events provided for encoding
166    #[snafu(display("No events provided for encoding"))]
167    NoEvents,
168
169    /// Failed to fetch schema from provider
170    #[snafu(display("Failed to fetch schema from provider: {message}"))]
171    SchemaFetchError {
172        /// Error message from the provider
173        message: String,
174    },
175
176    /// Null value encountered for non-nullable field
177    #[snafu(display("Null value for non-nullable field '{field_name}'"))]
178    NullConstraint {
179        /// The field name
180        field_name: String,
181    },
182
183    /// Arrow serializer requires a schema
184    #[snafu(display("Arrow serializer requires a schema"))]
185    MissingSchema,
186
187    /// IO error during encoding
188    #[snafu(display("IO error: {source}"), context(false))]
189    Io {
190        /// The underlying IO error
191        source: std::io::Error,
192    },
193
194    /// Arrow JSON decoding error
195    #[snafu(display("Arrow JSON decoding error: {source}"))]
196    ArrowJsonDecode {
197        /// The underlying Arrow error
198        source: arrow::error::ArrowError,
199    },
200
201    /// Invalid Map schema structure
202    #[snafu(display("Invalid Map schema for field '{field_name}': {reason}"))]
203    InvalidMapSchema {
204        /// The field name
205        field_name: String,
206        /// Description of the schema violation
207        reason: String,
208    },
209}
210
211/// Encodes a batch of events into Arrow IPC streaming format
212pub fn encode_events_to_arrow_ipc_stream(
213    events: &[Event],
214    schema: SchemaRef,
215) -> Result<Bytes, ArrowEncodingError> {
216    if events.is_empty() {
217        return Err(ArrowEncodingError::NoEvents);
218    }
219
220    let json_values = vector_log_events_to_json_values(events).map_err(|e| {
221        ArrowEncodingError::RecordBatchCreation {
222            source: ArrowError::JsonError(e.to_string()),
223        }
224    })?;
225
226    let record_batch = build_record_batch(schema, &json_values)?;
227
228    let mut buffer = BytesMut::new().writer();
229    let mut writer =
230        StreamWriter::try_new(&mut buffer, record_batch.schema_ref()).context(IpcWriteSnafu)?;
231    writer.write(&record_batch).context(IpcWriteSnafu)?;
232    writer.finish().context(IpcWriteSnafu)?;
233
234    Ok(buffer.into_inner().freeze())
235}
236
237/// Recursively makes a Field and all its nested fields nullable
238fn make_field_nullable(field: &Field) -> Result<Field, ArrowEncodingError> {
239    let new_data_type = match field.data_type() {
240        DataType::List(inner_field) => DataType::List(make_field_nullable(inner_field)?.into()),
241        DataType::Struct(fields) => DataType::Struct(
242            fields
243                .iter()
244                .map(|f| make_field_nullable(f))
245                .collect::<Result<Vec<_>, _>>()?
246                .into(),
247        ),
248        DataType::Map(inner, sorted) => {
249            // A Map's inner field is a "entries" Struct<Key, Value>
250            let DataType::Struct(fields) = inner.data_type() else {
251                return InvalidMapSchemaSnafu {
252                    field_name: field.name(),
253                    reason: format!("inner type must be Struct, found {:?}", inner.data_type()),
254                }
255                .fail();
256            };
257
258            ensure!(
259                fields.len() == 2,
260                InvalidMapSchemaSnafu {
261                    field_name: field.name(),
262                    reason: format!("expected 2 fields (key, value), found {}", fields.len()),
263                },
264            );
265            let key_field = &fields[0];
266            let value_field = &fields[1];
267
268            let new_struct_fields: Fields =
269                [key_field.clone(), make_field_nullable(value_field)?.into()].into();
270
271            // Reconstruct the inner "entries" field
272            // The inner field itself must be non-nullable (only the Map wrapper is nullable)
273            let new_inner_field = inner
274                .as_ref()
275                .clone()
276                .with_data_type(DataType::Struct(new_struct_fields))
277                .with_nullable(false);
278
279            DataType::Map(new_inner_field.into(), *sorted)
280        }
281        other => other.clone(),
282    };
283
284    Ok(field
285        .clone()
286        .with_data_type(new_data_type)
287        .with_nullable(true))
288}
289
290/// Returns true if the field is absent from the value's object map, or explicitly null.
291/// Find non-nullable schema fields that are missing or null in any of the given events.
292pub fn find_null_non_nullable_fields<'a>(
293    schema: &'a Schema,
294    values: &[serde_json::Value],
295) -> Vec<&'a str> {
296    schema
297        .fields()
298        .iter()
299        .filter(|field| {
300            !field.is_nullable()
301                && values.iter().any(|value| {
302                    value
303                        .as_object()
304                        .and_then(|map| map.get(field.name().as_str()))
305                        .is_none_or(serde_json::Value::is_null)
306                })
307        })
308        .map(|field| field.name().as_str())
309        .collect()
310}
311
312pub(crate) fn vector_log_events_to_json_values(
313    events: &[Event],
314) -> Result<Vec<serde_json::Value>, serde_json::Error> {
315    events
316        .iter()
317        .filter_map(Event::maybe_as_log)
318        .map(serde_json::to_value)
319        .collect()
320}
321
322/// Build an Arrow RecordBatch from a slice of events using the provided schema.
323pub(crate) fn build_record_batch(
324    schema: SchemaRef,
325    values: &[serde_json::Value],
326) -> Result<RecordBatch, ArrowEncodingError> {
327    if values.is_empty() {
328        return Err(ArrowEncodingError::NoEvents);
329    }
330
331    let missing = find_null_non_nullable_fields(&schema, values);
332    if !missing.is_empty() {
333        let error: vector_common::Error = Box::new(ArrowEncodingError::NullConstraint {
334            field_name: missing.join(", "),
335        });
336        vector_common::internal_event::emit(crate::internal_events::EncoderNullConstraintError {
337            error: &error,
338        });
339        return Err(ArrowEncodingError::NullConstraint {
340            field_name: missing.join(", "),
341        });
342    }
343
344    let mut decoder = ReaderBuilder::new(schema)
345        .build_decoder()
346        .inspect_err(|e| {
347            vector_common::internal_event::emit(crate::internal_events::EncoderRecordBatchError {
348                error: e,
349                error_code: "arrow_record_batch_creation",
350            });
351        })
352        .context(RecordBatchCreationSnafu)?;
353
354    decoder
355        .serialize(values)
356        .inspect_err(|e| {
357            vector_common::internal_event::emit(crate::internal_events::EncoderRecordBatchError {
358                error: e,
359                error_code: "arrow_json_decode",
360            });
361        })
362        .context(ArrowJsonDecodeSnafu)?;
363
364    decoder
365        .flush()
366        .inspect_err(|e| {
367            vector_common::internal_event::emit(crate::internal_events::EncoderRecordBatchError {
368                error: e,
369                error_code: "arrow_json_decode",
370            });
371        })
372        .context(ArrowJsonDecodeSnafu)?
373        .ok_or(ArrowEncodingError::NoEvents)
374}
375
376#[cfg(test)]
377mod tests {
378    use super::*;
379    use arrow::{
380        array::{Array, AsArray},
381        datatypes::TimeUnit,
382        ipc::reader::StreamReader,
383    };
384    use chrono::Utc;
385    use std::io::Cursor;
386    use vector_core::event::{LogEvent, Value};
387    use vrl::event_path;
388
389    /// Helper to encode events and return the decoded RecordBatch
390    fn encode_and_decode(
391        events: Vec<Event>,
392        schema: SchemaRef,
393    ) -> Result<RecordBatch, Box<dyn std::error::Error>> {
394        let bytes = encode_events_to_arrow_ipc_stream(&events, schema.clone())?;
395        let cursor = Cursor::new(bytes);
396        let mut reader = StreamReader::try_new(cursor, None)?;
397        Ok(reader.next().unwrap()?)
398    }
399
400    /// Create a simple event from key-value pairs
401    fn create_event<V>(fields: Vec<(&str, V)>) -> Event
402    where
403        V: Into<Value>,
404    {
405        let mut log = LogEvent::default();
406        for (key, value) in fields {
407            log.insert(&vrl::path::parse_target_path(key).unwrap(), value.into());
408        }
409        Event::Log(log)
410    }
411
412    mod comprehensive {
413        use super::*;
414
415        #[test]
416        fn test_encode_all_types() {
417            use arrow::datatypes::{
418                Decimal128Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type,
419                Int64Type, TimestampMillisecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
420            };
421            use vrl::value::ObjectMap;
422
423            let now = Utc::now();
424
425            // Create a struct (tuple) value with unnamed fields
426            let mut tuple_value = ObjectMap::new();
427            tuple_value.insert("f0".into(), Value::Bytes("nested_str".into()));
428            tuple_value.insert("f1".into(), Value::Integer(999));
429
430            // Create a named struct (named tuple) value
431            let mut named_tuple_value = ObjectMap::new();
432            named_tuple_value.insert("category".into(), Value::Bytes("test_category".into()));
433            named_tuple_value.insert("tag".into(), Value::Bytes("test_tag".into()));
434
435            // Create a list value
436            let list_value = Value::Array(vec![
437                Value::Integer(1),
438                Value::Integer(2),
439                Value::Integer(3),
440            ]);
441
442            // Create a map value
443            let mut map_value = ObjectMap::new();
444            map_value.insert("key1".into(), Value::Integer(100));
445            map_value.insert("key2".into(), Value::Integer(200));
446
447            let mut log = LogEvent::default();
448            // Primitive types
449            log.insert(event_path!("string_field"), "test");
450            log.insert(event_path!("int8_field"), 127);
451            log.insert(event_path!("int16_field"), 32000);
452            log.insert(event_path!("int32_field"), 1000000);
453            log.insert(event_path!("int64_field"), 42);
454            log.insert(event_path!("uint8_field"), 255);
455            log.insert(event_path!("uint16_field"), 65535);
456            log.insert(event_path!("uint32_field"), 4000000);
457            log.insert(event_path!("uint64_field"), 9000000000_i64);
458            log.insert(event_path!("float32_field"), 3.15);
459            log.insert(event_path!("float64_field"), 3.15);
460            log.insert(event_path!("bool_field"), true);
461            log.insert(event_path!("timestamp_field"), now);
462            log.insert(event_path!("decimal_field"), 99.99);
463            // Complex types
464            log.insert(event_path!("list_field"), list_value);
465            log.insert(event_path!("struct_field"), Value::Object(tuple_value));
466            log.insert(
467                event_path!("named_struct_field"),
468                Value::Object(named_tuple_value),
469            );
470            log.insert(event_path!("map_field"), Value::Object(map_value));
471
472            let events = vec![Event::Log(log)];
473
474            // Build schema with all supported types
475            let struct_fields = arrow::datatypes::Fields::from(vec![
476                Field::new("f0", DataType::Utf8, true),
477                Field::new("f1", DataType::Int64, true),
478            ]);
479
480            let named_struct_fields = arrow::datatypes::Fields::from(vec![
481                Field::new("category", DataType::Utf8, true),
482                Field::new("tag", DataType::Utf8, true),
483            ]);
484
485            let map_entries = Field::new(
486                "entries",
487                DataType::Struct(arrow::datatypes::Fields::from(vec![
488                    Field::new("keys", DataType::Utf8, false),
489                    Field::new("values", DataType::Int64, true),
490                ])),
491                false,
492            );
493
494            let schema = Schema::new(vec![
495                Field::new("string_field", DataType::Utf8, true),
496                Field::new("int8_field", DataType::Int8, true),
497                Field::new("int16_field", DataType::Int16, true),
498                Field::new("int32_field", DataType::Int32, true),
499                Field::new("int64_field", DataType::Int64, true),
500                Field::new("uint8_field", DataType::UInt8, true),
501                Field::new("uint16_field", DataType::UInt16, true),
502                Field::new("uint32_field", DataType::UInt32, true),
503                Field::new("uint64_field", DataType::UInt64, true),
504                Field::new("float32_field", DataType::Float32, true),
505                Field::new("float64_field", DataType::Float64, true),
506                Field::new("bool_field", DataType::Boolean, true),
507                Field::new(
508                    "timestamp_field",
509                    DataType::Timestamp(TimeUnit::Millisecond, None),
510                    true,
511                ),
512                Field::new("decimal_field", DataType::Decimal128(10, 2), true),
513                Field::new(
514                    "list_field",
515                    DataType::List(Field::new("item", DataType::Int64, true).into()),
516                    true,
517                ),
518                Field::new("struct_field", DataType::Struct(struct_fields), true),
519                Field::new(
520                    "named_struct_field",
521                    DataType::Struct(named_struct_fields),
522                    true,
523                ),
524                Field::new("map_field", DataType::Map(map_entries.into(), false), true),
525            ])
526            .into();
527
528            let batch = encode_and_decode(events, schema).expect("Failed to encode");
529
530            assert_eq!(batch.num_rows(), 1);
531            assert_eq!(batch.num_columns(), 18);
532
533            // Verify all primitive types
534            assert_eq!(batch.column(0).as_string::<i32>().value(0), "test");
535            assert_eq!(batch.column(1).as_primitive::<Int8Type>().value(0), 127);
536            assert_eq!(batch.column(2).as_primitive::<Int16Type>().value(0), 32000);
537            assert_eq!(
538                batch.column(3).as_primitive::<Int32Type>().value(0),
539                1000000
540            );
541            assert_eq!(batch.column(4).as_primitive::<Int64Type>().value(0), 42);
542            assert_eq!(batch.column(5).as_primitive::<UInt8Type>().value(0), 255);
543            assert_eq!(batch.column(6).as_primitive::<UInt16Type>().value(0), 65535);
544            assert_eq!(
545                batch.column(7).as_primitive::<UInt32Type>().value(0),
546                4000000
547            );
548            assert_eq!(
549                batch.column(8).as_primitive::<UInt64Type>().value(0),
550                9000000000
551            );
552            assert!((batch.column(9).as_primitive::<Float32Type>().value(0) - 3.15).abs() < 0.001);
553            assert!((batch.column(10).as_primitive::<Float64Type>().value(0) - 3.15).abs() < 0.001);
554            assert!(batch.column(11).as_boolean().value(0));
555            assert_eq!(
556                batch
557                    .column(12)
558                    .as_primitive::<TimestampMillisecondType>()
559                    .value(0),
560                now.timestamp_millis()
561            );
562            assert_eq!(
563                batch.column(13).as_primitive::<Decimal128Type>().value(0),
564                9999
565            );
566
567            let list_array = batch.column(14).as_list::<i32>();
568            assert!(!list_array.is_null(0));
569            let list_values = list_array.value(0);
570            assert_eq!(list_values.len(), 3);
571            let int_array = list_values.as_primitive::<Int64Type>();
572            assert_eq!(int_array.value(0), 1);
573            assert_eq!(int_array.value(1), 2);
574            assert_eq!(int_array.value(2), 3);
575
576            // Verify struct field (unnamed)
577            let struct_array = batch.column(15).as_struct();
578            assert!(!struct_array.is_null(0));
579            assert_eq!(
580                struct_array.column(0).as_string::<i32>().value(0),
581                "nested_str"
582            );
583            assert_eq!(
584                struct_array.column(1).as_primitive::<Int64Type>().value(0),
585                999
586            );
587
588            // Verify named struct field (named tuple)
589            let named_struct_array = batch.column(16).as_struct();
590            assert!(!named_struct_array.is_null(0));
591            assert_eq!(
592                named_struct_array.column(0).as_string::<i32>().value(0),
593                "test_category"
594            );
595            assert_eq!(
596                named_struct_array.column(1).as_string::<i32>().value(0),
597                "test_tag"
598            );
599
600            // Verify map field
601            let map_array = batch.column(17).as_map();
602            assert!(!map_array.is_null(0));
603            let map_value = map_array.value(0);
604            assert_eq!(map_value.len(), 2);
605        }
606    }
607
608    mod error_handling {
609        use super::*;
610
611        #[test]
612        fn test_encode_empty_events() {
613            let schema = Schema::new(vec![Field::new("message", DataType::Utf8, true)]).into();
614            let events: Vec<Event> = vec![];
615            let result = encode_events_to_arrow_ipc_stream(&events, schema);
616            assert!(matches!(result.unwrap_err(), ArrowEncodingError::NoEvents));
617        }
618
619        #[test]
620        fn test_missing_non_nullable_field_errors() {
621            let events = vec![create_event(vec![("other_field", "value")])];
622
623            let schema = Schema::new(vec![Field::new(
624                "required_field",
625                DataType::Utf8,
626                false, // non-nullable
627            )])
628            .into();
629
630            let result = encode_events_to_arrow_ipc_stream(&events, schema);
631            assert!(result.is_err());
632        }
633    }
634
635    mod temporal_types {
636        use super::*;
637        use arrow::datatypes::{
638            TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
639            TimestampSecondType,
640        };
641
642        #[test]
643        fn test_encode_timestamp_precisions() {
644            let now = Utc::now();
645            let mut log = LogEvent::default();
646            log.insert(event_path!("ts_second"), now);
647            log.insert(event_path!("ts_milli"), now);
648            log.insert(event_path!("ts_micro"), now);
649            log.insert(event_path!("ts_nano"), now);
650
651            let events = vec![Event::Log(log)];
652
653            let schema = Schema::new(vec![
654                Field::new(
655                    "ts_second",
656                    DataType::Timestamp(TimeUnit::Second, None),
657                    true,
658                ),
659                Field::new(
660                    "ts_milli",
661                    DataType::Timestamp(TimeUnit::Millisecond, None),
662                    true,
663                ),
664                Field::new(
665                    "ts_micro",
666                    DataType::Timestamp(TimeUnit::Microsecond, None),
667                    true,
668                ),
669                Field::new(
670                    "ts_nano",
671                    DataType::Timestamp(TimeUnit::Nanosecond, None),
672                    true,
673                ),
674            ])
675            .into();
676
677            let batch = encode_and_decode(events, schema).unwrap();
678
679            assert_eq!(batch.num_rows(), 1);
680            assert_eq!(batch.num_columns(), 4);
681
682            let ts_second = batch.column(0).as_primitive::<TimestampSecondType>();
683            assert!(!ts_second.is_null(0));
684            assert_eq!(ts_second.value(0), now.timestamp());
685
686            let ts_milli = batch.column(1).as_primitive::<TimestampMillisecondType>();
687            assert!(!ts_milli.is_null(0));
688            assert_eq!(ts_milli.value(0), now.timestamp_millis());
689
690            let ts_micro = batch.column(2).as_primitive::<TimestampMicrosecondType>();
691            assert!(!ts_micro.is_null(0));
692            assert_eq!(ts_micro.value(0), now.timestamp_micros());
693
694            let ts_nano = batch.column(3).as_primitive::<TimestampNanosecondType>();
695            assert!(!ts_nano.is_null(0));
696            assert_eq!(ts_nano.value(0), now.timestamp_nanos_opt().unwrap());
697        }
698
699        #[test]
700        fn test_encode_mixed_timestamp_string_native_and_integer() {
701            let now = Utc::now();
702
703            let mut log1 = LogEvent::default();
704            log1.insert(event_path!("ts"), "2025-10-22T10:18:44.256Z"); // RFC3339 String
705
706            let mut log2 = LogEvent::default();
707            log2.insert(event_path!("ts"), now); // Native Timestamp
708
709            let mut log3 = LogEvent::default();
710            log3.insert(event_path!("ts"), 1729594724256000000_i64); // Integer (nanoseconds)
711
712            let events = vec![Event::Log(log1), Event::Log(log2), Event::Log(log3)];
713
714            let schema = Schema::new(vec![Field::new(
715                "ts",
716                DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
717                true,
718            )])
719            .into();
720
721            let batch = encode_and_decode(events, schema).unwrap();
722
723            assert_eq!(batch.num_rows(), 3);
724
725            let ts_array = batch.column(0).as_primitive::<TimestampNanosecondType>();
726
727            // All three should be non-null
728            assert!(!ts_array.is_null(0));
729            assert!(!ts_array.is_null(1));
730            assert!(!ts_array.is_null(2));
731
732            // First one should match the parsed RFC3339 string
733            let expected = chrono::DateTime::parse_from_rfc3339("2025-10-22T10:18:44.256Z")
734                .unwrap()
735                .timestamp_nanos_opt()
736                .unwrap();
737            assert_eq!(ts_array.value(0), expected);
738
739            // Second one should match the native timestamp
740            assert_eq!(ts_array.value(1), now.timestamp_nanos_opt().unwrap());
741
742            // Third one should match the integer
743            assert_eq!(ts_array.value(2), 1729594724256000000_i64);
744        }
745    }
746
747    mod config_tests {
748        use super::*;
749        use tokio_util::codec::Encoder;
750
751        #[test]
752        fn test_config_allow_nullable_fields_overrides_schema() {
753            let mut log1 = LogEvent::default();
754            log1.insert(event_path!("strict_field"), 42);
755            let log2 = LogEvent::default();
756            let events = vec![Event::Log(log1), Event::Log(log2)];
757
758            let schema = Schema::new(vec![Field::new("strict_field", DataType::Int64, false)]);
759
760            let mut config = ArrowStreamSerializerConfig::new(schema);
761            config.allow_nullable_fields = true;
762
763            let mut serializer =
764                ArrowStreamSerializer::new(config).expect("Failed to create serializer");
765
766            let mut buffer = BytesMut::new();
767            serializer
768                .encode(events, &mut buffer)
769                .expect("Encoding should succeed when allow_nullable_fields is true");
770
771            let cursor = Cursor::new(buffer);
772            let mut reader = StreamReader::try_new(cursor, None).expect("Failed to create reader");
773            let batch = reader.next().unwrap().expect("Failed to read batch");
774
775            assert_eq!(batch.num_rows(), 2);
776
777            let binding = batch.schema();
778            let output_field = binding.field(0);
779            assert!(
780                output_field.is_nullable(),
781                "The output schema field should have been transformed to nullable=true"
782            );
783
784            let array = batch
785                .column(0)
786                .as_primitive::<arrow::datatypes::Int64Type>();
787
788            assert_eq!(array.value(0), 42);
789            assert!(!array.is_null(0));
790            assert!(
791                array.is_null(1),
792                "The missing value should be encoded as null"
793            );
794        }
795
796        #[test]
797        fn test_make_field_nullable_with_nested_types() {
798            let inner_struct_field = Field::new("nested_field", DataType::Int64, false);
799            let inner_struct =
800                DataType::Struct(arrow::datatypes::Fields::from(vec![inner_struct_field]));
801            let list_field = Field::new("item", inner_struct, false);
802            let list_type = DataType::List(list_field.into());
803            let outer_field = Field::new("inner_list", list_type, false);
804            let outer_struct = DataType::Struct(arrow::datatypes::Fields::from(vec![outer_field]));
805
806            let original_field = Field::new("root", outer_struct, false);
807            let nullable_field = make_field_nullable(&original_field).unwrap();
808
809            assert!(
810                nullable_field.is_nullable(),
811                "Root field should be nullable"
812            );
813
814            if let DataType::Struct(root_fields) = nullable_field.data_type() {
815                let inner_list_field = &root_fields[0];
816                assert!(inner_list_field.is_nullable());
817
818                if let DataType::List(list_item_field) = inner_list_field.data_type() {
819                    assert!(list_item_field.is_nullable());
820
821                    if let DataType::Struct(inner_struct_fields) = list_item_field.data_type() {
822                        let nested_field = &inner_struct_fields[0];
823                        assert!(nested_field.is_nullable());
824                    } else {
825                        panic!("Expected Struct type for list items");
826                    }
827                } else {
828                    panic!("Expected List type for inner_list");
829                }
830            } else {
831                panic!("Expected Struct type for root field");
832            }
833        }
834
835        #[test]
836        fn test_make_field_nullable_with_map_type() {
837            let key_field = Field::new("key", DataType::Utf8, false);
838            let value_field = Field::new("value", DataType::Int64, false);
839            let entries_struct =
840                DataType::Struct(arrow::datatypes::Fields::from(vec![key_field, value_field]));
841            let entries_field = Field::new("entries", entries_struct, false);
842            let map_type = DataType::Map(entries_field.into(), false);
843
844            let original_field = Field::new("my_map", map_type, false);
845            let nullable_field = make_field_nullable(&original_field).unwrap();
846
847            assert!(
848                nullable_field.is_nullable(),
849                "Root map field should be nullable"
850            );
851
852            if let DataType::Map(entries_field, _sorted) = nullable_field.data_type() {
853                assert!(
854                    !entries_field.is_nullable(),
855                    "Map entries field should be non-nullable"
856                );
857
858                if let DataType::Struct(struct_fields) = entries_field.data_type() {
859                    let key_field = &struct_fields[0];
860                    let value_field = &struct_fields[1];
861                    assert!(
862                        !key_field.is_nullable(),
863                        "Map key field should be non-nullable"
864                    );
865                    assert!(
866                        value_field.is_nullable(),
867                        "Map value field should be nullable"
868                    );
869                } else {
870                    panic!("Expected Struct type for map entries");
871                }
872            } else {
873                panic!("Expected Map type for my_map field");
874            }
875        }
876    }
877
878    mod null_non_nullable {
879        use super::*;
880
881        #[test]
882        fn test_missing_non_nullable_field_error_names_fields() {
883            let schema: SchemaRef = Schema::new(vec![
884                Field::new("required_field", DataType::Utf8, false),
885                Field::new("optional_field", DataType::Utf8, true),
886            ])
887            .into();
888
889            // Event is missing "required_field" entirely
890            let event = create_event(vec![("optional_field", "hello")]);
891
892            let result = encode_events_to_arrow_ipc_stream(&[event], schema);
893            let err = result.unwrap_err().to_string();
894            assert!(
895                err.contains("required_field"),
896                "Error should name the missing field, got: {err}"
897            );
898            assert!(
899                !err.contains("optional_field"),
900                "Error should not name nullable fields, got: {err}"
901            );
902        }
903
904        #[test]
905        fn test_null_value_in_non_nullable_field_error_names_fields() {
906            let schema: SchemaRef = Schema::new(vec![
907                Field::new("id", DataType::Int64, false),
908                Field::new("name", DataType::Utf8, false),
909            ])
910            .into();
911
912            // Event has "id" but "name" is null
913            let event = create_event(vec![("id", Value::Integer(1))]);
914
915            let result = encode_events_to_arrow_ipc_stream(&[event], schema);
916            let err = result.unwrap_err().to_string();
917            assert!(
918                err.contains("name"),
919                "Error should name the null field, got: {err}"
920            );
921        }
922
923        #[test]
924        fn test_find_null_non_nullable_fields_returns_empty_when_all_present() {
925            let schema = Schema::new(vec![
926                Field::new("a", DataType::Utf8, false),
927                Field::new("b", DataType::Int64, false),
928            ]);
929
930            let event = create_event(vec![
931                ("a", Value::Bytes("val".into())),
932                ("b", Value::Integer(42)),
933            ]);
934            let missing = find_null_non_nullable_fields(
935                &schema,
936                &vector_log_events_to_json_values(&[event]).unwrap(),
937            );
938            assert!(
939                missing.is_empty(),
940                "Expected no missing fields, got: {missing:?}"
941            );
942        }
943
944        #[test]
945        fn test_find_null_non_nullable_fields_detects_explicit_null() {
946            let schema = Schema::new(vec![Field::new("a", DataType::Utf8, false)]);
947
948            let event = create_event(vec![("a", Value::Null)]);
949            let missing = find_null_non_nullable_fields(
950                &schema,
951                &vector_log_events_to_json_values(&[event]).unwrap(),
952            );
953            assert_eq!(missing, vec!["a"]);
954        }
955    }
956}
codecs/encoding/format/arrow.rs

codecs/encoding/format/
arrow.rs
