use anyhow::{Context, Result, bail, ensure};
use arrow::{
array::*,
datatypes::{DataType, Field, Schema},
record_batch::RecordBatch,
};
use rmpv::Value;
use std::sync::Arc;
pub fn cell<'a>(row: &'a Value, name: &str) -> &'a Value {
row.as_map()
.and_then(|m| m.iter().rev().find(|(k, _)| k.as_str() == Some(name)))
.map_or(&Value::Nil, |(_, v)| v)
}
fn array(values: &[&Value]) -> Result<ArrayRef> {
if values.iter().all(|v| v.as_bool().is_some()) {
return Ok(Arc::new(BooleanArray::from(
values
.iter()
.map(|v| v.as_bool().unwrap())
.collect::<Vec<_>>(),
)));
}
ensure!(
values.iter().all(|v| v.is_number()),
"projected value must be numeric or boolean"
);
if values
.iter()
.any(|v| matches!(v, Value::F32(_) | Value::F64(_)))
{
Ok(Arc::new(Float64Array::from(
values
.iter()
.map(|v| v.as_f64().unwrap())
.collect::<Vec<_>>(),
)))
} else if values
.iter()
.any(|v| v.as_u64().is_some_and(|n| n > i64::MAX as u64))
{
if values.iter().any(|v| v.as_i64().is_some_and(|n| n < 0)) {
return Ok(Arc::new(Float64Array::from(
values
.iter()
.map(|v| v.as_f64().unwrap())
.collect::<Vec<_>>(),
)));
}
Ok(Arc::new(UInt64Array::from(
values.iter().map(|v| v.as_u64()).collect::<Vec<_>>(),
)))
} else {
Ok(Arc::new(Int64Array::from(
values.iter().map(|v| v.as_i64()).collect::<Vec<_>>(),
)))
}
}
fn sequence(a: &ArrayRef, row: usize) -> Result<(&ArrayRef, usize, usize)> {
if let Some(a) = a.as_any().downcast_ref::<ListArray>() {
return Ok((
a.values(),
a.value_offsets()[row] as usize,
a.value_offsets()[row + 1] as usize,
));
}
if let Some(a) = a.as_any().downcast_ref::<LargeListArray>() {
return Ok((
a.values(),
a.value_offsets()[row] as usize,
a.value_offsets()[row + 1] as usize,
));
}
if let Some(a) = a.as_any().downcast_ref::<FixedSizeListArray>() {
let start = a.value_offset(row) as usize;
return Ok((a.values(), start, start + a.value_length() as usize));
}
bail!("length projection requires list column")
}
fn dtype(field: &crate::config::Field) -> DataType {
match field.dtype.as_str() {
"i32" => DataType::Int32,
"i64" => DataType::Int64,
"f32" => DataType::Float32,
"f64" => DataType::Float64,
_ => DataType::UInt8,
}
}
fn convert(values: &ArrayRef, field: &crate::config::Field) -> Result<ArrayRef> {
if field.dtype == "bool" {
let booleans = arrow::compute::cast(values, &DataType::Boolean)?;
return Ok(arrow::compute::cast(&booleans, &DataType::UInt8)?);
}
Ok(arrow::compute::cast(values, &dtype(field))?)
}
fn padding(field: &crate::config::Field) -> Result<ArrayRef> {
let pad = field.pad.clone().unwrap_or_else(|| {
if field.dtype == "bool" {
false.into()
} else {
0.into()
}
});
let pad = match pad {
serde_json::Value::Bool(b) => Value::Boolean(b),
serde_json::Value::Number(n) => n
.as_i64()
.map(Value::from)
.unwrap_or_else(|| Value::F64(n.as_f64().unwrap())),
_ => bail!("invalid pad"),
};
convert(&array(&[&pad])?, field)
}
fn dense(
field: &crate::config::Field,
values: ArrayRef,
indices: &[Option<usize>],
) -> Result<ArrayRef> {
let pad = padding(field)?;
let data = values.to_data();
let pad = pad.to_data();
let mut builder = MutableArrayData::new(vec![&data, &pad], false, indices.len());
let mut i = 0;
while i < indices.len() {
if let Some(start) = indices[i] {
let mut end = i + 1;
while end < indices.len() && indices[end] == Some(start + end - i) {
end += 1;
}
builder.extend(0, start, start + end - i);
i = end;
} else {
builder.extend(1, 0, 1);
i += 1;
}
}
Ok(make_array(builder.freeze()))
}
fn batch(fields: &[crate::config::Field], arrays: Vec<ArrayRef>) -> Result<RecordBatch> {
let mut columns = Vec::new();
let mut schema = Vec::new();
for (field, values) in fields.iter().zip(arrays) {
ensure!(
values.null_count() == 0,
"field {}: null projected value or failed conversion",
field.name
);
let mut result = values;
for width in field.dimensions().into_iter().rev() {
result = Arc::new(FixedSizeListArray::try_new(
Arc::new(Field::new("item", result.data_type().clone(), false)),
width as i32,
result,
None,
)?);
}
schema.push(Field::new(&field.name, result.data_type().clone(), false));
columns.push(result);
}
Ok(RecordBatch::try_new(
Arc::new(Schema::new(schema)),
columns,
)?)
}
pub fn project_values(rows: &[&Value], fields: &[crate::config::Field]) -> Result<RecordBatch> {
let _projection = crate::profile::Timer::new(5);
let arrays = fields
.iter()
.map(|field| {
projected_values(rows, field)
.with_context(|| format!("field {} from {}", field.name, field.from))
})
.collect::<Result<Vec<_>>>()?;
batch(fields, arrays)
}
fn projected_values(rows: &[&Value], field: &crate::config::Field) -> Result<ArrayRef> {
let mut selected: Vec<_> = rows
.iter()
.map(|row| Some(cell(row, &field.from)))
.collect();
for width in field.dimensions() {
let mut next = Vec::with_capacity(selected.len() * width);
for value in selected {
let list = value
.map(|v| {
v.as_array()
.ok_or_else(|| anyhow::anyhow!("shape projection requires list"))
})
.transpose()?;
next.extend((0..width).map(|i| list.and_then(|v| v.get(i))));
}
selected = next;
}
let values: Vec<_> = selected.iter().flatten().copied().collect();
let mut index = 0;
let indices: Vec<_> = selected
.iter()
.map(|v| {
v.map(|_| {
let result = index;
index += 1;
result
})
})
.collect();
let values = {
let _arrow = crate::profile::Timer::new(4);
convert(&array(&values)?, field)?
};
dense(field, values, &indices)
}
pub fn project(input: &RecordBatch, fields: &[crate::config::Field]) -> Result<RecordBatch> {
let mut columns = Vec::new();
for field in fields {
let mut a = input
.column_by_name(&field.from)
.ok_or_else(|| anyhow::anyhow!("missing field {}", field.from))?;
if let &[width] = field.dimensions().as_slice()
&& let Some(list) = a.as_fixed_size_list_opt()
&& list.value_length() as usize == width
&& list.null_count() == 0
&& !list.is_empty()
&& (list.value_type().is_numeric() || list.value_type() == DataType::Boolean)
{
let start = list.value_offset(0) as usize;
columns.push(convert(
&list.values().slice(start, list.len() * width),
field,
)?);
continue;
}
let mut indices: Vec<_> = (0..input.num_rows()).map(Some).collect();
for width in field.dimensions() {
let mut next = Vec::with_capacity(indices.len() * width);
let mut child = a;
for row in indices {
if let Some(row) = row {
ensure!(!a.is_null(row), "field {}: null projected list", field.name);
let (values, start, end) =
sequence(a, row).with_context(|| format!("field {}", field.name))?;
child = values;
next.extend((0..width).map(|i| (start + i < end).then_some(start + i)));
} else {
next.extend(std::iter::repeat_n(None, width));
}
}
a = child;
indices = next;
}
let values = if indices.iter().all(Option::is_none) {
new_empty_array(&dtype(field))
} else {
ensure!(
a.data_type().is_numeric() || *a.data_type() == DataType::Boolean,
"field {} has incompatible item type",
field.name
);
convert(a, field)?
};
columns.push(dense(field, values, &indices)?);
}
batch(fields, columns)
}
pub struct Builder<'a> {
fields: &'a [crate::config::Field],
buffers: Vec<arrow::buffer::MutableBuffer>,
}
impl<'a> Builder<'a> {
pub fn new(fields: &'a [crate::config::Field], rows: usize) -> Self {
let buffers = fields
.iter()
.map(|f| {
let size = dtype(f).primitive_width().unwrap();
arrow::buffer::MutableBuffer::with_capacity(rows * f.width() * size)
})
.collect();
Self { fields, buffers }
}
pub fn push(&mut self, chunk: &RecordBatch) {
for (buffer, column) in self.buffers.iter_mut().zip(chunk.columns()) {
let mut leaf = column.clone();
while let Some(list) = leaf.as_fixed_size_list_opt() {
let start = if list.is_empty() {
0
} else {
list.value_offset(0) as usize
};
leaf = list
.values()
.slice(start, list.len() * list.value_length() as usize);
}
let data = leaf.to_data();
let size = data.data_type().primitive_width().unwrap();
buffer.extend_from_slice(
&data.buffers()[0].as_slice()[data.offset() * size..][..data.len() * size],
);
}
}
pub fn finish(self) -> Result<RecordBatch> {
let arrays = self
.fields
.iter()
.zip(self.buffers)
.map(|(field, buffer)| {
let dtype = dtype(field);
let len = buffer.len() / dtype.primitive_width().unwrap();
Ok(make_array(
ArrayData::builder(dtype)
.len(len)
.add_buffer(buffer.into())
.build()?,
))
})
.collect::<Result<Vec<_>>>()?;
batch(self.fields, arrays)
}
}