use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use arrow::array::{
Array, ArrayRef, AsArray, FixedSizeBinaryArray, RecordBatch, RecordBatchOptions,
};
use arrow::datatypes::{FieldRef, Schema};
use arrow::ipc::reader::StreamReader;
use arrow::ipc::writer::StreamWriter;
use uuid::Uuid;
use crate::catalog::INTERNAL_ROW_ID_FIELD_NAME;
use crate::expr::{Expr, visited_columns};
use crate::{ILError, ILResult};
pub fn schema_without_row_id(schema: &Schema) -> Schema {
let fields = schema
.fields
.iter()
.filter(|field| field.name() != INTERNAL_ROW_ID_FIELD_NAME)
.cloned()
.collect::<Vec<_>>();
Schema::new_with_metadata(fields, schema.metadata().clone())
}
pub fn extract_row_id_array_from_record_batch(
record_batch: &RecordBatch,
) -> ILResult<FixedSizeBinaryArray> {
let index = record_batch
.schema_ref()
.index_of(INTERNAL_ROW_ID_FIELD_NAME)?;
let row_id_array = record_batch
.column(index)
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.ok_or_else(|| ILError::internal("Row id column is not an FixedSizeBinaryArray"))?;
if row_id_array.value_length() != 16 {
return Err(ILError::internal(
"Row id column is not an FixedSizeBinaryArray with length 16",
));
}
Ok(row_id_array.clone())
}
pub fn extract_row_ids_from_record_batch(record_batch: &RecordBatch) -> ILResult<Vec<Uuid>> {
let row_id_array = extract_row_id_array_from_record_batch(record_batch)?;
fixed_size_binary_array_to_uuids(&row_id_array)
}
pub fn sort_record_batches(batches: &[RecordBatch], sort_col: &str) -> ILResult<Vec<RecordBatch>> {
if batches.is_empty() {
return Ok(Vec::new());
}
let record = arrow::compute::concat_batches(&batches[0].schema(), batches)?;
let sort_col_idx = record.schema().index_of(sort_col)?;
let sort_array = record.column(sort_col_idx);
let indices = arrow::compute::sort_to_indices(sort_array, None, None)?;
let sorted_arrays: Vec<ArrayRef> = record
.columns()
.iter()
.map(|col| arrow::compute::take(col, &indices, None))
.collect::<arrow::error::Result<_>>()?;
let options = RecordBatchOptions::new().with_row_count(Some(record.num_rows()));
let batch = RecordBatch::try_new_with_options(record.schema(), sorted_arrays, &options)?;
Ok(vec![batch])
}
pub fn array_to_uuids(array: &dyn Array) -> ILResult<Vec<Uuid>> {
let data_type = array.data_type();
let fixed_size_binary_array = array.as_fixed_size_binary_opt().ok_or_else(|| {
ILError::internal(format!("array is not an FixedSizeBinaryArray: {data_type}",))
})?;
fixed_size_binary_array_to_uuids(fixed_size_binary_array)
}
pub fn fixed_size_binary_array_to_uuids(array: &FixedSizeBinaryArray) -> ILResult<Vec<Uuid>> {
if array.value_length() != 16 {
return Err(ILError::internal(
"array is not an FixedSizeBinaryArray with length 16",
));
}
let mut uuids = Vec::with_capacity(array.len());
for i in 0..array.len() {
let uuid = Uuid::from_slice(array.value(i))?;
uuids.push(uuid);
}
Ok(uuids)
}
pub fn build_row_id_array<T, U>(iter: T) -> ILResult<FixedSizeBinaryArray>
where
T: Iterator<Item = U>,
U: AsRef<[u8]>,
{
let array =
FixedSizeBinaryArray::try_from_sparse_iter_with_size(iter.map(|item| Some(item)), 16)?;
Ok(array)
}
pub fn project_schema(schema: &Schema, projection: Option<&Vec<usize>>) -> ILResult<Schema> {
if let Some(projection) = projection {
let schema = schema.project(projection)?;
Ok(schema)
} else {
Ok(schema.clone())
}
}
pub fn build_projection_from_condition(schema: &Schema, condition: &Expr) -> ILResult<Vec<usize>> {
let visited_columns = visited_columns(condition);
if visited_columns.is_empty() {
return Ok(Vec::new());
}
let mut projection = Vec::new();
for col in visited_columns {
let idx = schema.index_of(&col)?;
projection.push(idx);
}
projection.sort();
Ok(projection)
}
pub fn serialize_array(array: ArrayRef, field: FieldRef) -> ILResult<Vec<u8>> {
let schema = Arc::new(Schema::new(vec![field]));
let batch = RecordBatch::try_new(schema.clone(), vec![array.clone()])?;
let mut buf = Vec::with_capacity(array.get_array_memory_size());
let mut writer = StreamWriter::try_new(&mut buf, &schema)?;
writer.write(&batch)?;
writer.finish()?;
Ok(buf)
}
pub fn deserialize_array(buf: &[u8], field: FieldRef) -> ILResult<ArrayRef> {
let schema = Arc::new(Schema::new(vec![field]));
let reader = StreamReader::try_new(buf, None)?;
if reader.schema() != schema {
return Err(ILError::internal(format!(
"Schema mismatch when deserializing array: {:?}, {:?}",
reader.schema(),
schema,
)));
}
let mut arrays = Vec::new();
for batch in reader {
let batch = batch?;
let array = batch.column(0).clone();
arrays.push(array);
}
let array = arrow::compute::concat(
arrays
.iter()
.map(|a| a.as_ref())
.collect::<Vec<_>>()
.as_slice(),
)?;
Ok(array)
}
pub fn rewrite_batch_schema(
batch: &RecordBatch,
field_name_id_map: &HashMap<String, Uuid>,
) -> ILResult<RecordBatch> {
let mut new_fields = Vec::new();
let batch_schema = batch.schema_ref();
for field in batch_schema.fields() {
if let Some(field_id) = field_name_id_map.get(field.name()) {
let new_field_name = hex::encode(field_id);
let new_field = field.as_ref().clone().with_name(new_field_name);
new_fields.push(Arc::new(new_field));
} else if field.name() == INTERNAL_ROW_ID_FIELD_NAME {
new_fields.push(field.clone());
} else {
return Err(ILError::invalid_input(format!("Invalid field {field}")));
}
}
let new_schema =
Arc::new(Schema::new(new_fields).with_metadata(batch_schema.metadata().clone()));
let new_batch = RecordBatch::try_new(new_schema, batch.columns().to_vec())?;
Ok(new_batch)
}
pub fn correct_batch_schema(
batch: &RecordBatch,
field_id_name_map: &HashMap<Uuid, String>,
) -> ILResult<RecordBatch> {
let mut new_fields = Vec::new();
let batch_schema = batch.schema_ref();
for field in batch_schema.fields() {
if field.name() == INTERNAL_ROW_ID_FIELD_NAME {
new_fields.push(field.clone());
} else if let Ok(field_id) = Uuid::parse_str(field.name()) {
let Some(correct_field_name) = field_id_name_map.get(&field_id) else {
return Err(ILError::internal(format!(
"Not found field name for field id {field_id}"
)));
};
let new_field = field.as_ref().clone().with_name(correct_field_name);
new_fields.push(Arc::new(new_field));
} else {
new_fields.push(field.clone());
}
}
let new_schema =
Arc::new(Schema::new(new_fields).with_metadata(batch_schema.metadata().clone()));
let new_batch = RecordBatch::try_new(new_schema, batch.columns().to_vec())?;
Ok(new_batch)
}
pub fn timestamp_ms_from_now(duration: Duration) -> i64 {
let start_time = SystemTime::now();
let end_time = start_time + duration;
end_time
.duration_since(SystemTime::UNIX_EPOCH)
.expect("Time went backwards")
.as_millis() as i64
}
pub fn serialize_row_ids(row_ids: &[Uuid]) -> Vec<u8> {
let mut buf = Vec::with_capacity(row_ids.len() * 16);
for row_id in row_ids {
buf.extend_from_slice(row_id.as_bytes());
}
buf
}
pub fn deserialize_row_ids(bytes: &[u8]) -> ILResult<Vec<Uuid>> {
if !bytes.len().is_multiple_of(16) {
return Err(ILError::internal(format!(
"row_ids bytes length {} is not a multiple of 16",
bytes.len()
)));
}
let mut row_ids = Vec::with_capacity(bytes.len() / 16);
for chunk in bytes.as_chunks::<16>().0 {
row_ids.push(Uuid::from_slice(chunk)?);
}
Ok(row_ids)
}