use std::{iter, sync::Arc};
use arrow_array::{ArrayRef, RecordBatch};
use arrow_schema::{FieldRef, Schema};
use reifydb_core::{
interface::catalog::{column::Column, sumtype::SumType},
internal_error,
value::{batch::batch_with, column::builder::ColumnBuilder},
};
use reifydb_value::{
error::Error,
fragment::Fragment,
params::Params,
value::{
Value,
column_view::{ColumnView, ViewData},
system_columns::{column_view, is_system_field},
},
};
use super::coerce::RowCoercer;
use crate::{Result, error::EngineError, vm::instruction::dml::series_insert::VariantTags};
const TAG: &str = "tag";
struct Segment<'a> {
columns: &'a [Column],
tagged: bool,
source_name: &'a str,
coercer: &'a RowCoercer,
}
pub type CoercedTags = (Vec<Option<u8>>, Option<(usize, Error)>);
pub(super) fn coerce_columns(
rows: &[Params],
batches: &[(usize, RecordBatch)],
columns: &[Column],
sumtype: Option<&SumType>,
source_name: &str,
coercer: &RowCoercer,
) -> Result<(Vec<(FieldRef, ArrayRef)>, Option<CoercedTags>)> {
let segment = Segment {
columns,
tagged: sumtype.is_some(),
source_name,
coercer,
};
let total = rows.len() + batches.iter().map(|(_, batch)| batch.num_rows()).sum::<usize>();
let mut builders: Vec<ColumnBuilder> =
columns.iter().map(|col| ColumnBuilder::with_capacity(col.constraint.get_type(), total)).collect();
let mut tags = sumtype.map(|sumtype| VariantTags::new(sumtype, total));
let mut taken = 0;
let mut first_row = 0;
for (at, batch) in batches {
first_row +=
collect_rows_to_columns(&segment, &rows[taken..*at], first_row, &mut builders, tags.as_mut())?;
taken = *at;
first_row += coerce_batch(&segment, batch, first_row, &mut builders, tags.as_mut())?;
}
collect_rows_to_columns(&segment, &rows[taken..], first_row, &mut builders, tags.as_mut())?;
let columns = builders.into_iter().zip(columns).map(|(builder, column)| builder.finish(&column.name)).collect();
Ok((columns, tags.map(VariantTags::finish)))
}
fn coerce_batch(
segment: &Segment<'_>,
batch: &RecordBatch,
first_row: usize,
builders: &mut [ColumnBuilder],
tags: Option<&mut VariantTags<'_>>,
) -> Result<usize> {
let rows = batch.num_rows();
if rows == 0 {
return Ok(0);
}
let views = segment.columns.iter().map(|col| column_view(batch, &col.name)).collect::<Result<Vec<_>>>()?;
let mut failure: Option<(usize, Error)> = None;
for ((builder, column), view) in builders.iter_mut().zip(segment.columns).zip(&views) {
let limit = failure.as_ref().map_or(rows, |(row, _)| *row);
let Some(view) = view else {
for _ in 0..rows {
builder.push_none();
}
continue;
};
let per_cell =
matches!(view.data, ViewData::Any { .. } | ViewData::Digest { .. } | ViewData::None { .. });
if per_cell {
if let Some(found) = coerce_cells(segment, column, view, first_row, limit, builder) {
failure = Some(found);
}
continue;
}
if view.base_type() == *column.constraint.get_type().inner_type() {
builder.append_values(view)?;
continue;
}
match segment.coercer.cast_column(view, column) {
Ok(cast) => builder.append_values(&ColumnView::try_from(&cast)?)?,
Err(_) => match coerce_cells(segment, column, view, first_row, limit, builder) {
Some(found) => failure = Some(found),
None if failure.is_some() => {}
None => {
return Err(internal_error!(
"bulk column {} failed its column cast where every cell cast passes",
column.name
));
}
},
}
}
if let Some((_, error)) = failure {
return Err(error);
}
if let Some(field) = batch.schema().fields().iter().find(|field| {
is_system_field(field)
|| !(segment.columns.iter().any(|c| &c.name == field.name())
|| (segment.tagged && field.name() == TAG))
}) {
return Err(EngineError::BulkInsertColumnNotFound {
fragment: Fragment::None,
table_name: segment.source_name.to_string(),
column: field.name().to_string(),
}
.into());
}
if let Some(tags) = tags {
match column_view(batch, TAG)? {
Some(view) => {
for row in 0..rows {
tags.push(view.get_value(row));
}
}
None => {
for _ in 0..rows {
tags.push(Value::none());
}
}
}
}
Ok(rows)
}
fn coerce_cells(
segment: &Segment<'_>,
column: &Column,
view: &ColumnView<'_>,
first_row: usize,
limit: usize,
builder: &mut ColumnBuilder,
) -> Option<(usize, Error)> {
for row in 0..limit {
match segment.coercer.coerce(view.get_value(row), column, segment.source_name, first_row + row) {
Ok(value) => builder.push_value(value),
Err(error) => return Some((row, error)),
}
}
None
}
fn collect_rows_to_columns(
segment: &Segment<'_>,
rows: &[Params],
first_row: usize,
builders: &mut [ColumnBuilder],
tags: Option<&mut VariantTags<'_>>,
) -> Result<usize> {
let columns = segment.columns;
let source_name = segment.source_name;
let num_cols = columns.len();
let too_many =
rows.iter().position(|params| matches!(params, Params::Positional(vals) if vals.len() > num_cols));
let accepted = &rows[..too_many.unwrap_or(rows.len())];
match params_batch(segment, accepted)? {
Some(batch) => {
coerce_batch(segment, &batch, first_row, builders, None)?;
}
None => coerce_rows(segment, accepted, first_row, builders)?,
}
if let Some(Params::Positional(vals)) = too_many.map(|row| &rows[row]) {
return Err(EngineError::BulkInsertTooManyValues {
fragment: Fragment::None,
expected: num_cols,
actual: vals.len(),
}
.into());
}
for params in rows {
if let Params::Named(map) = params {
for name in map.keys() {
if !(columns.iter().any(|c| &c.name == name) || (segment.tagged && name == TAG)) {
return Err(EngineError::BulkInsertColumnNotFound {
fragment: Fragment::None,
table_name: source_name.to_string(),
column: name.to_string(),
}
.into());
}
}
}
}
if let Some(tags) = tags {
for params in rows {
tags.push(match params {
Params::Named(map) => map.get(TAG).cloned().unwrap_or(Value::none()),
_ => Value::none(),
});
}
}
Ok(rows.len())
}
fn params_batch(segment: &Segment<'_>, rows: &[Params]) -> Result<Option<RecordBatch>> {
let cell = |params: &Params, index: usize, name: &str| match params {
Params::Named(map) => map.get(name).cloned(),
Params::Positional(vals) => vals.get(index).cloned(),
Params::None => None,
};
let mut fields = Vec::with_capacity(segment.columns.len());
for (index, column) in segment.columns.iter().enumerate() {
match params_column(&column.name, rows.iter().map(|params| cell(params, index, &column.name)).collect())
{
ParamsColumn::Typed(field) => fields.push(field),
ParamsColumn::Absent => {}
ParamsColumn::Mixed => return Ok(None),
}
}
batch_with(&Arc::new(Schema::empty()), fields, rows.len()).map(Some)
}
enum ParamsColumn {
Typed((FieldRef, ArrayRef)),
Absent,
Mixed,
}
fn params_column(name: &str, values: Vec<Option<Value>>) -> ParamsColumn {
let present = |value: &&Option<Value>| matches!(value, Some(value) if !matches!(value, Value::None { .. }));
let mut types = values.iter().filter(present).flatten().map(Value::get_type);
let Some(first) = types.next() else {
return ParamsColumn::Absent;
};
if !types.all(|other| other == first) {
return ParamsColumn::Mixed;
}
let mut builder = ColumnBuilder::with_capacity(first, values.len());
for value in values {
match value {
Some(value) if !matches!(value, Value::None { .. }) => builder.push_value(value),
_ => builder.push_none(),
}
}
ParamsColumn::Typed(builder.finish(name))
}
fn coerce_rows(segment: &Segment<'_>, rows: &[Params], first_row: usize, builders: &mut [ColumnBuilder]) -> Result<()> {
let columns = segment.columns;
let source_name = segment.source_name;
let coercer = segment.coercer;
for (row_idx, params) in rows.iter().enumerate() {
match params {
Params::Named(map) => {
for (col_idx, col) in columns.iter().enumerate() {
let value = map.get(&col.name).cloned().unwrap_or(Value::none());
builders[col_idx].push_value(coercer.coerce(
value,
col,
source_name,
first_row + row_idx,
)?);
}
}
Params::Positional(vals) => {
for ((col_data, col), val) in builders
.iter_mut()
.zip(columns.iter())
.zip(vals.iter().map(Some).chain(iter::repeat(None)))
{
let value = val.cloned().unwrap_or(Value::none());
col_data.push_value(coercer.coerce(
value,
col,
source_name,
first_row + row_idx,
)?);
}
}
Params::None => {
for col_data in builders.iter_mut() {
col_data.push_none();
}
}
}
}
Ok(())
}