use std::io::Write;
use std::sync::Arc;
use arrow::array::{ArrayRef, BooleanArray, Int64Array, StringDictionaryBuilder, StructArray};
use arrow::buffer::NullBuffer;
use arrow::datatypes::{DataType as ArrowDataType, Field, Fields, Int8Type, Schema};
use arrow::record_batch::RecordBatch;
use parquet::arrow::ArrowWriter;
use parquet::basic::{Compression, ZstdLevel};
use parquet::file::metadata::KeyValue;
use parquet::file::properties::WriterProperties;
use thiserror::Error;
use crate::export::arrow_convert::{build_typed_value_array, ArrowConvertError};
use crate::export::delta_schema::{derive_delta_schema, DeltaSchemaError, DeltaSchemaOpts};
use crate::schema::{CqlType, TableSchema};
use crate::storage::sstable::reader::delta_scan::{CellDelta, DeltaRecord, RangeBound};
use crate::types::Value;
#[derive(Debug, Error)]
pub enum DeltaParquetError {
#[error("delta schema error: {0}")]
Schema(#[from] DeltaSchemaError),
#[error("arrow error: {0}")]
Arrow(#[from] arrow::error::ArrowError),
#[error("parquet error: {0}")]
Parquet(#[from] ::parquet::errors::ParquetError),
#[error("{0}")]
InvalidValue(String),
#[error("invalid options: {0}")]
InvalidOptions(String),
#[error("writer already finalized")]
AlreadyFinalized,
}
impl From<ArrowConvertError> for DeltaParquetError {
fn from(e: ArrowConvertError) -> Self {
match e {
ArrowConvertError::Arrow(a) => DeltaParquetError::Arrow(a),
ArrowConvertError::InvalidValue(s) => DeltaParquetError::InvalidValue(s),
}
}
}
#[derive(Debug, Clone)]
pub struct DeltaParquetOptions {
pub row_group_size: usize,
pub compression: DeltaParquetCompression,
pub schema_opts: DeltaSchemaOpts,
pub source: String,
}
impl Default for DeltaParquetOptions {
fn default() -> Self {
Self {
row_group_size: 10_000,
compression: DeltaParquetCompression::default(),
schema_opts: DeltaSchemaOpts::default(),
source: String::new(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum DeltaParquetCompression {
#[default]
Snappy,
Zstd,
Uncompressed,
}
impl DeltaParquetCompression {
fn to_parquet(self) -> Compression {
match self {
DeltaParquetCompression::Snappy => Compression::SNAPPY,
DeltaParquetCompression::Zstd => Compression::ZSTD(ZstdLevel::default()),
DeltaParquetCompression::Uncompressed => Compression::UNCOMPRESSED,
}
}
}
pub struct DeltaParquetWriter<W: Write + Send> {
writer: Option<ArrowWriter<W>>,
schema: Arc<Schema>,
field_meta: Arc<FieldMeta>,
buffer: Vec<DeltaRecord>,
row_group_size: usize,
records_written: u64,
source: String,
schema_hash: String,
}
struct FieldMeta {
partition_keys: Vec<KeyColMeta>,
clustering_keys: Vec<KeyColMeta>,
value_cols: Vec<ValueColMeta>,
}
struct KeyColMeta {
name: String,
cql_type: CqlType,
}
struct ValueColMeta {
name: String,
cql_type: CqlType,
is_collection: bool,
is_static: bool,
}
impl<W: Write + Send> DeltaParquetWriter<W> {
pub fn new(
output: W,
table: &TableSchema,
options: DeltaParquetOptions,
) -> Result<Self, DeltaParquetError> {
if options.row_group_size == 0 {
return Err(DeltaParquetError::InvalidOptions(
"row_group_size must be > 0".into(),
));
}
let schema = derive_delta_schema(table, &options.schema_opts)?;
let schema = Arc::new(schema);
let field_meta = Arc::new(build_field_meta(table, &options.schema_opts)?);
let schema_hash = compute_schema_hash(table);
let props = WriterProperties::builder()
.set_compression(options.compression.to_parquet())
.set_max_row_group_size(options.row_group_size)
.build();
let arrow_writer = ArrowWriter::try_new(output, Arc::clone(&schema), Some(props))?;
Ok(Self {
writer: Some(arrow_writer),
schema,
field_meta,
buffer: Vec::with_capacity(options.row_group_size),
row_group_size: options.row_group_size,
records_written: 0,
source: options.source,
schema_hash,
})
}
pub fn write_record(&mut self, record: DeltaRecord) -> Result<(), DeltaParquetError> {
self.buffer.push(record);
self.records_written += 1;
if self.buffer.len() >= self.row_group_size {
let chunk =
std::mem::replace(&mut self.buffer, Vec::with_capacity(self.row_group_size));
self.flush_chunk(&chunk)?;
}
Ok(())
}
pub fn finalize(mut self) -> Result<(), DeltaParquetError> {
if !self.buffer.is_empty() {
let chunk = std::mem::take(&mut self.buffer);
self.flush_chunk(&chunk)?;
}
let mut writer = self
.writer
.take()
.ok_or(DeltaParquetError::AlreadyFinalized)?;
writer.append_key_value_metadata(KeyValue::new(
"cqlite.delta.version".to_string(),
"1".to_string(),
));
writer.append_key_value_metadata(KeyValue::new(
"cqlite.delta.source".to_string(),
self.source.clone(),
));
writer.append_key_value_metadata(KeyValue::new(
"cqlite.delta.schema_hash".to_string(),
self.schema_hash.clone(),
));
writer.append_key_value_metadata(KeyValue::new(
"cqlite.version".to_string(),
env!("CARGO_PKG_VERSION").to_string(),
));
writer.close()?;
Ok(())
}
pub fn records_written(&self) -> u64 {
self.records_written
}
fn flush_chunk(&mut self, records: &[DeltaRecord]) -> Result<(), DeltaParquetError> {
let writer = self
.writer
.as_mut()
.ok_or(DeltaParquetError::AlreadyFinalized)?;
let batch = build_record_batch(records, &self.schema, &self.field_meta)?;
writer.write(&batch)?;
Ok(())
}
}
impl<W: Write + Send> Drop for DeltaParquetWriter<W> {
fn drop(&mut self) {
debug_assert!(
self.writer.is_none(),
"DeltaParquetWriter dropped without calling finalize(); \
the Parquet footer was never written — output is corrupt"
);
}
}
pub fn write_delta_records_to_bytes(
records: impl IntoIterator<Item = DeltaRecord>,
table: &TableSchema,
options: DeltaParquetOptions,
) -> Result<Vec<u8>, DeltaParquetError> {
let mut buf = Vec::new();
let mut writer = DeltaParquetWriter::new(&mut buf, table, options)?;
for record in records {
writer.write_record(record)?;
}
writer.finalize()?;
Ok(buf)
}
static NULL_SENTINEL: Value = Value::Null;
#[inline]
fn as_value_ref(opt: Option<&Value>) -> Option<&Value> {
Some(opt.unwrap_or(&NULL_SENTINEL))
}
fn build_record_batch(
records: &[DeltaRecord],
schema: &Arc<Schema>,
meta: &FieldMeta,
) -> Result<RecordBatch, DeltaParquetError> {
let mut arrays: Vec<ArrayRef> = Vec::with_capacity(schema.fields().len());
for (pk_idx, pk) in meta.partition_keys.iter().enumerate() {
let values: Vec<Option<&Value>> = records
.iter()
.map(|r| {
let pk_vals = r.partition_key();
as_value_ref(pk_vals.get(pk_idx))
})
.collect();
let arr = build_typed_value_array(&pk.cql_type, &values)?;
arrays.push(arr);
}
for (ck_idx, ck) in meta.clustering_keys.iter().enumerate() {
let values: Vec<Option<&Value>> = records
.iter()
.map(|r| {
let ck_val = match r {
DeltaRecord::Upsert { keys, .. } => keys.clustering.get(ck_idx),
DeltaRecord::RowDelete { keys, .. } => keys.clustering.get(ck_idx),
_ => None,
};
as_value_ref(ck_val)
})
.collect();
let arr = build_typed_value_array(&ck.cql_type, &values)?;
arrays.push(arr);
}
for vcol in &meta.value_cols {
let arr = build_cell_struct_array(records, vcol)?;
arrays.push(arr);
}
{
let mut builder: StringDictionaryBuilder<Int8Type> = StringDictionaryBuilder::new();
for r in records {
builder.append_value(r.op_name());
}
arrays.push(Arc::new(builder.finish()));
}
{
let vals: Vec<Option<i64>> = records
.iter()
.map(|r| match r {
DeltaRecord::Upsert { liveness, .. } => liveness.as_ref().map(|l| l.writetime),
DeltaRecord::StaticUpsert { .. } => None,
DeltaRecord::RowDelete { deleted_at, .. } => Some(*deleted_at),
DeltaRecord::RangeDelete { deleted_at, .. } => Some(*deleted_at),
DeltaRecord::PartitionDelete { deleted_at, .. } => Some(*deleted_at),
})
.collect();
arrays.push(Arc::new(Int64Array::from(vals)));
}
let range_start_arr = build_range_bound_array(records, &meta.clustering_keys, true)?;
arrays.push(range_start_arr);
let range_end_arr = build_range_bound_array(records, &meta.clustering_keys, false)?;
arrays.push(range_end_arr);
let batch = RecordBatch::try_new(Arc::clone(schema), arrays)?;
Ok(batch)
}
fn build_cell_struct_array(
records: &[DeltaRecord],
vcol: &ValueColMeta,
) -> Result<ArrayRef, DeltaParquetError> {
let cell_deltas: Vec<Option<&CellDelta>> = records
.iter()
.map(|r| find_cell_delta(r, &vcol.name, vcol.is_static))
.collect();
let struct_valid: Vec<bool> = cell_deltas.iter().map(|d| d.is_some()).collect();
let values: Vec<Option<&Value>> = cell_deltas
.iter()
.map(|d| as_value_ref(d.and_then(|cd| cd.value.as_ref())))
.collect();
let value_arr = build_typed_value_array(&vcol.cql_type, &values)?;
let wt_vals: Vec<i64> = cell_deltas
.iter()
.map(|d| d.map(|cd| cd.writetime).unwrap_or(0))
.collect();
let wt_arr: ArrayRef = Arc::new(Int64Array::from(wt_vals));
let ea_vals: Vec<Option<i64>> = cell_deltas
.iter()
.map(|d| d.and_then(|cd| cd.expires_at))
.collect();
let ea_arr: ArrayRef = Arc::new(Int64Array::from(ea_vals));
let mut child_fields: Vec<Field> = vec![
Field::new("value", value_arr.data_type().clone(), true),
Field::new("writetime", ArrowDataType::Int64, false),
Field::new("expires_at", ArrowDataType::Int64, true),
];
let mut child_arrays: Vec<ArrayRef> = vec![value_arr, wt_arr, ea_arr];
if vcol.is_collection {
let replaced_vals: Vec<bool> = cell_deltas
.iter()
.map(|d| d.map(|cd| cd.replaced).unwrap_or(false))
.collect();
let replaced_arr: ArrayRef = Arc::new(BooleanArray::from(replaced_vals));
child_fields.push(Field::new("replaced", ArrowDataType::Boolean, false));
child_arrays.push(replaced_arr);
}
let null_buffer = NullBuffer::from(struct_valid);
let struct_arr =
StructArray::try_new(Fields::from(child_fields), child_arrays, Some(null_buffer))?;
Ok(Arc::new(struct_arr))
}
fn find_cell_delta<'a>(
record: &'a DeltaRecord,
col_name: &str,
is_static: bool,
) -> Option<&'a CellDelta> {
match record {
DeltaRecord::Upsert { cells, .. } if !is_static => cells
.iter()
.find(|(id, _)| id.name() == col_name)
.map(|(_, cd)| cd),
DeltaRecord::StaticUpsert { cells, .. } if is_static => cells
.iter()
.find(|(id, _)| id.name() == col_name)
.map(|(_, cd)| cd),
_ => None,
}
}
fn build_range_bound_array(
records: &[DeltaRecord],
clustering_keys: &[KeyColMeta],
is_start: bool,
) -> Result<ArrayRef, DeltaParquetError> {
let bounds: Vec<Option<&RangeBound>> = records
.iter()
.map(|r| match r {
DeltaRecord::RangeDelete { start, end, .. } => {
if is_start {
Some(start)
} else {
Some(end)
}
}
_ => None,
})
.collect();
let struct_valid: Vec<bool> = bounds.iter().map(|b| b.is_some()).collect();
let mut child_fields: Vec<Field> = Vec::new();
let mut child_arrays: Vec<ArrayRef> = Vec::new();
for (ck_idx, ck) in clustering_keys.iter().enumerate() {
let values: Vec<Option<&Value>> = bounds
.iter()
.map(|b| as_value_ref(b.and_then(|rb| rb.values.get(ck_idx))))
.collect();
let arr = build_typed_value_array(&ck.cql_type, &values)?;
child_fields.push(Field::new(&ck.name, arr.data_type().clone(), true));
child_arrays.push(arr);
}
let inclusive_vals: Vec<bool> = bounds
.iter()
.map(|b| b.map(|rb| rb.inclusive).unwrap_or(false))
.collect();
child_fields.push(Field::new("inclusive", ArrowDataType::Boolean, false));
child_arrays.push(Arc::new(BooleanArray::from(inclusive_vals)));
let null_buffer = NullBuffer::from(struct_valid);
let struct_arr =
StructArray::try_new(Fields::from(child_fields), child_arrays, Some(null_buffer))?;
Ok(Arc::new(struct_arr))
}
fn build_field_meta(
table: &TableSchema,
_opts: &DeltaSchemaOpts,
) -> Result<FieldMeta, DeltaParquetError> {
let pk_names: std::collections::HashSet<&str> = table
.partition_keys
.iter()
.map(|k| k.name.as_str())
.collect();
let ck_names: std::collections::HashSet<&str> = table
.clustering_keys
.iter()
.map(|k| k.name.as_str())
.collect();
let mut partition_keys = Vec::new();
for key in table.ordered_partition_keys() {
let cql_type = CqlType::parse(&key.data_type).map_err(|e| {
DeltaParquetError::InvalidValue(format!(
"cannot parse CQL type '{}' for partition key '{}': {}",
key.data_type, key.name, e
))
})?;
partition_keys.push(KeyColMeta {
name: key.name.clone(),
cql_type,
});
}
let mut clustering_keys = Vec::new();
for ck in table.ordered_clustering_keys() {
let cql_type = CqlType::parse(&ck.data_type).map_err(|e| {
DeltaParquetError::InvalidValue(format!(
"cannot parse CQL type '{}' for clustering key '{}': {}",
ck.data_type, ck.name, e
))
})?;
clustering_keys.push(KeyColMeta {
name: ck.name.clone(),
cql_type,
});
}
let mut value_cols = Vec::new();
for col in &table.columns {
if pk_names.contains(col.name.as_str()) || ck_names.contains(col.name.as_str()) {
continue;
}
let cql_type = CqlType::parse(&col.data_type).map_err(|e| {
DeltaParquetError::InvalidValue(format!(
"cannot parse CQL type '{}' for column '{}': {}",
col.data_type, col.name, e
))
})?;
let is_collection = matches!(
&cql_type,
CqlType::List(_) | CqlType::Set(_) | CqlType::Map(_, _)
);
value_cols.push(ValueColMeta {
name: col.name.clone(),
cql_type,
is_collection,
is_static: col.is_static,
});
}
Ok(FieldMeta {
partition_keys,
clustering_keys,
value_cols,
})
}
fn compute_schema_hash(table: &TableSchema) -> String {
const FNV_OFFSET_BASIS: u64 = 14695981039346656037;
const FNV_PRIME: u64 = 1099511628211;
fn fnv1a_update(mut hash: u64, bytes: &[u8]) -> u64 {
for &b in bytes {
hash ^= b as u64;
hash = hash.wrapping_mul(FNV_PRIME);
}
hash
}
let mut hash = FNV_OFFSET_BASIS;
hash = fnv1a_update(hash, table.keyspace.as_bytes());
hash = fnv1a_update(hash, b".");
hash = fnv1a_update(hash, table.table.as_bytes());
hash = fnv1a_update(hash, b":");
for k in &table.partition_keys {
hash = fnv1a_update(hash, k.name.as_bytes());
hash = fnv1a_update(hash, b" ");
hash = fnv1a_update(hash, k.data_type.as_bytes());
hash = fnv1a_update(hash, b",");
}
hash = fnv1a_update(hash, b";");
for k in &table.clustering_keys {
hash = fnv1a_update(hash, k.name.as_bytes());
hash = fnv1a_update(hash, b" ");
hash = fnv1a_update(hash, k.data_type.as_bytes());
hash = fnv1a_update(hash, b",");
}
hash = fnv1a_update(hash, b";");
for c in &table.columns {
hash = fnv1a_update(hash, c.name.as_bytes());
hash = fnv1a_update(hash, b" ");
hash = fnv1a_update(hash, c.data_type.as_bytes());
hash = fnv1a_update(hash, if c.is_static { b" static," } else { b"," });
}
format!("{:016x}", hash)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::schema::{ClusteringColumn, ClusteringOrder, Column, KeyColumn, TableSchema};
use crate::storage::sstable::reader::delta_scan::{CellDelta, RangeBound, RowKeys};
use crate::types::{ColumnId, Value};
use arrow::array::{
Array, BooleanArray, DictionaryArray, Int32Array, Int64Array, StringArray, StructArray,
};
use arrow::datatypes::Int8Type;
use bytes::Bytes;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::collections::HashMap;
fn example_schema() -> TableSchema {
TableSchema {
keyspace: "test_ks".into(),
table: "t".into(),
partition_keys: vec![KeyColumn {
name: "pk".into(),
data_type: "int".into(),
position: 0,
}],
clustering_keys: vec![ClusteringColumn {
name: "ck".into(),
data_type: "text".into(),
position: 0,
order: ClusteringOrder::Asc,
}],
columns: vec![
Column {
name: "val".into(),
data_type: "text".into(),
is_static: false,
nullable: true,
default: None,
},
Column {
name: "st".into(),
data_type: "text".into(),
is_static: true,
nullable: true,
default: None,
},
],
comments: HashMap::new(),
dropped_columns: HashMap::new(),
}
}
fn round_trip(
records: Vec<DeltaRecord>,
schema: &TableSchema,
opts: DeltaParquetOptions,
) -> RecordBatch {
let bytes = write_delta_records_to_bytes(records, schema, opts).expect("write failed");
assert!(!bytes.is_empty(), "output must not be empty");
assert_eq!(&bytes[0..4], b"PAR1", "must start with Parquet magic");
let bytes = Bytes::copy_from_slice(&bytes);
let builder = ParquetRecordBatchReaderBuilder::try_new(bytes).expect("reader builder");
let mut reader = builder.build().expect("reader build");
reader
.next()
.expect("at least one batch")
.expect("batch read ok")
}
fn read_footer_metadata(bytes: &[u8]) -> HashMap<String, String> {
let bytes = Bytes::copy_from_slice(bytes);
let builder = ParquetRecordBatchReaderBuilder::try_new(bytes).expect("reader builder");
let meta = builder.metadata().clone();
let kv = meta
.file_metadata()
.key_value_metadata()
.cloned()
.unwrap_or_default();
kv.into_iter()
.filter_map(|kv| kv.value.map(|v| (kv.key, v)))
.collect()
}
#[test]
fn test_upsert_partial_update() {
let record = DeltaRecord::Upsert {
keys: RowKeys::new(vec![Value::Integer(1)], vec![Value::Text("a".into())]),
liveness: None,
cells: vec![(
ColumnId::new("val"),
CellDelta::value(Value::Text("x".into()), 1_000_000),
)],
};
let schema = example_schema();
let batch = round_trip(vec![record], &schema, DeltaParquetOptions::default());
assert_eq!(batch.num_rows(), 1);
let pk_col = batch.column_by_name("pk").expect("pk column");
let pk_arr = pk_col.as_any().downcast_ref::<Int32Array>().expect("Int32");
assert_eq!(pk_arr.value(0), 1);
let ck_col = batch.column_by_name("ck").expect("ck column");
let ck_arr = ck_col
.as_any()
.downcast_ref::<StringArray>()
.expect("String");
assert_eq!(ck_arr.value(0), "a");
let op_col = batch.column_by_name("__op").expect("__op");
let op_dict = op_col
.as_any()
.downcast_ref::<DictionaryArray<Int8Type>>()
.expect("Dict<Int8, Utf8>");
let op_values = op_dict
.values()
.as_any()
.downcast_ref::<StringArray>()
.expect("String values");
let op_key = op_dict.key(0).expect("key 0");
assert_eq!(op_values.value(op_key as usize), "upsert");
let ts_col = batch.column_by_name("__ts").expect("__ts");
let ts_arr = ts_col.as_any().downcast_ref::<Int64Array>().expect("Int64");
assert!(
ts_arr.is_null(0),
"__ts must be null for UPDATE (no liveness)"
);
let val_col = batch.column_by_name("val").expect("val column");
let val_struct = val_col
.as_any()
.downcast_ref::<StructArray>()
.expect("StructArray");
assert!(val_struct.is_valid(0), "val struct must be non-null");
let value_field = val_struct.column_by_name("value").expect("value field");
let value_str = value_field
.as_any()
.downcast_ref::<StringArray>()
.expect("String");
assert_eq!(value_str.value(0), "x");
let wt_field = val_struct.column_by_name("writetime").expect("writetime");
let wt_arr = wt_field
.as_any()
.downcast_ref::<Int64Array>()
.expect("Int64");
assert_eq!(wt_arr.value(0), 1_000_000);
let ea_field = val_struct.column_by_name("expires_at").expect("expires_at");
let ea_arr = ea_field
.as_any()
.downcast_ref::<Int64Array>()
.expect("Int64");
assert!(ea_arr.is_null(0), "expires_at must be null (no TTL)");
let st_col = batch.column_by_name("st").expect("st column");
let st_struct = st_col
.as_any()
.downcast_ref::<StructArray>()
.expect("StructArray");
assert!(
st_struct.is_null(0),
"st struct must be null (absent column)"
);
}
#[test]
fn test_cell_tombstone_upsert() {
let record = DeltaRecord::Upsert {
keys: RowKeys::new(vec![Value::Integer(1)], vec![Value::Text("a".into())]),
liveness: None,
cells: vec![(ColumnId::new("val"), CellDelta::tombstone(2_000_000))],
};
let schema = example_schema();
let batch = round_trip(vec![record], &schema, DeltaParquetOptions::default());
assert_eq!(batch.num_rows(), 1);
let val_col = batch.column_by_name("val").expect("val column");
let val_struct = val_col
.as_any()
.downcast_ref::<StructArray>()
.expect("StructArray");
assert!(
val_struct.is_valid(0),
"val struct must be NON-null for cell tombstone (column is present)"
);
let value_field = val_struct.column_by_name("value").expect("value");
assert!(
value_field.is_null(0),
"value sub-field must be null for cell tombstone"
);
let wt_field = val_struct.column_by_name("writetime").expect("writetime");
let wt_arr = wt_field
.as_any()
.downcast_ref::<Int64Array>()
.expect("Int64");
assert_eq!(wt_arr.value(0), 2_000_000);
}
#[test]
fn test_null_struct_absent_vs_cell_tombstone() {
let rec0 = DeltaRecord::Upsert {
keys: RowKeys::new(vec![Value::Integer(1)], vec![Value::Text("a".into())]),
liveness: None,
cells: vec![(
ColumnId::new("val"),
CellDelta::value(Value::Text("alive".into()), 100),
)],
};
let rec1 = DeltaRecord::Upsert {
keys: RowKeys::new(vec![Value::Integer(2)], vec![Value::Text("b".into())]),
liveness: None,
cells: vec![(ColumnId::new("val"), CellDelta::tombstone(200))],
};
let schema = example_schema();
let batch = round_trip(vec![rec0, rec1], &schema, DeltaParquetOptions::default());
assert_eq!(batch.num_rows(), 2);
let val_col = batch.column_by_name("val").expect("val column");
let val_struct = val_col
.as_any()
.downcast_ref::<StructArray>()
.expect("StructArray");
let value_field = val_struct.column_by_name("value").expect("value");
assert!(val_struct.is_valid(0), "row0 val struct must be non-null");
assert!(!value_field.is_null(0), "row0 val.value must be non-null");
assert!(
val_struct.is_valid(1),
"row1 val struct must be non-null (tombstone)"
);
assert!(
value_field.is_null(1),
"row1 val.value must be null (tombstone)"
);
let st_col = batch.column_by_name("st").expect("st column");
let st_struct = st_col
.as_any()
.downcast_ref::<StructArray>()
.expect("StructArray");
assert!(st_struct.is_null(0), "row0 st must be null (absent)");
assert!(st_struct.is_null(1), "row1 st must be null (absent)");
}
#[test]
fn test_static_upsert() {
let record = DeltaRecord::StaticUpsert {
partition_key: RowKeys::partition_only(vec![Value::Integer(1)]),
cells: vec![(
ColumnId::new("st"),
CellDelta::value(Value::Text("S".into()), 3_000_000),
)],
};
let schema = example_schema();
let batch = round_trip(vec![record], &schema, DeltaParquetOptions::default());
assert_eq!(batch.num_rows(), 1);
let ck_col = batch.column_by_name("ck").expect("ck column");
assert!(ck_col.is_null(0), "ck must be null for static_upsert");
let op_col = batch.column_by_name("__op").expect("__op");
let op_dict = op_col
.as_any()
.downcast_ref::<DictionaryArray<Int8Type>>()
.expect("Dict");
let op_values = op_dict
.values()
.as_any()
.downcast_ref::<StringArray>()
.expect("String");
let op_key = op_dict.key(0).expect("key 0");
assert_eq!(op_values.value(op_key as usize), "static_upsert");
let st_col = batch.column_by_name("st").expect("st column");
let st_struct = st_col
.as_any()
.downcast_ref::<StructArray>()
.expect("Struct");
assert!(st_struct.is_valid(0), "st struct must be non-null");
let st_value = st_struct.column_by_name("value").expect("value");
let st_str = st_value
.as_any()
.downcast_ref::<StringArray>()
.expect("String");
assert_eq!(st_str.value(0), "S");
let val_col = batch.column_by_name("val").expect("val column");
let val_struct = val_col
.as_any()
.downcast_ref::<StructArray>()
.expect("Struct");
assert!(
val_struct.is_null(0),
"val struct must be null for static_upsert"
);
}
#[test]
fn test_row_delete() {
let record = DeltaRecord::RowDelete {
keys: RowKeys::new(vec![Value::Integer(1)], vec![Value::Text("a".into())]),
deleted_at: 4_000_000,
};
let schema = example_schema();
let batch = round_trip(vec![record], &schema, DeltaParquetOptions::default());
assert_eq!(batch.num_rows(), 1);
let op_col = batch.column_by_name("__op").expect("__op");
let op_dict = op_col
.as_any()
.downcast_ref::<DictionaryArray<Int8Type>>()
.expect("Dict");
let op_values = op_dict
.values()
.as_any()
.downcast_ref::<StringArray>()
.expect("String");
let op_key = op_dict.key(0).expect("key 0");
assert_eq!(op_values.value(op_key as usize), "row_delete");
let ts_col = batch.column_by_name("__ts").expect("__ts");
let ts_arr = ts_col.as_any().downcast_ref::<Int64Array>().expect("Int64");
assert!(!ts_arr.is_null(0), "__ts must be non-null for row_delete");
assert_eq!(ts_arr.value(0), 4_000_000);
let val_col = batch.column_by_name("val").expect("val");
let val_struct = val_col
.as_any()
.downcast_ref::<StructArray>()
.expect("Struct");
assert!(val_struct.is_null(0), "val must be null for row_delete");
}
#[test]
fn test_range_delete() {
let record = DeltaRecord::RangeDelete {
partition_key: RowKeys::partition_only(vec![Value::Integer(1)]),
start: RangeBound::inclusive(vec![Value::Text("a".into())]),
end: RangeBound::exclusive(vec![Value::Text("m".into())]),
deleted_at: 5_000_000,
};
let schema = example_schema();
let batch = round_trip(vec![record], &schema, DeltaParquetOptions::default());
assert_eq!(batch.num_rows(), 1);
let op_col = batch.column_by_name("__op").expect("__op");
let op_dict = op_col
.as_any()
.downcast_ref::<DictionaryArray<Int8Type>>()
.expect("Dict");
let op_values = op_dict
.values()
.as_any()
.downcast_ref::<StringArray>()
.expect("String");
let op_key = op_dict.key(0).expect("key 0");
assert_eq!(op_values.value(op_key as usize), "range_delete");
let ts_col = batch.column_by_name("__ts").expect("__ts");
let ts_arr = ts_col.as_any().downcast_ref::<Int64Array>().expect("Int64");
assert_eq!(ts_arr.value(0), 5_000_000);
let ck_col = batch.column_by_name("ck").expect("ck");
assert!(ck_col.is_null(0), "ck must be null for range_delete");
let rs_col = batch
.column_by_name("__range_start")
.expect("__range_start");
let rs_struct = rs_col
.as_any()
.downcast_ref::<StructArray>()
.expect("Struct");
assert!(rs_struct.is_valid(0), "__range_start must be non-null");
let rs_ck = rs_struct.column_by_name("ck").expect("ck in range_start");
let rs_ck_str = rs_ck
.as_any()
.downcast_ref::<StringArray>()
.expect("String");
assert_eq!(rs_ck_str.value(0), "a");
let rs_inc = rs_struct
.column_by_name("inclusive")
.expect("inclusive in range_start");
let rs_inc_bool = rs_inc
.as_any()
.downcast_ref::<BooleanArray>()
.expect("Bool");
assert!(rs_inc_bool.value(0), "range_start must be inclusive");
let re_col = batch.column_by_name("__range_end").expect("__range_end");
let re_struct = re_col
.as_any()
.downcast_ref::<StructArray>()
.expect("Struct");
assert!(re_struct.is_valid(0), "__range_end must be non-null");
let re_ck = re_struct.column_by_name("ck").expect("ck in range_end");
let re_ck_str = re_ck
.as_any()
.downcast_ref::<StringArray>()
.expect("String");
assert_eq!(re_ck_str.value(0), "m");
let re_inc = re_struct
.column_by_name("inclusive")
.expect("inclusive in range_end");
let re_inc_bool = re_inc
.as_any()
.downcast_ref::<BooleanArray>()
.expect("Bool");
assert!(!re_inc_bool.value(0), "range_end must be exclusive");
}
#[test]
fn test_partition_delete() {
let record = DeltaRecord::PartitionDelete {
partition_key: RowKeys::partition_only(vec![Value::Integer(1)]),
deleted_at: 6_000_000,
};
let schema = example_schema();
let batch = round_trip(vec![record], &schema, DeltaParquetOptions::default());
assert_eq!(batch.num_rows(), 1);
let op_col = batch.column_by_name("__op").expect("__op");
let op_dict = op_col
.as_any()
.downcast_ref::<DictionaryArray<Int8Type>>()
.expect("Dict");
let op_values = op_dict
.values()
.as_any()
.downcast_ref::<StringArray>()
.expect("String");
let op_key = op_dict.key(0).expect("key 0");
assert_eq!(op_values.value(op_key as usize), "partition_delete");
let ts_col = batch.column_by_name("__ts").expect("__ts");
let ts_arr = ts_col.as_any().downcast_ref::<Int64Array>().expect("Int64");
assert_eq!(ts_arr.value(0), 6_000_000);
let ck_col = batch.column_by_name("ck").expect("ck");
assert!(ck_col.is_null(0), "ck must be null for partition_delete");
let rs_col = batch
.column_by_name("__range_start")
.expect("__range_start");
let rs_struct = rs_col
.as_any()
.downcast_ref::<StructArray>()
.expect("Struct");
assert!(
rs_struct.is_null(0),
"__range_start must be null for partition_delete"
);
let re_col = batch.column_by_name("__range_end").expect("__range_end");
let re_struct = re_col
.as_any()
.downcast_ref::<StructArray>()
.expect("Struct");
assert!(
re_struct.is_null(0),
"__range_end must be null for partition_delete"
);
}
#[test]
fn test_all_five_op_shapes_in_one_batch() {
let schema = example_schema();
let records = vec![
DeltaRecord::Upsert {
keys: RowKeys::new(vec![Value::Integer(1)], vec![Value::Text("a".into())]),
liveness: None,
cells: vec![(
ColumnId::new("val"),
CellDelta::value(Value::Text("x".into()), 100),
)],
},
DeltaRecord::StaticUpsert {
partition_key: RowKeys::partition_only(vec![Value::Integer(1)]),
cells: vec![(
ColumnId::new("st"),
CellDelta::value(Value::Text("S".into()), 200),
)],
},
DeltaRecord::RowDelete {
keys: RowKeys::new(vec![Value::Integer(2)], vec![Value::Text("b".into())]),
deleted_at: 300,
},
DeltaRecord::RangeDelete {
partition_key: RowKeys::partition_only(vec![Value::Integer(3)]),
start: RangeBound::inclusive(vec![Value::Text("c".into())]),
end: RangeBound::exclusive(vec![Value::Text("z".into())]),
deleted_at: 400,
},
DeltaRecord::PartitionDelete {
partition_key: RowKeys::partition_only(vec![Value::Integer(4)]),
deleted_at: 500,
},
];
let batch = round_trip(records, &schema, DeltaParquetOptions::default());
assert_eq!(batch.num_rows(), 5);
let op_col = batch.column_by_name("__op").expect("__op");
let op_dict = op_col
.as_any()
.downcast_ref::<DictionaryArray<Int8Type>>()
.expect("Dict");
let op_values = op_dict
.values()
.as_any()
.downcast_ref::<StringArray>()
.expect("String");
let expected_ops = [
"upsert",
"static_upsert",
"row_delete",
"range_delete",
"partition_delete",
];
for (i, expected_op) in expected_ops.iter().enumerate() {
let key = op_dict.key(i).expect("key");
assert_eq!(
op_values.value(key as usize),
*expected_op,
"row {i} __op mismatch"
);
}
}
#[test]
fn test_footer_metadata_all_four_keys() {
let schema = example_schema();
let record = DeltaRecord::Upsert {
keys: RowKeys::new(vec![Value::Integer(1)], vec![Value::Text("a".into())]),
liveness: None,
cells: vec![(
ColumnId::new("val"),
CellDelta::value(Value::Text("x".into()), 100),
)],
};
let opts = DeltaParquetOptions {
source: "nb-5-big-Data.db-gen1".to_string(),
..Default::default()
};
let bytes = write_delta_records_to_bytes(vec![record], &schema, opts).expect("write");
let metadata = read_footer_metadata(&bytes);
assert_eq!(
metadata.get("cqlite.delta.version").map(String::as_str),
Some("1"),
"cqlite.delta.version must be '1'"
);
assert_eq!(
metadata.get("cqlite.delta.source").map(String::as_str),
Some("nb-5-big-Data.db-gen1"),
"cqlite.delta.source must match the options"
);
let hash = metadata
.get("cqlite.delta.schema_hash")
.expect("cqlite.delta.schema_hash must be present");
assert!(!hash.is_empty(), "schema_hash must not be empty");
assert_eq!(
metadata.get("cqlite.version").map(String::as_str),
Some(env!("CARGO_PKG_VERSION")),
"cqlite.version must equal CARGO_PKG_VERSION"
);
}
#[test]
fn test_streaming_multiple_row_groups() {
let schema = example_schema();
let records: Vec<DeltaRecord> = (0..25i32)
.map(|i| DeltaRecord::Upsert {
keys: RowKeys::new(vec![Value::Integer(i)], vec![Value::text(format!("ck{i}"))]),
liveness: None,
cells: vec![(
ColumnId::new("val"),
CellDelta::value(Value::text(format!("v{i}")), i as i64 * 1000),
)],
})
.collect();
let opts = DeltaParquetOptions {
row_group_size: 10,
..Default::default()
};
let bytes = write_delta_records_to_bytes(records, &schema, opts).expect("write");
let bytes2 = Bytes::copy_from_slice(&bytes);
let builder = ParquetRecordBatchReaderBuilder::try_new(bytes2).expect("reader builder");
let total: usize = builder
.build()
.expect("reader")
.map(|b| b.expect("batch").num_rows())
.sum();
assert_eq!(total, 25, "all 25 records must be written");
}
#[test]
fn test_schema_hash_deterministic() {
let schema = example_schema();
let h1 = compute_schema_hash(&schema);
let h2 = compute_schema_hash(&schema);
assert_eq!(h1, h2, "schema hash must be deterministic");
}
#[test]
fn test_schema_hash_fixed_value() {
const FNV_OFFSET_BASIS: u64 = 14695981039346656037;
const FNV_PRIME: u64 = 1099511628211;
fn fnv(mut h: u64, b: &[u8]) -> u64 {
for &byte in b {
h ^= byte as u64;
h = h.wrapping_mul(FNV_PRIME);
}
h
}
let mut h = FNV_OFFSET_BASIS;
h = fnv(h, b"test_ks");
h = fnv(h, b".");
h = fnv(h, b"t");
h = fnv(h, b":");
h = fnv(h, b"pk");
h = fnv(h, b" ");
h = fnv(h, b"int");
h = fnv(h, b",");
h = fnv(h, b";");
h = fnv(h, b"ck");
h = fnv(h, b" ");
h = fnv(h, b"text");
h = fnv(h, b",");
h = fnv(h, b";");
h = fnv(h, b"val");
h = fnv(h, b" ");
h = fnv(h, b"text");
h = fnv(h, b",");
h = fnv(h, b"st");
h = fnv(h, b" ");
h = fnv(h, b"text");
h = fnv(h, b" static,");
let expected = format!("{:016x}", h);
let actual = compute_schema_hash(&example_schema());
assert_eq!(
actual, expected,
"FNV-1a 64-bit schema hash must match known stable value"
);
assert_eq!(actual.len(), 16, "hash must be 16 hex characters");
}
#[test]
fn test_empty_record_list() {
let schema = example_schema();
let bytes = write_delta_records_to_bytes(vec![], &schema, DeltaParquetOptions::default())
.expect("write empty");
assert!(!bytes.is_empty());
assert_eq!(&bytes[0..4], b"PAR1");
}
}