use std::sync::Arc;
use datafusion::arrow::array::{ArrayRef, BooleanArray, Int64Array, StringArray};
use datafusion::arrow::datatypes::SchemaRef;
use datafusion::arrow::error::ArrowError;
use datafusion::arrow::record_batch::{RecordBatch, RecordBatchOptions};
use crate::LixError;
pub(super) enum Col<R: 'static> {
Utf8(for<'a> fn(&'a R) -> Option<&'a str>),
Utf8Owned(fn(&R) -> Option<String>),
Utf8Fallible(fn(&R) -> Result<Option<String>, LixError>),
Bool(fn(&R) -> Option<bool>),
I64(fn(&R) -> Option<i64>),
}
pub(super) enum ColumnTableError {
UnsupportedColumn(String),
Arrow(ArrowError),
ArrowZeroColumn(ArrowError),
Row(LixError),
}
pub(super) struct ColumnTable<R: 'static> {
pub(super) columns: &'static [(&'static str, Col<R>)],
}
impl<R> ColumnTable<R> {
pub(super) fn col(&self, name: &str) -> Option<&Col<R>> {
self.columns
.iter()
.find(|(column_name, _)| *column_name == name)
.map(|(_, col)| col)
}
pub(super) fn build(
&self,
schema: SchemaRef,
rows: &[R],
) -> Result<RecordBatch, ColumnTableError> {
if schema.fields().is_empty() {
let options = RecordBatchOptions::new().with_row_count(Some(rows.len()));
return RecordBatch::try_new_with_options(schema, vec![], &options)
.map_err(ColumnTableError::ArrowZeroColumn);
}
let arrays = schema
.fields()
.iter()
.map(|field| {
let name = field.name().as_str();
let col = self
.col(name)
.ok_or_else(|| ColumnTableError::UnsupportedColumn(name.to_string()))?;
build_array(col, rows)
})
.collect::<Result<Vec<_>, ColumnTableError>>()?;
RecordBatch::try_new(schema, arrays).map_err(ColumnTableError::Arrow)
}
}
#[expect(trivial_casts)]
pub(super) fn build_array<R>(col: &Col<R>, rows: &[R]) -> Result<ArrayRef, ColumnTableError> {
Ok(match col {
Col::Utf8(get) => string_array(rows.iter().map(get)),
Col::Utf8Owned(get) => Arc::new(StringArray::from_iter(rows.iter().map(get))) as ArrayRef,
Col::Utf8Fallible(get) => Arc::new(StringArray::from(
rows.iter()
.map(get)
.collect::<Result<Vec<_>, LixError>>()
.map_err(ColumnTableError::Row)?,
)) as ArrayRef,
Col::Bool(get) => Arc::new(BooleanArray::from_iter(rows.iter().map(get))) as ArrayRef,
Col::I64(get) => Arc::new(Int64Array::from_iter(rows.iter().map(get))) as ArrayRef,
})
}
#[expect(trivial_casts)]
pub(super) fn string_array<'a>(values: impl Iterator<Item = Option<&'a str>>) -> ArrayRef {
Arc::new(StringArray::from_iter(values)) as ArrayRef
}