reifydb_engine/bulk_insert/

validation.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

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(())
}