diff --git a/changelog.d/25772_clickhouse_arrow_drop_json_roundtrip.enhancement.md b/changelog.d/25772_clickhouse_arrow_drop_json_roundtrip.enhancement.md new file mode 100644 index 0000000000000..3a1407ca27857 --- /dev/null +++ b/changelog.d/25772_clickhouse_arrow_drop_json_roundtrip.enhancement.md @@ -0,0 +1,3 @@ +The `clickhouse` sink's `arrow_stream` batch encoding now serializes events directly into the Arrow record batch instead of first materializing each event as an intermediate JSON value. This removes one allocation and one serialization pass per event on the encode hot path, reducing CPU by roughly 30% on wide or deeply nested events. Encoding behavior is otherwise unchanged, with one exception: non-finite floats (`inf`/`-inf`) are now passed through to the destination, whereas previously they were encoded as null. + +authors: benjamin-awd diff --git a/lib/codecs/src/encoding/format/arrow.rs b/lib/codecs/src/encoding/format/arrow.rs index 0c67317b7f926..966418810cd70 100644 --- a/lib/codecs/src/encoding/format/arrow.rs +++ b/lib/codecs/src/encoding/format/arrow.rs @@ -6,7 +6,6 @@ use arrow::{ datatypes::{DataType, Field, Fields, Schema, SchemaRef}, - error::ArrowError, ipc::writer::StreamWriter, json::reader::ReaderBuilder, record_batch::RecordBatch, @@ -15,7 +14,7 @@ use async_trait::async_trait; use bytes::{BufMut, Bytes, BytesMut}; use snafu::{ResultExt, Snafu, ensure}; use vector_config::configurable_component; -use vector_core::event::Event; +use vector_core::event::{Event, LogEvent, Value}; /// Provides Arrow schema for encoding. /// @@ -98,12 +97,8 @@ impl ArrowStreamSerializer { &self, events: &[Event], ) -> Result { - let values = vector_log_events_to_json_values(events).map_err(|e| { - ArrowEncodingError::RecordBatchCreation { - source: arrow::error::ArrowError::JsonError(e.to_string()), - } - })?; - build_record_batch(self.schema.clone(), &values) + let events: Vec<&LogEvent> = events.iter().filter_map(Event::maybe_as_log).collect(); + build_record_batch(self.schema.clone(), &events) } /// Create a new ArrowStreamSerializer with the given configuration @@ -217,13 +212,9 @@ pub fn encode_events_to_arrow_ipc_stream( return Err(ArrowEncodingError::NoEvents); } - let json_values = vector_log_events_to_json_values(events).map_err(|e| { - ArrowEncodingError::RecordBatchCreation { - source: ArrowError::JsonError(e.to_string()), - } - })?; + let events: Vec<&LogEvent> = events.iter().filter_map(Event::maybe_as_log).collect(); - let record_batch = build_record_batch(schema, &json_values)?; + let record_batch = build_record_batch(schema, &events)?; let mut buffer = BytesMut::new().writer(); let mut writer = @@ -289,26 +280,25 @@ fn make_field_nullable(field: &Field) -> Result { /// Returns true if the field is absent from the value's object map, or explicitly null. /// Find non-nullable schema fields that are missing or null in any of the given events. -pub fn find_null_non_nullable_fields<'a>( - schema: &'a Schema, - values: &[serde_json::Value], -) -> Vec<&'a str> { +pub fn find_null_non_nullable_fields<'a>(schema: &'a Schema, events: &[&LogEvent]) -> Vec<&'a str> { schema .fields() .iter() .filter(|field| { !field.is_nullable() - && values.iter().any(|value| { - value + && events.iter().any(|event| { + event + .value() .as_object() .and_then(|map| map.get(field.name().as_str())) - .is_none_or(serde_json::Value::is_null) + .is_none_or(Value::is_null) }) }) .map(|field| field.name().as_str()) .collect() } +#[cfg(feature = "parquet")] pub(crate) fn vector_log_events_to_json_values( events: &[Event], ) -> Result, serde_json::Error> { @@ -322,13 +312,13 @@ pub(crate) fn vector_log_events_to_json_values( /// Build an Arrow RecordBatch from a slice of events using the provided schema. pub(crate) fn build_record_batch( schema: SchemaRef, - values: &[serde_json::Value], + events: &[&LogEvent], ) -> Result { - if values.is_empty() { + if events.is_empty() { return Err(ArrowEncodingError::NoEvents); } - let missing = find_null_non_nullable_fields(&schema, values); + let missing = find_null_non_nullable_fields(&schema, events); if !missing.is_empty() { let error: vector_common::Error = Box::new(ArrowEncodingError::NullConstraint { field_name: missing.join(", "), @@ -341,6 +331,15 @@ pub(crate) fn build_record_batch( }); } + decode_rows_to_record_batch(schema, events) +} + +/// Serialize rows through Arrow's JSON decoder into one RecordBatch. Parquet `AutoInfer` +/// encodes from the rows it inferred the schema from, so the two stay in sync (e.g. `inf`->`null`). +pub(crate) fn decode_rows_to_record_batch( + schema: SchemaRef, + rows: &[T], +) -> Result { let mut decoder = ReaderBuilder::new(schema) .build_decoder() .inspect_err(|e| { @@ -352,7 +351,7 @@ pub(crate) fn build_record_batch( .context(RecordBatchCreationSnafu)?; decoder - .serialize(values) + .serialize(rows) .inspect_err(|e| { vector_common::internal_event::emit(crate::internal_events::EncoderRecordBatchError { error: e, @@ -927,10 +926,7 @@ mod tests { ("a", Value::Bytes("val".into())), ("b", Value::Integer(42)), ]); - let missing = find_null_non_nullable_fields( - &schema, - &vector_log_events_to_json_values(&[event]).unwrap(), - ); + let missing = find_null_non_nullable_fields(&schema, &[event.as_log()]); assert!( missing.is_empty(), "Expected no missing fields, got: {missing:?}" @@ -942,10 +938,7 @@ mod tests { let schema = Schema::new(vec![Field::new("a", DataType::Utf8, false)]); let event = create_event(vec![("a", Value::Null)]); - let missing = find_null_non_nullable_fields( - &schema, - &vector_log_events_to_json_values(&[event]).unwrap(), - ); + let missing = find_null_non_nullable_fields(&schema, &[event.as_log()]); assert_eq!(missing, vec!["a"]); } } diff --git a/lib/codecs/src/encoding/format/parquet.rs b/lib/codecs/src/encoding/format/parquet.rs index d954340b890e0..c0250229e1adf 100644 --- a/lib/codecs/src/encoding/format/parquet.rs +++ b/lib/codecs/src/encoding/format/parquet.rs @@ -27,9 +27,9 @@ use vector_common::internal_event::{ ComponentEventsDropped, Count, InternalEventHandle, Registered, UNINTENTIONAL, emit, register, }; use vector_config::configurable_component; -use vector_core::event::Event; +use vector_core::event::{Event, LogEvent}; -use super::arrow::{ArrowEncodingError, build_record_batch}; +use super::arrow::{ArrowEncodingError, build_record_batch, decode_rows_to_record_batch}; use crate::encoding::format::arrow::vector_log_events_to_json_values; use crate::internal_events::{ArrowWriterError, JsonSerializationError, SchemaGenerationError}; @@ -319,15 +319,9 @@ impl tokio_util::codec::Encoder> for ParquetSerializer { return Ok(()); } - let json_values = match vector_log_events_to_json_values(&events) { - Ok(values) => values, - Err(e) => { - emit(JsonSerializationError { error: &e }); - return Err(Box::new(e)); - } - }; + let logs: Vec<&LogEvent> = events.iter().filter_map(Event::maybe_as_log).collect(); - let non_log_count = events.len() - json_values.len(); + let non_log_count = events.len() - logs.len(); if non_log_count > 0 { warn!( @@ -338,17 +332,15 @@ impl tokio_util::codec::Encoder> for ParquetSerializer { self.events_dropped_handle.emit(Count(non_log_count)) } - if json_values.is_empty() { + if logs.is_empty() { return Ok(()); } - match self.schema_mode { + let record_batch = match self.schema_mode { // In strict mode, check for extra top-level fields not in the schema. ParquetSchemaMode::Strict => { - for event in &events { - if let Some(log) = event.maybe_as_log() - && let Some(object_map) = log.as_map() - { + for log in &logs { + if let Some(object_map) = log.as_map() { for top_level in object_map.keys() { if !self.schema_field_names.contains(top_level.as_str()) { return Err(Box::new(ArrowEncodingError::SchemaFetchError { @@ -360,18 +352,29 @@ impl tokio_util::codec::Encoder> for ParquetSerializer { } } } + build_record_batch(Arc::clone(&self.schema), &logs).map_err(Box::new)? } ParquetSchemaMode::AutoInfer => { + // Schema inference is driven by Arrow's JSON schema inference, which still + // requires `serde_json::Value`s. + let json_values = match vector_log_events_to_json_values(&events) { + Ok(values) => values, + Err(e) => { + emit(JsonSerializationError { error: &e }); + return Err(Box::new(e)); + } + }; let schema = ParquetSchemaGenerator::infer_schema(&json_values)?; self.schema = Arc::new(ParquetSchemaGenerator::try_normalize_schema( &events, schema, )); + decode_rows_to_record_batch(Arc::clone(&self.schema), &json_values) + .map_err(Box::new)? } - ParquetSchemaMode::Relaxed => {} - } - - let record_batch = - build_record_batch(Arc::clone(&self.schema), &json_values).map_err(Box::new)?; + ParquetSchemaMode::Relaxed => { + build_record_batch(Arc::clone(&self.schema), &logs).map_err(Box::new)? + } + }; Self::write_record_batch(&record_batch, buffer, &self.writer_props).map_err(Box::new)?; @@ -627,6 +630,16 @@ mod tests { } } + #[test] + fn autoinfer_encodes_non_finite_floats() { + let events = vec![ + create_event(vec![("ratio", f64::INFINITY)]), + create_event(vec![("ratio", f64::NEG_INFINITY)]), + ]; + let (_schema, num_rows) = encode_autoinfer_and_read_schema(events); + assert_eq!(num_rows, 2); + } + #[test] fn test_parquet_empty_events() { let mut serializer = ParquetSerializer::new(ParquetSerializerConfig {