update

luoluoyuyu · luoluoyuyu · commit 55a50261cefa · 2026-04-02T21:15:51.000+08:00
diff --git a/src/runtime/streaming/format/deserializer.rs b/src/runtime/streaming/format/deserializer.rs
@@ -12,84 +12,176 @@
 
 
 use anyhow::{anyhow, Result};
-use arrow_array::builder::StringBuilder;
-use arrow_array::RecordBatch;
+use arrow_array::builder::{BinaryBuilder, StringBuilder, TimestampNanosecondBuilder};
+use arrow_array::{ArrayRef, RecordBatch};
 use arrow_json::reader::ReaderBuilder;
-use arrow_schema::SchemaRef;
+use arrow_schema::{Schema, SchemaRef};
+use std::collections::HashMap;
 use std::sync::Arc;
 
+use crate::sql::common::TIMESTAMP_FIELD;
+
 use super::config::{BadDataPolicy, Format};
 
 pub struct DataDeserializer {
     format: Format,
-    schema: SchemaRef,
+    final_schema: SchemaRef,
+    decoder_schema: SchemaRef,
     bad_data_policy: BadDataPolicy,
 }
 
 impl DataDeserializer {
     pub fn new(format: Format, schema: SchemaRef, bad_data_policy: BadDataPolicy) -> Self {
+        let decoder_schema = schema_without_timestamp(schema.as_ref());
         Self {
             format,
-            schema,
+            final_schema: schema,
+            decoder_schema,
             bad_data_policy,
         }
     }
 
     pub fn deserialize_batch(&self, messages: &[&[u8]]) -> Result<RecordBatch> {
+        self.deserialize_batch_with_kafka_timestamps(messages, &[])
+    }
+
+    pub fn deserialize_batch_with_kafka_timestamps(
+        &self,
+        messages: &[&[u8]],
+        kafka_timestamps_ms: &[u64],
+    ) -> Result<RecordBatch> {
         match &self.format {
-            Format::Json(_) => self.deserialize_json(messages),
-            Format::RawString => self.deserialize_raw_string(messages),
-            Format::RawBytes => self.deserialize_raw_bytes(messages),
+            Format::Json(_) => self.deserialize_json(messages, kafka_timestamps_ms),
+            Format::RawString => self.deserialize_raw_string(messages, kafka_timestamps_ms),
+            Format::RawBytes => self.deserialize_raw_bytes(messages, kafka_timestamps_ms),
         }
     }
 
-    fn deserialize_json(&self, messages: &[&[u8]]) -> Result<RecordBatch> {
+    fn deserialize_json(&self, messages: &[&[u8]], kafka_timestamps_ms: &[u64]) -> Result<RecordBatch> {
         let mut buffer = Vec::with_capacity(messages.len() * 256);
         for msg in messages {
             buffer.extend_from_slice(msg);
             buffer.push(b'\n');
         }
 
         let allow_bad_data = self.bad_data_policy == BadDataPolicy::Drop;
-        let mut decoder = ReaderBuilder::new(self.schema.clone())
-            .with_strict_mode(!allow_bad_data)
+        let mut decoder = ReaderBuilder::new(self.decoder_schema.clone())
+            .with_strict_mode(false)
             .build_decoder()?;
 
         decoder.decode(&buffer)?;
 
-        let batch = if allow_bad_data {
-            let (batch, _mask, _, _errors) = decoder.flush_with_bad_data()?.unwrap();
-            batch
+        let (batch, valid_indices) = if allow_bad_data {
+            let Some((batch, mask, _, _errors)) = decoder.flush_with_bad_data()? else {
+                return Ok(RecordBatch::new_empty(self.final_schema.clone()));
+            };
+            let mut indices = Vec::with_capacity(batch.num_rows());
+            for i in 0..mask.len() {
+                if mask.value(i) {
+                    indices.push(i);
+                }
+            }
+            (batch, indices)
         } else {
-            decoder
+            let batch = decoder
                 .flush()?
-                .ok_or_else(|| anyhow!("JSON decoder returned no batch"))?
+                .unwrap_or_else(|| RecordBatch::new_empty(self.decoder_schema.clone()));
+            let indices: Vec<usize> = (0..batch.num_rows()).collect();
+            (batch, indices)
         };
 
-        Ok(batch)
+        self.rebuild_with_timestamp(batch, kafka_timestamps_ms, &valid_indices)
     }
 
-    fn deserialize_raw_string(&self, messages: &[&[u8]]) -> Result<RecordBatch> {
+    fn deserialize_raw_string(&self, messages: &[&[u8]], kafka_timestamps_ms: &[u64]) -> Result<RecordBatch> {
+        let value_idx = self
+            .decoder_schema
+            .index_of("value")
+            .map_err(|_| anyhow!("RawString format requires a 'value' column"))?;
+
         let mut builder = StringBuilder::with_capacity(messages.len(), messages.len() * 64);
         for msg in messages {
             builder.append_value(String::from_utf8_lossy(msg));
         }
 
-        let array = Arc::new(builder.finish());
-        RecordBatch::try_new(self.schema.clone(), vec![array])
-            .map_err(|e| anyhow!("build RawString batch: {e}"))
+        let mut columns = vec![None; self.decoder_schema.fields().len()];
+        columns[value_idx] = Some(Arc::new(builder.finish()) as ArrayRef);
+        let decoded_columns = columns
+            .into_iter()
+            .map(|c| c.ok_or_else(|| anyhow!("missing RawString decoded column")))
+            .collect::<Result<Vec<_>>>()?;
+        let decoded_batch = RecordBatch::try_new(self.decoder_schema.clone(), decoded_columns)
+            .map_err(|e| anyhow!("build RawString decoded batch: {e}"))?;
+        let valid_indices: Vec<usize> = (0..decoded_batch.num_rows()).collect();
+        self.rebuild_with_timestamp(decoded_batch, kafka_timestamps_ms, &valid_indices)
     }
 
-    fn deserialize_raw_bytes(&self, messages: &[&[u8]]) -> Result<RecordBatch> {
-        use arrow_array::builder::BinaryBuilder;
-
+    fn deserialize_raw_bytes(&self, messages: &[&[u8]], kafka_timestamps_ms: &[u64]) -> Result<RecordBatch> {
+        let value_idx = self
+            .decoder_schema
+            .index_of("value")
+            .map_err(|_| anyhow!("RawBytes format requires a 'value' column"))?;
         let mut builder = BinaryBuilder::with_capacity(messages.len(), messages.len() * 64);
         for msg in messages {
             builder.append_value(msg);
         }
 
-        let array = Arc::new(builder.finish());
-        RecordBatch::try_new(self.schema.clone(), vec![array])
-            .map_err(|e| anyhow!("build RawBytes batch: {e}"))
+        let mut columns = vec![None; self.decoder_schema.fields().len()];
+        columns[value_idx] = Some(Arc::new(builder.finish()) as ArrayRef);
+        let decoded_columns = columns
+            .into_iter()
+            .map(|c| c.ok_or_else(|| anyhow!("missing RawBytes decoded column")))
+            .collect::<Result<Vec<_>>>()?;
+        let decoded_batch = RecordBatch::try_new(self.decoder_schema.clone(), decoded_columns)
+            .map_err(|e| anyhow!("build RawBytes decoded batch: {e}"))?;
+        let valid_indices: Vec<usize> = (0..decoded_batch.num_rows()).collect();
+        self.rebuild_with_timestamp(decoded_batch, kafka_timestamps_ms, &valid_indices)
+    }
+
+    fn rebuild_with_timestamp(
+        &self,
+        decoded_batch: RecordBatch,
+        kafka_timestamps_ms: &[u64],
+        valid_indices: &[usize],
+    ) -> Result<RecordBatch> {
+        let mut by_name: HashMap<String, ArrayRef> = decoded_batch
+            .schema()
+            .fields()
+            .iter()
+            .zip(decoded_batch.columns().iter())
+            .map(|(f, a)| (f.name().to_string(), a.clone()))
+            .collect();
+
+        let mut ts_builder = TimestampNanosecondBuilder::with_capacity(valid_indices.len());
+        for idx in valid_indices {
+            let ms = kafka_timestamps_ms.get(*idx).copied().unwrap_or(0);
+            ts_builder.append_value((ms as i64).saturating_mul(1_000_000));
+        }
+        let timestamp_col: ArrayRef = Arc::new(ts_builder.finish());
+
+        let mut columns = Vec::with_capacity(self.final_schema.fields().len());
+        for field in self.final_schema.fields() {
+            if field.name() == TIMESTAMP_FIELD {
+                columns.push(timestamp_col.clone());
+            } else {
+                let array = by_name
+                    .remove(field.name())
+                    .ok_or_else(|| anyhow!("decoded JSON missing field '{}'", field.name()))?;
+                columns.push(array);
+            }
+        }
+
+        RecordBatch::try_new(self.final_schema.clone(), columns)
+            .map_err(|e| anyhow!("build JSON batch with _timestamp: {e}"))
     }
 }
+
+fn schema_without_timestamp(schema: &Schema) -> SchemaRef {
+    let fields = schema
+        .fields()
+        .iter()
+        .filter(|f| f.name() != TIMESTAMP_FIELD)
+        .cloned()
+        .collect::<Vec<_>>();
+    Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))
+}
diff --git a/src/runtime/streaming/format/serializer.rs b/src/runtime/streaming/format/serializer.rs
@@ -18,6 +18,8 @@ use arrow_json::EncoderOptions;
 use arrow_schema::{DataType, Field, SchemaRef};
 use std::sync::Arc;
 
+use crate::sql::common::TIMESTAMP_FIELD;
+
 use super::config::{Format, JsonFormat};
 use super::json_encoder::CustomEncoderFactory;
 
@@ -32,7 +34,7 @@ impl DataSerializer {
             .fields()
             .iter()
             .enumerate()
-            .filter(|(_, f)| !f.name().starts_with('_'))
+            .filter(|(_, f)| f.name() != TIMESTAMP_FIELD)
             .map(|(i, _)| i)
             .collect();
 
diff --git a/src/runtime/streaming/operators/source/kafka/mod.rs b/src/runtime/streaming/operators/source/kafka/mod.rs
@@ -59,6 +59,8 @@ pub trait BatchDeserializer: Send + 'static {
 pub struct BufferedDeserializer {
     inner: DataDeserializer,
     buffer: Vec<Vec<u8>>,
+    /// Parallel to `buffer`: Kafka message timestamp (ms) per row for filling `_timestamp`.
+    kafka_timestamps_ms: Vec<u64>,
     batch_size: usize,
 }
 
@@ -67,6 +69,7 @@ impl BufferedDeserializer {
         Self {
             inner: DataDeserializer::new(format, schema, bad_data_policy),
             buffer: Vec::with_capacity(batch_size),
+            kafka_timestamps_ms: Vec::with_capacity(batch_size),
             batch_size,
         }
     }
@@ -76,10 +79,11 @@ impl BatchDeserializer for BufferedDeserializer {
     fn deserialize_slice(
         &mut self,
         payload: &[u8],
-        _timestamp: u64,
+        timestamp: u64,
         _metadata: Option<HashMap<&str, FieldValueType<'_>>>,
     ) -> Result<()> {
         self.buffer.push(payload.to_vec());
+        self.kafka_timestamps_ms.push(timestamp);
         Ok(())
     }
 
@@ -93,8 +97,11 @@ impl BatchDeserializer for BufferedDeserializer {
         }
 
         let refs: Vec<&[u8]> = self.buffer.iter().map(|v| v.as_slice()).collect();
-        let batch = self.inner.deserialize_batch(&refs)?;
+        let batch = self
+            .inner
+            .deserialize_batch_with_kafka_timestamps(&refs, &self.kafka_timestamps_ms)?;
         self.buffer.clear();
+        self.kafka_timestamps_ms.clear();
         Ok(Some(batch))
     }
 
@@ -277,7 +284,11 @@ impl SourceOperator for KafkaSourceOperator {
             Ok(Ok(msg)) => {
                 let partition = msg.partition();
                 let offset = msg.offset();
-                let timestamp = msg.timestamp().to_millis().unwrap_or(0);
+                let timestamp = msg.timestamp().to_millis().ok_or_else(|| {
+                    anyhow!(
+                        "Failed to read timestamp from Kafka record: message has no timestamp"
+                    )
+                })?;
 
                 self.current_offsets.insert(partition, offset);
 
