fix(delta): strip system columns and normalize timestamps for writes

luoluoyuyu · cursoragent · luoluoyuyu · commit efa47cf27ba8 · 2026-05-20T10:12:11.000+08:00
Exclude _timestamp/_updating_meta from Delta catalog and Parquet schema,
cast user columns to Delta-compatible types, and support multiple
timestamp units in watermark max extraction.

Co-authored-by: Cursor &lt;cursoragent@cursor.com&gt;
diff --git a/src/streaming_runtime_operators/src/factory/connector/delta.rs b/src/streaming_runtime_operators/src/factory/connector/delta.rs
@@ -24,7 +24,9 @@ use crate::factory::connector::sink_props_codec::{
 };
 use crate::factory::global::Registry;
 use crate::factory::operator_constructor::OperatorConstructor;
-use crate::operators::sink::delta::{DeltaFormat, DeltaSinkOperator};
+use crate::operators::sink::delta::{
+    DeltaFormat, DeltaSinkOperator, strip_streaming_system_columns_arc,
+};
 use crate::operators::sink::filesystem::compression_from_str;
 use crate::sql::common::FsSchema;
 use crate::sql::common::constants::connection_format_value;
@@ -47,7 +49,7 @@ impl OperatorConstructor for DeltaSinkDispatcher {
             .map(|fs| FsSchema::try_from(fs.clone()))
             .transpose()
             .map_err(|e| anyhow::anyhow!("invalid fs_schema for delta sink: {e}"))?
-            .map(|fs| fs.schema);
+            .and_then(|fs| strip_streaming_system_columns_arc(fs.schema));
 
         let format = props
             .get(opt::FORMAT)
diff --git a/src/streaming_runtime_operators/src/operators/sink/delta/delta_commit.rs b/src/streaming_runtime_operators/src/operators/sink/delta/delta_commit.rs
@@ -14,9 +14,10 @@ use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
 
+use arrow::compute::cast;
 use arrow_array::RecordBatch;
 use arrow_ipc::writer::StreamWriter;
-use arrow_schema::Schema as ArrowSchema;
+use arrow_schema::{DataType, Field, FieldRef, Schema as ArrowSchema, TimeUnit};
 use deltalake::errors::DeltaTableError;
 use deltalake::kernel::engine::arrow_conversion::TryIntoKernel as _;
 use deltalake::kernel::schema::cast::normalize_for_delta;
@@ -27,8 +28,42 @@ use deltalake::{DeltaTable, open_table_with_storage_options};
 use tracing::{info, instrument};
 use url::Url;
 
+use crate::sql::common::{TIMESTAMP_FIELD, UPDATING_META_FIELD};
+
 use super::DeltaSinkError;
 
+/// Streaming-internal columns that must not be persisted to external sinks.
+pub fn is_streaming_system_column(name: &str) -> bool {
+    name == TIMESTAMP_FIELD || name == UPDATING_META_FIELD
+}
+
+/// Remove `_timestamp` / `_updating_meta` from a schema (e.g. connector `fs_schema` may inject them).
+pub fn strip_streaming_system_columns(schema: &ArrowSchema) -> ArrowSchema {
+    let fields: Vec<FieldRef> = schema
+        .fields()
+        .iter()
+        .filter(|f| !is_streaming_system_column(f.name()))
+        .cloned()
+        .collect();
+    ArrowSchema::new(fields)
+}
+
+pub fn strip_streaming_system_columns_arc(schema: Arc<ArrowSchema>) -> Option<Arc<ArrowSchema>> {
+    let stripped = strip_streaming_system_columns(schema.as_ref());
+    if stripped.fields().is_empty() {
+        return None;
+    }
+    let had_system = schema
+        .fields()
+        .iter()
+        .any(|f| is_streaming_system_column(f.name()));
+    if had_system {
+        Some(Arc::new(stripped))
+    } else {
+        Some(schema)
+    }
+}
+
 pub struct UncommittedDataFile {
     pub path: String,
     pub size_bytes: u64,
@@ -42,6 +77,8 @@ pub struct DeltaTableCommitter {
     table: Option<DeltaTable>,
     /// Precomputed from catalog `fs_schema` at startup, or lazily from the first batch.
     delta_columns: Option<Vec<StructField>>,
+    /// Arrow 55 schema for Parquet writes (timestamps normalized to microsecond).
+    write_schema: Option<Arc<ArrowSchema>>,
 }
 
 impl DeltaTableCommitter {
@@ -50,7 +87,11 @@ impl DeltaTableCommitter {
         storage_options: HashMap<String, String>,
         catalog_schema: Option<Arc<ArrowSchema>>,
     ) -> Result<Self, DeltaSinkError> {
-        let delta_columns = catalog_schema
+        let user_schema = catalog_schema.and_then(strip_streaming_system_columns_arc);
+        let write_schema = user_schema
+            .as_ref()
+            .map(|s| Arc::new(normalize_arrow_schema_for_delta(s)));
+        let delta_columns = user_schema
             .as_deref()
             .map(arrow_schema_to_delta_columns)
             .transpose()?;
@@ -61,13 +102,30 @@ impl DeltaTableCommitter {
             uncommitted: Vec::new(),
             table: None,
             delta_columns,
+            write_schema,
         })
     }
 
-    /// Fallback when catalog schema is absent: derive columns from the first flushed batch.
+    pub fn write_schema(&self) -> Option<Arc<ArrowSchema>> {
+        self.write_schema.clone()
+    }
+
+    /// Fallback when catalog schema is absent: derive user columns from the first flushed batch.
     pub fn update_schema(&mut self, schema: Arc<ArrowSchema>) -> Result<(), DeltaSinkError> {
+        let user_schema = strip_streaming_system_columns_arc(schema).ok_or_else(|| {
+            DeltaSinkError::CommitterFailed(
+                "cannot derive delta table schema: no user columns after removing streaming \
+                 system columns (_timestamp, _updating_meta)"
+                    .into(),
+            )
+        })?;
         if self.delta_columns.is_none() {
-            self.delta_columns = Some(arrow_schema_to_delta_columns(&schema)?);
+            self.delta_columns = Some(arrow_schema_to_delta_columns(user_schema.as_ref())?);
+        }
+        if self.write_schema.is_none() {
+            self.write_schema = Some(Arc::new(normalize_arrow_schema_for_delta(
+                user_schema.as_ref(),
+            )));
         }
         Ok(())
     }
@@ -231,6 +289,105 @@ impl DeltaTableCommitter {
     }
 }
 
+/// Normalize Arrow 55 schema for Delta Parquet writes (align with deltalake `normalize_for_delta`).
+pub fn normalize_arrow_schema_for_delta(schema: &ArrowSchema) -> ArrowSchema {
+    let fields: Vec<FieldRef> = schema
+        .fields()
+        .iter()
+        .map(|f| Arc::new(normalize_field_for_delta(f.as_ref())))
+        .collect();
+    ArrowSchema::new(fields)
+}
+
+fn normalize_field_for_delta(field: &Field) -> Field {
+    let data_type = normalize_datatype_for_delta(field.data_type());
+    if data_type == *field.data_type() {
+        field.clone()
+    } else {
+        field.clone().with_data_type(data_type)
+    }
+}
+
+fn normalize_datatype_for_delta(dt: &DataType) -> DataType {
+    match dt {
+        DataType::Date64 => DataType::Date32,
+        DataType::Timestamp(TimeUnit::Second, tz)
+        | DataType::Timestamp(TimeUnit::Millisecond, tz)
+        | DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
+            DataType::Timestamp(TimeUnit::Microsecond, tz.clone())
+        }
+        DataType::Struct(fields) => {
+            let normalized: Vec<FieldRef> = fields
+                .iter()
+                .map(|f| Arc::new(normalize_field_for_delta(f.as_ref())))
+                .collect();
+            DataType::Struct(normalized.into())
+        }
+        DataType::List(inner) => {
+            DataType::List(Arc::new(normalize_field_for_delta(inner.as_ref())))
+        }
+        DataType::LargeList(inner) => {
+            DataType::LargeList(Arc::new(normalize_field_for_delta(inner.as_ref())))
+        }
+        DataType::FixedSizeList(inner, len) => {
+            DataType::FixedSizeList(Arc::new(normalize_field_for_delta(inner.as_ref())), *len)
+        }
+        DataType::Map(entries, sorted) => DataType::Map(
+            Arc::new(normalize_field_for_delta(entries.as_ref())),
+            *sorted,
+        ),
+        _ => dt.clone(),
+    }
+}
+
+/// Cast record batches so on-disk Parquet matches the Delta table schema.
+pub fn cast_batches_for_delta_write(
+    batches: &[RecordBatch],
+    target_schema: &ArrowSchema,
+) -> Result<Vec<RecordBatch>, DeltaSinkError> {
+    let target = Arc::new(target_schema.clone());
+    batches
+        .iter()
+        .map(|batch| cast_batch_for_delta_write(batch, &target))
+        .collect()
+}
+
+fn cast_batch_for_delta_write(
+    batch: &RecordBatch,
+    target_schema: &Arc<ArrowSchema>,
+) -> Result<RecordBatch, DeltaSinkError> {
+    if batch.schema().as_ref() == target_schema.as_ref() {
+        return Ok(batch.clone());
+    }
+
+    let mut columns = Vec::with_capacity(target_schema.fields().len());
+    for field in target_schema.fields() {
+        let col = batch.column_by_name(field.name()).ok_or_else(|| {
+            DeltaSinkError::CommitterFailed(format!(
+                "batch missing column '{}' required by delta write schema",
+                field.name()
+            ))
+        })?;
+        let casted = if col.data_type() == field.data_type() {
+            col.clone()
+        } else {
+            cast(col, field.data_type()).map_err(|e| {
+                DeltaSinkError::CommitterFailed(format!(
+                    "failed to cast column '{}' from {:?} to {:?}: {e}",
+                    field.name(),
+                    col.data_type(),
+                    field.data_type()
+                ))
+            })?
+        };
+        columns.push(casted);
+    }
+
+    RecordBatch::try_new(target_schema.clone(), columns).map_err(|e| {
+        DeltaSinkError::CommitterFailed(format!("failed to build delta write batch: {e}"))
+    })
+}
+
 /// Bridge arrow 55 (runtime) schema to deltalake kernel schema via IPC.
 fn arrow_schema_to_delta_columns(schema: &ArrowSchema) -> Result<Vec<StructField>, DeltaSinkError> {
     let empty = RecordBatch::new_empty(Arc::new(schema.clone()));
diff --git a/src/streaming_runtime_operators/src/operators/sink/delta/mod.rs b/src/streaming_runtime_operators/src/operators/sink/delta/mod.rs
@@ -12,6 +12,8 @@
 
 mod delta_commit;
 
+pub use delta_commit::strip_streaming_system_columns_arc;
+
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -21,7 +23,8 @@ use arrow_schema::Schema as ArrowSchema;
 use async_trait::async_trait;
 use bytes::Bytes;
 use delta_commit::{
-    DeltaTableCommitter, UncommittedDataFile, build_delta_storage_options, resolve_delta_table_uri,
+    DeltaTableCommitter, UncommittedDataFile, build_delta_storage_options,
+    cast_batches_for_delta_write, resolve_delta_table_uri,
 };
 use object_store::aws::AmazonS3Builder;
 use object_store::path::Path as ObjectStorePath;
@@ -94,6 +97,8 @@ pub struct DeltaSinkOperator {
     s3_bucket: Option<String>,
     sink_path: String,
     catalog_schema: Option<Arc<ArrowSchema>>,
+    /// Normalized schema for Parquet files (microsecond timestamps, etc.).
+    parquet_write_schema: Option<Arc<ArrowSchema>>,
 }
 
 impl DeltaSinkOperator {
@@ -171,7 +176,8 @@ impl DeltaSinkOperator {
             storage_options,
             s3_bucket,
             sink_path: path,
-            catalog_schema,
+            catalog_schema: catalog_schema.and_then(strip_streaming_system_columns_arc),
+            parquet_write_schema: None,
         })
     }
 
@@ -186,23 +192,35 @@ impl DeltaSinkOperator {
             return Ok(());
         }
 
-        let fallback_schema = self
-            .catalog_schema
-            .is_none()
-            .then(|| self.pending[0].schema());
+        let fallback_schema = self.catalog_schema.is_none().then(|| {
+            strip_streaming_system_columns_arc(self.pending[0].schema())
+                .expect("batch has no user columns after removing streaming system columns")
+        });
         let batches = std::mem::take(&mut self.pending);
         self.pending_bytes = 0;
 
         let record_count: u64 = batches.iter().map(|b| b.num_rows() as u64).sum();
         let format = self.format;
         let compression = self.parquet_compression;
+        let parquet_write_schema = self.parquet_write_schema.clone();
 
-        let bytes = tokio::task::spawn_blocking(move || match format {
-            DeltaFormat::Csv => FormatEncoder::encode_csv(&batches),
-            DeltaFormat::Parquet => FormatEncoder::encode_parquet(&batches, compression),
-            DeltaFormat::JsonL => FormatEncoder::encode_jsonl(&batches),
-            DeltaFormat::Avro => FormatEncoder::encode_avro(&batches),
-            DeltaFormat::Orc => FormatEncoder::encode_orc(&batches),
+        let bytes = tokio::task::spawn_blocking(move || {
+            let batches = if format == DeltaFormat::Parquet {
+                if let Some(ref schema) = parquet_write_schema {
+                    cast_batches_for_delta_write(&batches, schema)?
+                } else {
+                    batches
+                }
+            } else {
+                batches
+            };
+            match format {
+                DeltaFormat::Csv => FormatEncoder::encode_csv(&batches),
+                DeltaFormat::Parquet => FormatEncoder::encode_parquet(&batches, compression),
+                DeltaFormat::JsonL => FormatEncoder::encode_jsonl(&batches),
+                DeltaFormat::Avro => FormatEncoder::encode_avro(&batches),
+                DeltaFormat::Orc => FormatEncoder::encode_orc(&batches),
+            }
         })
         .await
         .map_err(|e| DeltaSinkError::SerializationPanic(e.to_string()))?
@@ -302,11 +320,13 @@ impl Operator for DeltaSinkOperator {
         self.table_uri = Some(table_uri.clone());
 
         if self.format == DeltaFormat::Parquet {
-            self.committer = Some(DeltaTableCommitter::try_new(
+            let committer = DeltaTableCommitter::try_new(
                 table_uri,
                 self.storage_options.clone(),
                 self.catalog_schema.clone(),
-            )?);
+            )?;
+            self.parquet_write_schema = committer.write_schema();
+            self.committer = Some(committer);
         } else {
             warn!(
                 format = ?self.format,
diff --git a/src/streaming_runtime_operators/src/operators/watermark/watermark_generator.rs b/src/streaming_runtime_operators/src/operators/watermark/watermark_generator.rs
@@ -13,8 +13,12 @@
 use anyhow::{Result, anyhow};
 use arrow::compute::kernels::aggregate;
 use arrow_array::cast::AsArray;
-use arrow_array::types::TimestampNanosecondType;
+use arrow_array::types::{
+    TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
+    TimestampSecondType,
+};
 use arrow_array::{RecordBatch, TimestampNanosecondArray};
+use arrow_schema::{DataType, TimeUnit};
 use bincode::{Decode, Encode};
 use datafusion::physical_expr::PhysicalExpr;
 use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec;
@@ -78,9 +82,32 @@ impl WatermarkGeneratorOperator {
 
     fn extract_max_timestamp(&self, batch: &RecordBatch) -> Option<SystemTime> {
         let ts_column = batch.column(self.timestamp_index);
-        let arr = ts_column.as_primitive::<TimestampNanosecondType>();
-        let max_ts = aggregate::max(arr)?;
-        Some(from_nanos(max_ts as u128))
+        match ts_column.data_type() {
+            DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+                let arr = ts_column.as_primitive::<TimestampNanosecondType>();
+                aggregate::max(arr).map(|v| from_nanos(v as u128))
+            }
+            DataType::Timestamp(TimeUnit::Microsecond, _) => {
+                let arr = ts_column.as_primitive::<TimestampMicrosecondType>();
+                aggregate::max(arr).map(|v| from_nanos((v as u128) * 1_000))
+            }
+            DataType::Timestamp(TimeUnit::Millisecond, _) => {
+                let arr = ts_column.as_primitive::<TimestampMillisecondType>();
+                aggregate::max(arr).map(|v| from_nanos((v as u128) * 1_000_000))
+            }
+            DataType::Timestamp(TimeUnit::Second, _) => {
+                let arr = ts_column.as_primitive::<TimestampSecondType>();
+                aggregate::max(arr).map(|v| from_nanos((v as u128) * 1_000_000_000))
+            }
+            other => {
+                debug!(
+                    ?other,
+                    index = self.timestamp_index,
+                    "skip max timestamp: column is not a timestamp array"
+                );
+                None
+            }
+        }
     }
 
     fn evaluate_watermark(&self, batch: &RecordBatch) -> Result<SystemTime> {
