diff --git a/core/src/writer.rs b/core/src/writer.rs index 4c1e353..c8812d5 100644 --- a/core/src/writer.rs +++ b/core/src/writer.rs @@ -18,7 +18,7 @@ use std::io; use arrow_array::*; -use arrow_schema::Schema; +use arrow_schema::{DataType, Field, Schema}; use crate::bucket_writer::{BucketWriter, PagedBucketOutput}; use crate::schema::MosaicSchema; @@ -48,6 +48,31 @@ fn check_zstd_block_size(size: usize, field: &str) -> io::Result<()> { Ok(()) } +fn data_types_match(expected: &DataType, actual: &DataType) -> bool { + match (expected, actual) { + (DataType::List(expected_field), DataType::List(actual_field)) => { + fields_match(expected_field, actual_field) + } + ( + DataType::Map(expected_field, expected_sorted), + DataType::Map(actual_field, actual_sorted), + ) => expected_sorted == actual_sorted && fields_match(expected_field, actual_field), + (DataType::Struct(expected_fields), DataType::Struct(actual_fields)) => { + expected_fields.len() == actual_fields.len() + && expected_fields.iter().zip(actual_fields.iter()).all( + |(expected_field, actual_field)| fields_match(expected_field, actual_field), + ) + } + _ => expected == actual, + } +} + +fn fields_match(expected: &Field, actual: &Field) -> bool { + expected.name() == actual.name() + && expected.is_nullable() == actual.is_nullable() + && data_types_match(expected.data_type(), actual.data_type()) +} + pub trait OutputFile { fn write(&mut self, data: &[u8]) -> io::Result<()>; fn flush(&mut self) -> io::Result<()>; @@ -299,6 +324,7 @@ impl MosaicWriter { ), )); } + self.validate_batch_schema(batch)?; for (i, col) in self.schema.columns.iter().enumerate() { if !col.nullable && batch.column(self.batch_col_map[i]).null_count() > 0 { @@ -341,6 +367,36 @@ impl MosaicWriter { Ok(()) } + fn validate_batch_schema(&self, batch: &RecordBatch) -> io::Result<()> { + let batch_schema = batch.schema(); + for (i, col) in self.schema.columns.iter().enumerate() { + let batch_field = batch_schema.field(self.batch_col_map[i]); + if batch_field.name() != &col.name { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "field name mismatch at column {}: schema has '{}' but batch has '{}'", + i, + col.name, + batch_field.name() + ), + )); + } + if !data_types_match(&col.data_type, batch_field.data_type()) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "field type mismatch for column '{}': schema has {:?} but batch has {:?}", + col.name, + col.data_type, + batch_field.data_type() + ), + )); + } + } + Ok(()) + } + fn flush_row_group(&mut self) -> io::Result<()> { if self.current_row_group_rows == 0 { return Ok(()); @@ -831,7 +887,7 @@ mod tests { let batch2 = RecordBatch::try_new( Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), + Field::new("id", DataType::Int32, true), Field::new("name", DataType::Utf8, true), ])), vec![ @@ -842,6 +898,57 @@ mod tests { .unwrap(); let result = writer.write_batch(&batch2); assert!(result.is_ok()); + + let batch3 = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + ])), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec![None, Some("world")])), + ], + ) + .unwrap(); + let result = writer.write_batch(&batch3); + assert!(result.is_ok()); + } + + #[test] + fn test_write_batch_rejects_reordered_schema() { + let arrow_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let out = MemOutputFile::new(); + let mut writer = MosaicWriter::new(out, &arrow_schema, WriterOptions::default()).unwrap(); + + let batch = RecordBatch::try_new( + Arc::new(arrow_schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![10, 20])), + ], + ) + .unwrap(); + writer.write_batch(&batch).unwrap(); + + let reordered_schema = Schema::new(vec![ + Field::new("b", DataType::Int32, false), + Field::new("a", DataType::Int32, false), + ]); + let reordered_batch = RecordBatch::try_new( + Arc::new(reordered_schema), + vec![ + Arc::new(Int32Array::from(vec![30, 40])), + Arc::new(Int32Array::from(vec![3, 4])), + ], + ) + .unwrap(); + + let err = writer.write_batch(&reordered_batch).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidInput); + assert!(err.to_string().contains("field name mismatch")); } #[test]