use arrow::datatypes::{DataType as ArrowDataType, Field, Fields, Schema};
use thiserror::Error;
use crate::export::arrow_convert::cql_type_to_arrow_data_type;
use crate::schema::{CqlType, TableSchema};
const RANGE_BOUND_INCLUSIVE_FIELD: &str = "inclusive";
#[derive(Debug, Error)]
pub enum DeltaSchemaError {
#[error(
"Column '{column}' collides with envelope reserved name '{reserved}'. \
Use DeltaSchemaOpts::envelope_prefix to choose a different prefix \
(e.g. envelope_prefix = \"_cqlite_\" gives \"_cqlite_op\", \"_cqlite_ts\", etc.)."
)]
ColumnCollision {
column: String,
reserved: String,
},
#[error(
"Counter tables cannot be projected to the delta envelope. \
Table '{keyspace}.{table}' contains counter column(s): {columns}. \
Counter semantics (distributed add/subtract) are incompatible with \
the per-cell writetime delta model."
)]
CounterTable {
keyspace: String,
table: String,
columns: String,
},
#[error("CQL type parse error for column '{column}': {source}")]
CqlTypeParse {
column: String,
#[source]
source: crate::error::Error,
},
}
#[derive(Debug, Clone)]
pub struct DeltaSchemaOpts {
pub envelope_prefix: String,
}
impl Default for DeltaSchemaOpts {
fn default() -> Self {
Self {
envelope_prefix: "__".to_string(),
}
}
}
impl DeltaSchemaOpts {
pub fn with_prefix(prefix: impl Into<String>) -> Self {
Self {
envelope_prefix: prefix.into(),
}
}
pub fn op_col(&self) -> String {
format!("{}op", self.envelope_prefix)
}
pub fn ts_col(&self) -> String {
format!("{}ts", self.envelope_prefix)
}
pub fn range_start_col(&self) -> String {
format!("{}range_start", self.envelope_prefix)
}
pub fn range_end_col(&self) -> String {
format!("{}range_end", self.envelope_prefix)
}
fn reserved_names(&self) -> [String; 4] {
[
self.op_col(),
self.ts_col(),
self.range_start_col(),
self.range_end_col(),
]
}
}
pub fn derive_delta_schema(
table: &TableSchema,
opts: &DeltaSchemaOpts,
) -> Result<Schema, DeltaSchemaError> {
let counter_cols: Vec<String> = table
.columns
.iter()
.filter(|col| {
CqlType::parse(&col.data_type)
.map(|t| is_counter_type(&t))
.unwrap_or(false)
})
.map(|col| col.name.clone())
.collect();
if !counter_cols.is_empty() {
return Err(DeltaSchemaError::CounterTable {
keyspace: table.keyspace.clone(),
table: table.table.clone(),
columns: counter_cols.join(", "),
});
}
let reserved = opts.reserved_names();
let all_column_names = table
.partition_keys
.iter()
.map(|k| &k.name)
.chain(table.clustering_keys.iter().map(|k| &k.name))
.chain(table.columns.iter().map(|c| &c.name));
for col_name in all_column_names {
for res in &reserved {
if col_name == res {
return Err(DeltaSchemaError::ColumnCollision {
column: col_name.clone(),
reserved: res.clone(),
});
}
}
}
for ck in &table.clustering_keys {
if ck.name == RANGE_BOUND_INCLUSIVE_FIELD {
return Err(DeltaSchemaError::ColumnCollision {
column: ck.name.clone(),
reserved: RANGE_BOUND_INCLUSIVE_FIELD.to_string(),
});
}
}
let mut fields: Vec<Field> = Vec::new();
let ordered_pk = table.ordered_partition_keys();
for key_col in &ordered_pk {
let cql_type =
CqlType::parse(&key_col.data_type).map_err(|e| DeltaSchemaError::CqlTypeParse {
column: key_col.name.clone(),
source: e,
})?;
let arrow_type = cql_type_to_arrow_data_type(&cql_type);
fields.push(Field::new(&key_col.name, arrow_type, false));
}
let ordered_ck = table.ordered_clustering_keys();
for ck_col in &ordered_ck {
let cql_type =
CqlType::parse(&ck_col.data_type).map_err(|e| DeltaSchemaError::CqlTypeParse {
column: ck_col.name.clone(),
source: e,
})?;
let arrow_type = cql_type_to_arrow_data_type(&cql_type);
fields.push(Field::new(&ck_col.name, arrow_type, true));
}
let pk_names: std::collections::HashSet<&str> =
ordered_pk.iter().map(|k| k.name.as_str()).collect();
let ck_names: std::collections::HashSet<&str> =
ordered_ck.iter().map(|k| k.name.as_str()).collect();
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| DeltaSchemaError::CqlTypeParse {
column: col.name.clone(),
source: e,
})?;
let cell_field = build_cell_struct_field(&col.name, &cql_type);
fields.push(cell_field);
}
fields.push(build_op_field(&opts.op_col()));
fields.push(Field::new(opts.ts_col(), ArrowDataType::Int64, true));
fields.push(build_range_bound_field(
&opts.range_start_col(),
&ordered_ck,
)?);
fields.push(build_range_bound_field(&opts.range_end_col(), &ordered_ck)?);
Ok(Schema::new(fields))
}
fn is_counter_type(cql_type: &CqlType) -> bool {
match cql_type {
CqlType::Counter => true,
CqlType::Frozen(inner) => is_counter_type(inner),
_ => false,
}
}
fn is_collection_type(cql_type: &CqlType) -> bool {
match cql_type {
CqlType::List(_) | CqlType::Set(_) | CqlType::Map(_, _) => true,
_ => false,
}
}
fn build_cell_struct_field(col_name: &str, cql_type: &CqlType) -> Field {
let value_arrow_type = cql_type_to_arrow_data_type(cql_type);
let is_collection = is_collection_type(cql_type);
let mut struct_fields = vec![
Field::new("value", value_arrow_type, true),
Field::new("writetime", ArrowDataType::Int64, false),
Field::new("expires_at", ArrowDataType::Int64, true),
];
if is_collection {
struct_fields.push(Field::new("replaced", ArrowDataType::Boolean, false));
}
Field::new(
col_name,
ArrowDataType::Struct(Fields::from(struct_fields)),
true, )
}
fn build_op_field(col_name: &str) -> Field {
Field::new(
col_name,
ArrowDataType::Dictionary(Box::new(ArrowDataType::Int8), Box::new(ArrowDataType::Utf8)),
false,
)
}
fn build_range_bound_field(
col_name: &str,
clustering_keys: &[&crate::schema::ClusteringColumn],
) -> Result<Field, DeltaSchemaError> {
let mut struct_fields: Vec<Field> = Vec::with_capacity(clustering_keys.len() + 1);
for ck_col in clustering_keys {
let cql_type =
CqlType::parse(&ck_col.data_type).map_err(|e| DeltaSchemaError::CqlTypeParse {
column: ck_col.name.clone(),
source: e,
})?;
let arrow_type = cql_type_to_arrow_data_type(&cql_type);
struct_fields.push(Field::new(&ck_col.name, arrow_type, true));
}
struct_fields.push(Field::new(
RANGE_BOUND_INCLUSIVE_FIELD,
ArrowDataType::Boolean,
false,
));
Ok(Field::new(
col_name,
ArrowDataType::Struct(Fields::from(struct_fields)),
true, ))
}
#[cfg(test)]
#[path = "delta_schema_tests.rs"]
mod tests;