use crate::query::{ColumnInfo, QueryRow};
use crate::schema::CqlType;
use crate::types::{DataType, Value};
use crate::util::value_fmt::ValueFormatter;
use arrow::array::{
ArrayRef, BinaryArray, BooleanArray, Date32Array, Float32Array, Float64Array, Int16Array,
Int32Array, Int64Array, Int8Array, ListArray, MapArray, StringArray, StructArray,
Time64NanosecondArray, TimestampMillisecondArray,
};
use arrow::buffer::{NullBuffer, OffsetBuffer};
use arrow::datatypes::{DataType as ArrowDataType, Field, Fields, Schema, TimeUnit};
use arrow::record_batch::RecordBatch;
use std::collections::HashMap;
use std::sync::Arc;
use thiserror::Error;
pub(crate) const DECIMAL_FIXED_SCALE: i32 = 9;
pub(crate) const DECIMAL_MAX_PRECISION: u8 = 38;
pub(crate) const ARROW_EXTENSION_NAME_KEY: &str = "ARROW:extension:name";
pub(crate) const ARROW_UUID_EXTENSION_NAME: &str = "arrow.uuid";
#[derive(Debug, Error)]
pub enum ArrowConvertError {
#[error("Arrow error: {0}")]
Arrow(#[from] arrow::error::ArrowError),
#[error("{0}")]
InvalidValue(String),
}
pub fn build_arrow_schema(columns: &[ColumnInfo]) -> Result<Schema, ArrowConvertError> {
let fields: Vec<Field> = columns.iter().map(column_to_field).collect();
Ok(Schema::new(fields))
}
pub fn rows_to_record_batch(
columns: &[ColumnInfo],
rows: &[QueryRow],
) -> Result<RecordBatch, ArrowConvertError> {
let schema = build_arrow_schema(columns)?;
let arrays = convert_to_arrays(columns, rows)?;
let batch = RecordBatch::try_new(Arc::new(schema), arrays)?;
Ok(batch)
}
pub(crate) fn bigint_to_i128(n: &num_bigint::BigInt) -> Result<i128, ArrowConvertError> {
let tc_bytes = n.to_signed_bytes_be();
if tc_bytes.len() > 16 {
return Err(ArrowConvertError::InvalidValue(
"BigInt value requires more than 16 bytes; cannot fit in i128".to_string(),
));
}
let pad: u8 = if n.sign() == num_bigint::Sign::Minus {
0xFF
} else {
0x00
};
let mut buf = [pad; 16];
buf[16 - tc_bytes.len()..].copy_from_slice(&tc_bytes);
Ok(i128::from_be_bytes(buf))
}
pub(crate) fn column_to_field(col: &ColumnInfo) -> Field {
if let Some(cql_type) = &col.cql_type {
if let Some(field) = cql_type_to_arrow_field(&col.name, cql_type, col.nullable) {
return field;
}
}
let arrow_type = data_type_to_arrow(&col.data_type);
Field::new(&col.name, arrow_type, col.nullable)
}
pub(crate) fn cql_type_to_arrow_field(
name: &str,
cql_type: &CqlType,
nullable: bool,
) -> Option<Field> {
match cql_type {
CqlType::Date => Some(Field::new(name, ArrowDataType::Date32, nullable)),
CqlType::Time => Some(Field::new(
name,
ArrowDataType::Time64(TimeUnit::Nanosecond),
nullable,
)),
CqlType::Decimal => Some(Field::new(
name,
ArrowDataType::Decimal128(DECIMAL_MAX_PRECISION, DECIMAL_FIXED_SCALE as i8),
nullable,
)),
CqlType::Varint => {
Some(Field::new(
name,
ArrowDataType::Decimal128(DECIMAL_MAX_PRECISION, 0),
nullable,
))
}
CqlType::Duration => {
Some(Field::new(name, ArrowDataType::Utf8, nullable))
}
CqlType::Uuid | CqlType::TimeUuid => {
let mut meta = HashMap::new();
meta.insert(
ARROW_EXTENSION_NAME_KEY.to_string(),
ARROW_UUID_EXTENSION_NAME.to_string(),
);
Some(Field::new(name, ArrowDataType::FixedSizeBinary(16), nullable).with_metadata(meta))
}
CqlType::Inet => Some(Field::new(name, ArrowDataType::Utf8, nullable)),
CqlType::Counter => Some(Field::new(name, ArrowDataType::Int64, nullable)),
CqlType::List(inner) | CqlType::Set(inner) => {
let item_type = cql_type_to_arrow_data_type(inner);
let item_field = Arc::new(Field::new("item", item_type, true));
Some(Field::new(name, ArrowDataType::List(item_field), nullable))
}
CqlType::Frozen(inner) => cql_type_to_arrow_field(name, inner, nullable),
CqlType::Map(key_type, val_type) => {
let key_arrow = cql_type_to_arrow_data_type(key_type);
let val_arrow = cql_type_to_arrow_data_type(val_type);
let entries_field = Arc::new(Field::new(
"entries",
ArrowDataType::Struct(Fields::from(vec![
Field::new("key", key_arrow, false),
Field::new("value", val_arrow, true),
])),
false,
));
Some(Field::new(
name,
ArrowDataType::Map(entries_field, false),
nullable,
))
}
CqlType::Tuple(element_types) => {
if element_types.is_empty() {
return Some(Field::new(name, ArrowDataType::Utf8, nullable));
}
let struct_type = cql_type_to_arrow_data_type(cql_type);
Some(Field::new(name, struct_type, nullable))
}
CqlType::Udt(_udt_name, udt_fields) => {
if udt_fields.is_empty() {
return Some(Field::new(name, ArrowDataType::Utf8, nullable));
}
let struct_type = cql_type_to_arrow_data_type(cql_type);
Some(Field::new(name, struct_type, nullable))
}
_ => None,
}
}
pub(crate) fn cql_type_to_arrow_data_type(cql_type: &CqlType) -> ArrowDataType {
match cql_type {
CqlType::Boolean => ArrowDataType::Boolean,
CqlType::TinyInt => ArrowDataType::Int8,
CqlType::SmallInt => ArrowDataType::Int16,
CqlType::Int => ArrowDataType::Int32,
CqlType::BigInt => ArrowDataType::Int64,
CqlType::Counter => ArrowDataType::Int64,
CqlType::Float => ArrowDataType::Float32,
CqlType::Double => ArrowDataType::Float64,
CqlType::Text | CqlType::Ascii | CqlType::Varchar => ArrowDataType::Utf8,
CqlType::Blob => ArrowDataType::Binary,
CqlType::Timestamp => ArrowDataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
CqlType::Date => ArrowDataType::Date32,
CqlType::Time => ArrowDataType::Time64(TimeUnit::Nanosecond),
CqlType::Decimal => {
ArrowDataType::Decimal128(DECIMAL_MAX_PRECISION, DECIMAL_FIXED_SCALE as i8)
}
CqlType::Varint => ArrowDataType::Decimal128(DECIMAL_MAX_PRECISION, 0),
CqlType::Duration => ArrowDataType::Utf8,
CqlType::Uuid | CqlType::TimeUuid => ArrowDataType::FixedSizeBinary(16),
CqlType::Inet => ArrowDataType::Utf8,
CqlType::List(inner) | CqlType::Set(inner) => {
let item_type = cql_type_to_arrow_data_type(inner);
ArrowDataType::List(Arc::new(Field::new("item", item_type, true)))
}
CqlType::Frozen(inner) => cql_type_to_arrow_data_type(inner),
CqlType::Map(key_type, val_type) => {
let key_arrow = cql_type_to_arrow_data_type(key_type);
let val_arrow = cql_type_to_arrow_data_type(val_type);
ArrowDataType::Map(
Arc::new(Field::new(
"entries",
ArrowDataType::Struct(Fields::from(vec![
Field::new("key", key_arrow, false),
Field::new("value", val_arrow, true),
])),
false,
)),
false,
)
}
CqlType::Tuple(element_types) => {
if element_types.is_empty() {
return ArrowDataType::Utf8;
}
let struct_fields: Vec<Field> = element_types
.iter()
.enumerate()
.map(|(i, t)| {
Field::new(
format!("field_{i}"),
cql_type_to_arrow_data_type(t),
true, )
})
.collect();
ArrowDataType::Struct(Fields::from(struct_fields))
}
CqlType::Udt(_udt_name, udt_fields) => {
if udt_fields.is_empty() {
return ArrowDataType::Utf8;
}
let struct_fields: Vec<Field> = udt_fields
.iter()
.map(|(field_name, field_type)| {
Field::new(
field_name.as_str(),
cql_type_to_arrow_data_type(field_type),
true, )
})
.collect();
ArrowDataType::Struct(Fields::from(struct_fields))
}
CqlType::Custom(_) => ArrowDataType::Utf8,
}
}
pub(crate) fn build_typed_value_array(
cql_type: &CqlType,
values: &[Option<&Value>],
) -> Result<ArrayRef, ArrowConvertError> {
let effective_type = unwrap_frozen_type(cql_type);
match effective_type {
CqlType::Boolean => {
let arr: Vec<Option<bool>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Boolean(b) => Ok(Some(*b)),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Boolean value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<bool>>, ArrowConvertError>>()?;
Ok(Arc::new(BooleanArray::from(arr)))
}
CqlType::TinyInt => {
let arr: Vec<Option<i8>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::TinyInt(i) => Ok(Some(*i)),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected TinyInt value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<i8>>, ArrowConvertError>>()?;
Ok(Arc::new(Int8Array::from(arr)))
}
CqlType::SmallInt => {
let arr: Vec<Option<i16>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::SmallInt(i) => Ok(Some(*i)),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected SmallInt value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<i16>>, ArrowConvertError>>()?;
Ok(Arc::new(Int16Array::from(arr)))
}
CqlType::Int => {
let arr: Vec<Option<i32>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Integer(i) => Ok(Some(*i)),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Int value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<i32>>, ArrowConvertError>>()?;
Ok(Arc::new(Int32Array::from(arr)))
}
CqlType::BigInt => {
let arr: Vec<Option<i64>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::BigInt(i) => Ok(Some(*i)),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected BigInt value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
Ok(Arc::new(Int64Array::from(arr)))
}
CqlType::Counter => {
let arr: Vec<Option<i64>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Counter(c) => Ok(Some(*c)),
Value::BigInt(i) => Ok(Some(*i)),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Counter value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
Ok(Arc::new(Int64Array::from(arr)))
}
CqlType::Float => {
let arr: Vec<Option<f32>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Float32(f) => Ok(Some(*f)),
Value::Float(f) => Ok(Some(*f as f32)),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Float value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<f32>>, ArrowConvertError>>()?;
Ok(Arc::new(Float32Array::from(arr)))
}
CqlType::Double => {
let arr: Vec<Option<f64>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Float(f) => Ok(Some(*f)),
Value::Float32(f) => Ok(Some(*f as f64)),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Double value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<f64>>, ArrowConvertError>>()?;
Ok(Arc::new(Float64Array::from(arr)))
}
CqlType::Text | CqlType::Ascii | CqlType::Varchar => {
let arr: Vec<Option<String>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Text(s) => Ok(Some(s.clone())),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Text value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
Ok(Arc::new(StringArray::from(arr)))
}
CqlType::Blob => {
let byte_slices: Vec<Option<Vec<u8>>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Blob(b) => Ok(Some(b.clone())),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Blob value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<Vec<u8>>>, ArrowConvertError>>()?;
let refs: Vec<Option<&[u8]>> = byte_slices.iter().map(|o| o.as_deref()).collect();
Ok(Arc::new(BinaryArray::from(refs)))
}
CqlType::Timestamp => {
let arr: Vec<Option<i64>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Timestamp(ts) => Ok(Some(*ts)),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Timestamp value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
Ok(Arc::new(
TimestampMillisecondArray::from(arr).with_timezone("UTC"),
))
}
CqlType::Date => {
let arr: Vec<Option<i32>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Date(d) => Ok(Some(*d)),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Date value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<i32>>, ArrowConvertError>>()?;
Ok(Arc::new(Date32Array::from(arr)))
}
CqlType::Time => {
let arr: Vec<Option<i64>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Time(t) => Ok(Some(*t)),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Time value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
Ok(Arc::new(Time64NanosecondArray::from(arr)))
}
CqlType::Decimal => {
let mut builder = arrow::array::Decimal128Builder::new()
.with_precision_and_scale(DECIMAL_MAX_PRECISION, DECIMAL_FIXED_SCALE as i8)?;
for opt in values {
let v = unwrap_frozen_value(*opt);
match v {
Some(Value::Decimal { scale, unscaled }) => {
let rescaled = rescale_decimal(*scale, unscaled)?;
builder.append_value(rescaled);
}
Some(Value::Null) | None => builder.append_null(),
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"expected Decimal value in element, got {:?}",
other
)));
}
}
}
Ok(Arc::new(builder.finish()))
}
CqlType::Varint => {
use num_bigint::BigInt;
let mut builder = arrow::array::Decimal128Builder::new()
.with_precision_and_scale(DECIMAL_MAX_PRECISION, 0)?;
for opt in values {
let v = unwrap_frozen_value(*opt);
match v {
Some(Value::Varint(bytes)) => {
if bytes.is_empty() {
builder.append_value(0);
} else {
let bigint = BigInt::from_signed_bytes_be(bytes);
let max_abs = BigInt::from(10i64).pow(38u32) - BigInt::from(1i64);
let abs_val = if bigint.sign() == num_bigint::Sign::Minus {
-bigint.clone()
} else {
bigint.clone()
};
if abs_val > max_abs {
return Err(ArrowConvertError::InvalidValue(
"varint element exceeds Decimal128(38, 0) range".to_string(),
));
}
let i128_val = bigint_to_i128(&bigint)?;
builder.append_value(i128_val);
}
}
Some(Value::Null) | None => builder.append_null(),
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"expected Varint value in element, got {:?}",
other
)));
}
}
}
Ok(Arc::new(builder.finish()))
}
CqlType::Duration => {
let arr: Vec<Option<String>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Duration { .. } => Ok(Some(ValueFormatter::format_value(v))),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Duration value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
Ok(Arc::new(StringArray::from(arr)))
}
CqlType::Uuid | CqlType::TimeUuid => {
let mut builder = arrow::array::FixedSizeBinaryBuilder::new(16);
for opt in values {
let v = unwrap_frozen_value(*opt);
match v {
Some(Value::Uuid(bytes)) => builder.append_value(bytes)?,
Some(Value::Null) | None => builder.append_null(),
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"expected Uuid value in element, got {:?}",
other
)));
}
}
}
Ok(Arc::new(builder.finish()))
}
CqlType::Inet => {
let arr: Vec<Option<String>> = values
.iter()
.filter_map(|opt| {
let v = unwrap_frozen_value(*opt)?;
Some(match v {
Value::Inet(bytes) => Ok(Some(ValueFormatter::format_value(&Value::Inet(
bytes.clone(),
)))),
Value::Null => Ok(None),
other => Err(ArrowConvertError::InvalidValue(format!(
"expected Inet value in element, got {:?}",
other
))),
})
})
.collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
Ok(Arc::new(StringArray::from(arr)))
}
CqlType::List(inner) | CqlType::Set(inner) => {
let element_type = cql_type_to_arrow_data_type(inner);
let item_field = Arc::new(Field::new("item", element_type, true));
let mut offsets: Vec<i32> = vec![0];
let mut flat_elements: Vec<Option<&Value>> = Vec::new();
let mut null_bitmap: Vec<bool> = Vec::new();
for opt in values {
let v = unwrap_frozen_value(*opt);
match v {
Some(Value::List(items)) | Some(Value::Set(items)) => {
null_bitmap.push(true);
for item in items {
flat_elements.push(Some(item));
}
offsets.push(flat_elements.len() as i32);
}
Some(Value::Null) | None => {
null_bitmap.push(false);
offsets.push(flat_elements.len() as i32);
}
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"expected List/Set value, got {:?}",
other
)));
}
}
}
let elements_array = build_typed_value_array(inner, &flat_elements)?;
let offset_buffer = OffsetBuffer::new(offsets.into());
let null_buffer = NullBuffer::from(null_bitmap);
Ok(Arc::new(ListArray::new(
item_field,
offset_buffer,
elements_array,
Some(null_buffer),
)))
}
CqlType::Frozen(inner) => build_typed_value_array(inner, values),
CqlType::Map(key_type, val_type) => {
let key_arrow = cql_type_to_arrow_data_type(key_type);
let val_arrow = cql_type_to_arrow_data_type(val_type);
let mut offsets: Vec<i32> = vec![0];
let mut flat_keys: Vec<Option<&Value>> = Vec::new();
let mut flat_vals: Vec<Option<&Value>> = Vec::new();
let mut null_bitmap: Vec<bool> = Vec::new();
for opt in values {
let v = unwrap_frozen_value(*opt);
match v {
Some(Value::Map(pairs)) => {
null_bitmap.push(true);
for (k, val) in pairs {
if matches!(k, Value::Null) {
return Err(ArrowConvertError::InvalidValue(
"null key in map is not allowed in Arrow MapArray".to_string(),
));
}
flat_keys.push(Some(k));
flat_vals.push(Some(val));
}
offsets.push(flat_keys.len() as i32);
}
Some(Value::Null) | None => {
null_bitmap.push(false);
offsets.push(flat_keys.len() as i32);
}
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"expected Map value, got {:?}",
other
)));
}
}
}
let key_array = build_typed_value_array(key_type, &flat_keys)?;
let val_array = build_typed_value_array(val_type, &flat_vals)?;
let struct_fields = Fields::from(vec![
Field::new("key", key_arrow, false),
Field::new("value", val_arrow, true),
]);
let entries_array =
StructArray::new(struct_fields.clone(), vec![key_array, val_array], None);
let map_field = Arc::new(Field::new(
"entries",
ArrowDataType::Struct(struct_fields),
false,
));
let offset_buffer = OffsetBuffer::new(offsets.into());
let null_buffer = NullBuffer::from(null_bitmap);
Ok(Arc::new(MapArray::new(
map_field,
offset_buffer,
entries_array,
Some(null_buffer),
false,
)))
}
CqlType::Tuple(element_types) => {
if element_types.is_empty() {
let arr: Vec<Option<String>> = values
.iter()
.map(|opt| match unwrap_frozen_value(*opt) {
Some(Value::Null) | None => Ok(None),
Some(v @ Value::Tuple(_)) => Ok(Some(ValueFormatter::format_value(v))),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"expected Tuple value, got {:?}",
other
))),
})
.collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
return Ok(Arc::new(StringArray::from(arr)));
}
let n_rows = values.len();
let n_fields = element_types.len();
let unwrapped: Vec<Option<&Value>> =
values.iter().map(|opt| unwrap_frozen_value(*opt)).collect();
for v in unwrapped.iter() {
match v {
Some(Value::Tuple(_)) | Some(Value::Null) | None => {}
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"expected Tuple value, got {:?}",
other
)));
}
}
}
let null_bitmap: Vec<bool> = unwrapped
.iter()
.map(|v| !matches!(v, Some(Value::Null) | None))
.collect();
let null_sentinel = Value::Null;
let mut child_arrays: Vec<ArrayRef> = Vec::with_capacity(n_fields);
for (field_idx, element_type) in element_types.iter().enumerate() {
let child_values: Vec<Option<&Value>> = (0..n_rows)
.map(|row_idx| {
match unwrapped[row_idx] {
Some(Value::Tuple(items)) => {
Some(
items
.get(field_idx)
.map(|v| v as &Value)
.unwrap_or(&null_sentinel),
)
}
_ => Some(&null_sentinel),
}
})
.collect();
let child_arr = build_typed_value_array(element_type, &child_values)?;
child_arrays.push(child_arr);
}
let struct_fields: Fields = Fields::from(
element_types
.iter()
.enumerate()
.map(|(i, t)| {
Field::new(format!("field_{i}"), cql_type_to_arrow_data_type(t), true)
})
.collect::<Vec<_>>(),
);
let null_buffer = NullBuffer::from(null_bitmap);
Ok(Arc::new(StructArray::new(
struct_fields,
child_arrays,
Some(null_buffer),
)))
}
CqlType::Udt(_udt_name, udt_fields) => {
if udt_fields.is_empty() {
let arr: Vec<Option<String>> = values
.iter()
.map(|opt| match unwrap_frozen_value(*opt) {
Some(Value::Null) | None => Ok(None),
Some(v @ Value::Udt(_)) => Ok(Some(ValueFormatter::format_value(v))),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"expected Udt value, got {:?}",
other
))),
})
.collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
return Ok(Arc::new(StringArray::from(arr)));
}
let n_rows = values.len();
let unwrapped: Vec<Option<&Value>> =
values.iter().map(|opt| unwrap_frozen_value(*opt)).collect();
for v in unwrapped.iter() {
match v {
Some(Value::Udt(_)) | Some(Value::Null) | None => {}
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"expected Udt value, got {:?}",
other
)));
}
}
}
let null_bitmap: Vec<bool> = unwrapped
.iter()
.map(|v| !matches!(v, Some(Value::Null) | None))
.collect();
let null_sentinel = Value::Null;
let mut child_arrays: Vec<ArrayRef> = Vec::with_capacity(udt_fields.len());
for (field_name, field_type) in udt_fields.iter() {
let child_values: Vec<Option<&Value>> = (0..n_rows)
.map(|row_idx| match unwrapped[row_idx] {
Some(Value::Udt(udt_val)) => {
Some(
udt_val
.fields
.iter()
.find(|f| &f.name == field_name)
.and_then(|f| f.value.as_ref().map(|v| v as &Value))
.unwrap_or(&null_sentinel),
)
}
_ => Some(&null_sentinel),
})
.collect();
let child_arr = build_typed_value_array(field_type, &child_values)?;
child_arrays.push(child_arr);
}
let struct_fields: Fields = Fields::from(
udt_fields
.iter()
.map(|(field_name, field_type)| {
Field::new(
field_name.as_str(),
cql_type_to_arrow_data_type(field_type),
true,
)
})
.collect::<Vec<_>>(),
);
let null_buffer = NullBuffer::from(null_bitmap);
Ok(Arc::new(StructArray::new(
struct_fields,
child_arrays,
Some(null_buffer),
)))
}
CqlType::Custom(_) => {
let arr: Vec<Option<String>> = values
.iter()
.map(|opt| match opt {
Some(Value::Null) | None => None,
Some(v) => Some(ValueFormatter::format_value(v)),
})
.collect();
Ok(Arc::new(StringArray::from(arr)))
}
}
}
pub(crate) fn unwrap_frozen_type(cql_type: &CqlType) -> &CqlType {
let mut t = cql_type;
while let CqlType::Frozen(inner) = t {
t = inner.as_ref();
}
t
}
pub(crate) fn unwrap_frozen_value(v: Option<&Value>) -> Option<&Value> {
match v {
Some(Value::Frozen(inner)) => Some(inner.as_ref()),
other => other,
}
}
pub(crate) fn data_type_to_arrow(data_type: &DataType) -> ArrowDataType {
match data_type {
DataType::Null => ArrowDataType::Null,
DataType::Boolean => ArrowDataType::Boolean,
DataType::TinyInt => ArrowDataType::Int8,
DataType::SmallInt => ArrowDataType::Int16,
DataType::Integer => ArrowDataType::Int32,
DataType::BigInt => ArrowDataType::Int64,
DataType::Float32 => ArrowDataType::Float32,
DataType::Float => ArrowDataType::Float64,
DataType::Text => ArrowDataType::Utf8,
DataType::Blob => ArrowDataType::Binary,
DataType::Timestamp => ArrowDataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
DataType::Uuid => ArrowDataType::FixedSizeBinary(16),
DataType::Json => ArrowDataType::Utf8,
DataType::List => {
ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Utf8, true)))
}
DataType::Set => {
ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Utf8, true)))
}
DataType::Map => ArrowDataType::Map(
Arc::new(Field::new(
"entries",
ArrowDataType::Struct(Fields::from(vec![
Field::new("key", ArrowDataType::Utf8, false),
Field::new("value", ArrowDataType::Utf8, true),
])),
false,
)),
false,
),
DataType::Tuple => ArrowDataType::Utf8, DataType::Udt => ArrowDataType::Utf8, DataType::Frozen => ArrowDataType::Utf8,
DataType::Tombstone => ArrowDataType::Utf8,
}
}
pub(crate) fn convert_to_arrays(
columns: &[ColumnInfo],
rows: &[QueryRow],
) -> Result<Vec<ArrayRef>, ArrowConvertError> {
columns
.iter()
.map(|col| convert_column_to_array(col, rows))
.collect()
}
pub(crate) fn convert_column_to_array(
col: &ColumnInfo,
rows: &[QueryRow],
) -> Result<ArrayRef, ArrowConvertError> {
if let Some(cql_type) = &col.cql_type {
let effective = unwrap_frozen_type(cql_type);
match effective {
CqlType::Date => return build_date32_array(col, rows),
CqlType::Time => return build_time64_ns_array(col, rows),
CqlType::Decimal => return build_decimal128_array(col, rows),
CqlType::Varint => return build_varint_as_decimal128_array(col, rows),
CqlType::Duration => return build_duration_utf8_array(col, rows),
CqlType::Uuid | CqlType::TimeUuid => return build_uuid_fixed_binary_array(col, rows),
CqlType::Inet => return build_inet_utf8_array(col, rows),
CqlType::Counter => return build_int64_array(col, rows),
CqlType::List(_)
| CqlType::Set(_)
| CqlType::Map(_, _)
| CqlType::Tuple(_)
| CqlType::Udt(_, _) => {
let column_values: Vec<Option<&Value>> = rows
.iter()
.map(|row| row.values.get(col.name.as_str()))
.collect();
return build_typed_value_array(cql_type, &column_values);
}
_ => {}
}
}
match &col.data_type {
DataType::Boolean => build_boolean_array(col, rows),
DataType::TinyInt => build_int8_array(col, rows),
DataType::SmallInt => build_int16_array(col, rows),
DataType::Integer => build_int32_array(col, rows),
DataType::BigInt => build_int64_array(col, rows),
DataType::Float32 => build_float32_array(col, rows),
DataType::Float => build_float64_array(col, rows),
DataType::Text | DataType::Json => build_string_array(col, rows),
DataType::Blob => build_binary_array(col, rows),
DataType::Timestamp => build_timestamp_array(col, rows),
DataType::Uuid => build_uuid_array(col, rows),
DataType::List | DataType::Set => build_list_array(col, rows),
DataType::Map => build_map_array(col, rows),
DataType::Tuple
| DataType::Udt
| DataType::Frozen
| DataType::Tombstone
| DataType::Null => {
build_string_array(col, rows) }
}
}
pub(crate) fn rescale_decimal(scale: i32, unscaled: &[u8]) -> Result<i128, ArrowConvertError> {
use num_bigint::BigInt;
if unscaled.is_empty() {
return Ok(0i128);
}
if scale > DECIMAL_FIXED_SCALE {
return Err(ArrowConvertError::InvalidValue(format!(
"decimal scale {scale} exceeds the fixed export scale {DECIMAL_FIXED_SCALE}; \
refusing to truncate (would lose precision)"
)));
}
let bigint = BigInt::from_signed_bytes_be(unscaled);
let delta = DECIMAL_FIXED_SCALE - scale;
let rescaled = if delta == 0 {
bigint
} else {
let factor = BigInt::from(10i64).pow(delta as u32);
bigint * factor
};
let max_abs = BigInt::from(10i64).pow(38u32) - BigInt::from(1i64);
let abs_rescaled = if rescaled.sign() == num_bigint::Sign::Minus {
-rescaled.clone()
} else {
rescaled.clone()
};
if abs_rescaled > max_abs {
return Err(ArrowConvertError::InvalidValue(format!(
"Decimal value exceeds Decimal128(38, {DECIMAL_FIXED_SCALE}) range after rescaling"
)));
}
bigint_to_i128(&rescaled)
}
fn build_boolean_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<bool>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::Boolean(b)) => Ok(Some(*b)),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Boolean value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<bool>>, ArrowConvertError>>()?;
Ok(Arc::new(BooleanArray::from(values)))
}
fn build_int8_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<i8>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::TinyInt(i)) => Ok(Some(*i)),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected TinyInt value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<i8>>, ArrowConvertError>>()?;
Ok(Arc::new(Int8Array::from(values)))
}
fn build_int16_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<i16>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::SmallInt(i)) => Ok(Some(*i)),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected SmallInt value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<i16>>, ArrowConvertError>>()?;
Ok(Arc::new(Int16Array::from(values)))
}
fn build_int32_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let allow_compat = col.cql_type.is_none();
let values: Vec<Option<i32>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::Integer(i)) => Ok(Some(*i)),
Some(Value::Date(d)) if allow_compat => Ok(Some(*d)), Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Int value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<i32>>, ArrowConvertError>>()?;
Ok(Arc::new(Int32Array::from(values)))
}
fn build_int64_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let effective = col.cql_type.as_ref().map(unwrap_frozen_type);
let allow_counter = matches!(effective, None | Some(CqlType::Counter));
let allow_compat = effective.is_none();
let values: Vec<Option<i64>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::BigInt(i)) => Ok(Some(*i)),
Some(Value::Counter(c)) if allow_counter => Ok(Some(*c)),
Some(Value::Time(t)) if allow_compat => Ok(Some(*t)), Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected BigInt value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
Ok(Arc::new(Int64Array::from(values)))
}
fn build_float32_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<f32>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::Float32(f)) => Ok(Some(*f)),
Some(Value::Float(f)) => Ok(Some(*f as f32)),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Float value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<f32>>, ArrowConvertError>>()?;
Ok(Arc::new(Float32Array::from(values)))
}
fn build_float64_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<f64>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::Float(f)) => Ok(Some(*f)),
Some(Value::Float32(f)) => Ok(Some(*f as f64)),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Double value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<f64>>, ArrowConvertError>>()?;
Ok(Arc::new(Float64Array::from(values)))
}
fn build_string_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let strict_text = matches!(
col.cql_type.as_ref().map(unwrap_frozen_type),
Some(CqlType::Text | CqlType::Ascii | CqlType::Varchar)
);
let values: Vec<Option<String>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::Null) => Ok(None),
Some(Value::Text(s)) => Ok(Some(s.clone())),
Some(Value::Json(j)) if !strict_text => Ok(Some(j.to_string())),
Some(other) if strict_text => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Text value, got {:?}",
col.name, other
))),
Some(other) => Ok(Some(ValueFormatter::format_value(other))),
},
)
.collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
Ok(Arc::new(StringArray::from(values)))
}
fn build_binary_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<&[u8]>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::Blob(b)) => Ok(Some(b.as_slice())),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Blob value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<&[u8]>>, ArrowConvertError>>()?;
Ok(Arc::new(BinaryArray::from(values)))
}
fn build_timestamp_array(
col: &ColumnInfo,
rows: &[QueryRow],
) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<i64>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::Timestamp(ts)) => Ok(Some(*ts)),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Timestamp value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
Ok(Arc::new(
TimestampMillisecondArray::from(values).with_timezone("UTC"),
))
}
fn build_uuid_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<[u8; 16]>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::Uuid(uuid)) => Ok(Some(*uuid)),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Uuid value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<[u8; 16]>>, ArrowConvertError>>()?;
let mut builder = arrow::array::FixedSizeBinaryBuilder::new(16);
for opt in values {
match opt {
Some(uuid) => builder.append_value(uuid)?,
None => builder.append_null(),
}
}
Ok(Arc::new(builder.finish()))
}
fn build_date32_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<i32>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::Date(days)) => Ok(Some(*days)),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Date value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<i32>>, ArrowConvertError>>()?;
Ok(Arc::new(Date32Array::from(values)))
}
fn build_time64_ns_array(
col: &ColumnInfo,
rows: &[QueryRow],
) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<i64>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::Time(nanos)) => Ok(Some(*nanos)),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Time value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
Ok(Arc::new(Time64NanosecondArray::from(values)))
}
fn build_decimal128_array(
col: &ColumnInfo,
rows: &[QueryRow],
) -> Result<ArrayRef, ArrowConvertError> {
let mut builder = arrow::array::Decimal128Builder::new()
.with_precision_and_scale(DECIMAL_MAX_PRECISION, DECIMAL_FIXED_SCALE as i8)?;
for row in rows {
match unwrap_frozen_value(row.values.get(col.name.as_str())) {
Some(Value::Decimal { scale, unscaled }) => {
let rescaled = rescale_decimal(*scale, unscaled).map_err(|e| {
ArrowConvertError::InvalidValue(format!("Column '{}': {e}", col.name))
})?;
builder.append_value(rescaled);
}
Some(Value::Null) | None => {
builder.append_null();
}
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"Column '{}': expected Decimal value, got {:?}",
col.name, other
)));
}
}
}
Ok(Arc::new(builder.finish()))
}
fn build_varint_as_decimal128_array(
col: &ColumnInfo,
rows: &[QueryRow],
) -> Result<ArrayRef, ArrowConvertError> {
use num_bigint::BigInt;
let mut builder = arrow::array::Decimal128Builder::new()
.with_precision_and_scale(DECIMAL_MAX_PRECISION, 0)?;
for row in rows {
match unwrap_frozen_value(row.values.get(col.name.as_str())) {
Some(Value::Varint(bytes)) => {
if bytes.is_empty() {
builder.append_value(0);
} else {
let bigint = BigInt::from_signed_bytes_be(bytes);
let max_abs = BigInt::from(10i64).pow(38u32) - BigInt::from(1i64);
let abs_val = if bigint.sign() == num_bigint::Sign::Minus {
-bigint.clone()
} else {
bigint.clone()
};
if abs_val > max_abs {
return Err(ArrowConvertError::InvalidValue(format!(
"Column '{}': varint value exceeds Decimal128(38, 0) range",
col.name
)));
}
let i128_val = bigint_to_i128(&bigint).map_err(|e| {
ArrowConvertError::InvalidValue(format!("Column '{}': {e}", col.name))
})?;
builder.append_value(i128_val);
}
}
Some(Value::Null) | None => {
builder.append_null();
}
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"Column '{}': expected Varint value, got {:?}",
col.name, other
)));
}
}
}
Ok(Arc::new(builder.finish()))
}
fn build_duration_utf8_array(
col: &ColumnInfo,
rows: &[QueryRow],
) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<String>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(v @ Value::Duration { .. }) => Ok(Some(ValueFormatter::format_value(v))),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Duration value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
Ok(Arc::new(StringArray::from(values)))
}
fn build_uuid_fixed_binary_array(
col: &ColumnInfo,
rows: &[QueryRow],
) -> Result<ArrayRef, ArrowConvertError> {
let mut builder = arrow::array::FixedSizeBinaryBuilder::new(16);
for row in rows {
match unwrap_frozen_value(row.values.get(col.name.as_str())) {
Some(Value::Uuid(bytes)) => builder.append_value(bytes)?,
Some(Value::Null) | None => builder.append_null(),
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"Column '{}': expected Uuid value, got {:?}",
col.name, other
)));
}
}
}
Ok(Arc::new(builder.finish()))
}
fn build_inet_utf8_array(
col: &ColumnInfo,
rows: &[QueryRow],
) -> Result<ArrayRef, ArrowConvertError> {
let values: Vec<Option<String>> = rows
.iter()
.map(
|row| match unwrap_frozen_value(row.values.get(col.name.as_str())) {
None => Ok(None),
Some(Value::Inet(bytes)) => Ok(Some(ValueFormatter::format_value(&Value::Inet(
bytes.clone(),
)))),
Some(Value::Null) => Ok(None),
Some(other) => Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Inet value, got {:?}",
col.name, other
))),
},
)
.collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
Ok(Arc::new(StringArray::from(values)))
}
fn build_list_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let mut offsets: Vec<i32> = vec![0];
let mut values: Vec<Option<String>> = Vec::new();
let mut null_bitmap: Vec<bool> = Vec::new();
for row in rows {
match unwrap_frozen_value(row.values.get(col.name.as_str())) {
Some(Value::List(items)) | Some(Value::Set(items)) => {
null_bitmap.push(true);
for item in items {
values.push(Some(ValueFormatter::format_value(item)));
}
offsets.push(values.len() as i32);
}
Some(Value::Null) | None => {
null_bitmap.push(false);
offsets.push(values.len() as i32);
}
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected List/Set value, got {:?}",
col.name, other
)));
}
}
}
let values_array = Arc::new(StringArray::from(values)) as ArrayRef;
let field = Arc::new(Field::new("item", ArrowDataType::Utf8, true));
let offset_buffer = OffsetBuffer::new(offsets.into());
let null_buffer = NullBuffer::from(null_bitmap);
Ok(Arc::new(ListArray::new(
field,
offset_buffer,
values_array,
Some(null_buffer),
)))
}
fn build_map_array(col: &ColumnInfo, rows: &[QueryRow]) -> Result<ArrayRef, ArrowConvertError> {
let mut offsets: Vec<i32> = vec![0];
let mut keys: Vec<Option<String>> = Vec::new();
let mut values: Vec<Option<String>> = Vec::new();
let mut null_bitmap: Vec<bool> = Vec::new();
for row in rows {
match unwrap_frozen_value(row.values.get(col.name.as_str())) {
Some(Value::Map(pairs)) => {
null_bitmap.push(true);
for (k, v) in pairs {
keys.push(Some(ValueFormatter::format_value(k)));
values.push(Some(ValueFormatter::format_value(v)));
}
offsets.push(keys.len() as i32);
}
Some(Value::Null) | None => {
null_bitmap.push(false);
offsets.push(keys.len() as i32);
}
Some(other) => {
return Err(ArrowConvertError::InvalidValue(format!(
"column '{}': expected Map value, got {:?}",
col.name, other
)));
}
}
}
let key_array = Arc::new(StringArray::from(keys)) as ArrayRef;
let value_array = Arc::new(StringArray::from(values)) as ArrayRef;
let struct_fields = Fields::from(vec![
Field::new("key", ArrowDataType::Utf8, false),
Field::new("value", ArrowDataType::Utf8, true),
]);
let entries_array = StructArray::new(struct_fields.clone(), vec![key_array, value_array], None);
let map_field = Arc::new(Field::new(
"entries",
ArrowDataType::Struct(struct_fields),
false,
));
let offset_buffer = OffsetBuffer::new(offsets.into());
let null_buffer = NullBuffer::from(null_bitmap);
Ok(Arc::new(MapArray::new(
map_field,
offset_buffer,
entries_array,
Some(null_buffer),
false,
)))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::query::{ColumnInfo, QueryRow};
use crate::schema::CqlType;
use crate::types::{DataType, Value};
use crate::RowKey;
use arrow::array::{Array, Int32Array};
fn col(name: &str, data_type: DataType, cql_type: Option<CqlType>) -> ColumnInfo {
ColumnInfo {
name: name.to_string(),
data_type,
nullable: true,
position: 0,
table_name: None,
cql_type,
}
}
fn row_one(name: &str, value: Value) -> QueryRow {
let mut values: HashMap<Arc<str>, Value> = HashMap::new();
values.insert(name.into(), value);
QueryRow {
values,
key: RowKey::new(Vec::new()),
metadata: Default::default(),
cell_metadata: None,
}
}
fn row_absent() -> QueryRow {
QueryRow {
values: HashMap::new(),
key: RowKey::new(Vec::new()),
metadata: Default::default(),
cell_metadata: None,
}
}
fn is_invalid_value(res: Result<arrow::record_batch::RecordBatch, ArrowConvertError>) -> bool {
matches!(res, Err(ArrowConvertError::InvalidValue(_)))
}
#[test]
fn typed_scalar_type_mismatch_is_error() {
let columns = vec![col("d", DataType::Timestamp, Some(CqlType::Date))];
let rows = vec![row_one("d", Value::Text("not-a-date".into()))];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn flat_builder_type_mismatch_is_error() {
let columns = vec![col("n", DataType::Integer, None)];
let rows = vec![row_one("n", Value::Text("nope".into()))];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn collection_expected_list_got_scalar_is_error() {
let columns = vec![col(
"l",
DataType::List,
Some(CqlType::List(Box::new(CqlType::Int))),
)];
let rows = vec![row_one("l", Value::Integer(5))];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn collection_mistyped_element_is_error() {
let columns = vec![col(
"l",
DataType::List,
Some(CqlType::List(Box::new(CqlType::Int))),
)];
let rows = vec![row_one(
"l",
Value::List(vec![Value::Integer(1), Value::Text("bad".into())]),
)];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn collection_expected_map_got_scalar_is_error() {
let columns = vec![col(
"m",
DataType::Map,
Some(CqlType::Map(
Box::new(CqlType::Text),
Box::new(CqlType::Int),
)),
)];
let rows = vec![row_one("m", Value::Integer(7))];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn null_and_absent_still_build_ok() {
let columns = vec![col("n", DataType::Integer, None)];
let rows = vec![row_one("n", Value::Null), row_absent()];
let batch = rows_to_record_batch(&columns, &rows).expect("null/absent must build");
assert_eq!(batch.num_rows(), 2);
assert_eq!(batch.column(0).null_count(), 2);
}
#[test]
fn correctly_typed_value_builds_ok() {
let columns = vec![col("n", DataType::Integer, None)];
let rows = vec![row_one("n", Value::Integer(42))];
let batch = rows_to_record_batch(&columns, &rows).expect("well-typed value must build");
let arr = batch
.column(0)
.as_any()
.downcast_ref::<Int32Array>()
.expect("Int32Array");
assert_eq!(arr.value(0), 42);
assert_eq!(arr.null_count(), 0);
}
#[test]
fn decimal_scale_above_fixed_is_error() {
let columns = vec![col("d", DataType::Blob, Some(CqlType::Decimal))];
let unscaled = num_bigint::BigInt::from(123_456_789_012i64).to_signed_bytes_be();
let rows = vec![row_one(
"d",
Value::Decimal {
scale: 12,
unscaled,
},
)];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn decimal_scale_within_fixed_builds_ok() {
use arrow::array::Decimal128Array;
let columns = vec![col("d", DataType::Blob, Some(CqlType::Decimal))];
let unscaled = num_bigint::BigInt::from(123_456i64).to_signed_bytes_be();
let rows = vec![row_one("d", Value::Decimal { scale: 3, unscaled })];
let batch = rows_to_record_batch(&columns, &rows).expect("in-range decimal must build");
let arr = batch
.column(0)
.as_any()
.downcast_ref::<Decimal128Array>()
.expect("Decimal128Array");
assert_eq!(arr.value(0), 123_456_000_000i128);
assert_eq!(arr.null_count(), 0);
}
#[test]
fn decimal_null_and_absent_still_null() {
let columns = vec![col("d", DataType::Blob, Some(CqlType::Decimal))];
let rows = vec![row_one("d", Value::Null), row_absent()];
let batch = rows_to_record_batch(&columns, &rows).expect("null/absent decimal must build");
assert_eq!(batch.num_rows(), 2);
assert_eq!(batch.column(0).null_count(), 2);
}
#[test]
fn float32_column_accepts_wide_float_value() {
let flat = vec![col("h", DataType::Float32, None)];
let rows = vec![row_one("h", Value::Float(1.84f32 as f64))];
let batch = rows_to_record_batch(&flat, &rows).expect("wide float must narrow, not error");
let arr = batch
.column(0)
.as_any()
.downcast_ref::<Float32Array>()
.expect("Float32Array");
assert_eq!(arr.value(0), 1.84f32);
assert_eq!(arr.null_count(), 0);
let typed = vec![col("h", DataType::Float32, Some(CqlType::Float))];
let rows = vec![row_one("h", Value::Float(1.84f32 as f64))];
let batch =
rows_to_record_batch(&typed, &rows).expect("wide float (typed) must narrow, not error");
let arr = batch
.column(0)
.as_any()
.downcast_ref::<Float32Array>()
.expect("Float32Array");
assert_eq!(arr.value(0), 1.84f32);
}
#[test]
fn tuple_expected_tuple_got_scalar_is_error() {
let columns = vec![col(
"t",
DataType::Text,
Some(CqlType::Tuple(vec![CqlType::Int, CqlType::Text])),
)];
let rows = vec![row_one("t", Value::Text("not-a-tuple".into()))];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn tuple_null_and_absent_still_build_ok() {
let columns = vec![col(
"t",
DataType::Text,
Some(CqlType::Tuple(vec![CqlType::Int, CqlType::Text])),
)];
let rows = vec![row_one("t", Value::Null), row_absent()];
let batch = rows_to_record_batch(&columns, &rows).expect("null/absent tuple must build");
assert_eq!(batch.num_rows(), 2);
assert_eq!(batch.column(0).null_count(), 2);
}
#[test]
fn udt_expected_udt_got_scalar_is_error() {
let columns = vec![col(
"u",
DataType::Text,
Some(CqlType::Udt(
"my_type".into(),
vec![("a".into(), CqlType::Int), ("b".into(), CqlType::Text)],
)),
)];
let rows = vec![row_one("u", Value::Integer(9))];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn udt_null_and_absent_still_build_ok() {
let columns = vec![col(
"u",
DataType::Text,
Some(CqlType::Udt(
"my_type".into(),
vec![("a".into(), CqlType::Int), ("b".into(), CqlType::Text)],
)),
)];
let rows = vec![row_one("u", Value::Null), row_absent()];
let batch = rows_to_record_batch(&columns, &rows).expect("null/absent UDT must build");
assert_eq!(batch.num_rows(), 2);
assert_eq!(batch.column(0).null_count(), 2);
}
#[test]
fn empty_field_udt_expected_udt_got_scalar_is_error() {
let columns = vec![col(
"u",
DataType::Text,
Some(CqlType::Udt("unresolved".into(), vec![])),
)];
let rows = vec![row_one("u", Value::Integer(9))];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn empty_field_tuple_expected_tuple_got_scalar_is_error() {
let columns = vec![col("t", DataType::Text, Some(CqlType::Tuple(vec![])))];
let rows = vec![row_one("t", Value::Text("nope".into()))];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn authoritative_text_column_type_mismatch_is_error() {
for cql in [CqlType::Text, CqlType::Ascii, CqlType::Varchar] {
let columns = vec![col("s", DataType::Text, Some(cql))];
let rows = vec![row_one("s", Value::Integer(1))];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
}
#[test]
fn authoritative_text_column_rejects_json() {
let columns = vec![col("s", DataType::Text, Some(CqlType::Text))];
let rows = vec![row_one("s", Value::Json(serde_json::json!({"a": 1})))];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn frozen_wrapped_scalar_values_build_ok() {
let text_cols = vec![col(
"s",
DataType::Text,
Some(CqlType::Frozen(Box::new(CqlType::Text))),
)];
let text_rows = vec![row_one(
"s",
Value::Frozen(Box::new(Value::Text("hi".into()))),
)];
let batch =
rows_to_record_batch(&text_cols, &text_rows).expect("frozen<text> value must build");
let arr = batch
.column(0)
.as_any()
.downcast_ref::<StringArray>()
.expect("StringArray");
assert_eq!(arr.value(0), "hi");
let date_cols = vec![col(
"d",
DataType::Integer,
Some(CqlType::Frozen(Box::new(CqlType::Date))),
)];
let date_rows = vec![row_one("d", Value::Frozen(Box::new(Value::Date(19_000))))];
let batch =
rows_to_record_batch(&date_cols, &date_rows).expect("frozen<date> value must build");
assert_eq!(batch.num_rows(), 1);
assert_eq!(batch.column(0).null_count(), 0);
}
#[test]
fn authoritative_text_column_builds_ok() {
let columns = vec![col("s", DataType::Text, Some(CqlType::Text))];
let rows = vec![
row_one("s", Value::Text("hi".into())),
row_one("s", Value::Null),
row_absent(),
];
let batch = rows_to_record_batch(&columns, &rows).expect("well-typed text must build");
let arr = batch
.column(0)
.as_any()
.downcast_ref::<StringArray>()
.expect("StringArray");
assert_eq!(arr.value(0), "hi");
assert_eq!(arr.null_count(), 2);
}
#[test]
fn authoritative_int_column_rejects_date() {
let columns = vec![col("n", DataType::Integer, Some(CqlType::Int))];
let rows = vec![row_one("n", Value::Date(19_000))];
assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
}
#[test]
fn opaque_int_column_accepts_date() {
let columns = vec![col("n", DataType::Integer, None)];
let rows = vec![row_one("n", Value::Date(19_000))];
let batch = rows_to_record_batch(&columns, &rows).expect("opaque int accepts Date");
let arr = batch
.column(0)
.as_any()
.downcast_ref::<Int32Array>()
.expect("Int32Array");
assert_eq!(arr.value(0), 19_000);
}
#[test]
fn authoritative_bigint_counter_reject_mismatch() {
let bigint_time = vec![col("b", DataType::BigInt, Some(CqlType::BigInt))];
assert!(is_invalid_value(rows_to_record_batch(
&bigint_time,
&[row_one("b", Value::Time(123))]
)));
let counter_time = vec![col("c", DataType::BigInt, Some(CqlType::Counter))];
assert!(is_invalid_value(rows_to_record_batch(
&counter_time,
&[row_one("c", Value::Time(123))]
)));
let bigint_counter = vec![col("b", DataType::BigInt, Some(CqlType::BigInt))];
assert!(is_invalid_value(rows_to_record_batch(
&bigint_counter,
&[row_one("b", Value::Counter(7))]
)));
}
#[test]
fn authoritative_counter_column_accepts_counter() {
let columns = vec![col("c", DataType::BigInt, Some(CqlType::Counter))];
let rows = vec![row_one("c", Value::Counter(42))];
let batch = rows_to_record_batch(&columns, &rows).expect("counter accepts Counter");
assert_eq!(batch.num_rows(), 1);
assert_eq!(batch.column(0).null_count(), 0);
}
}