use arrow::array::Array;
use arrow::buffer::BooleanBuffer;
use arrow::record_batch::RecordBatch;
use crate::{ExclusiveModeAst, MutuallyExclusivePlan};
pub(super) fn violations(plan: &MutuallyExclusivePlan, batch: &RecordBatch) -> BooleanBuffer {
let rows = batch.num_rows();
let mut once = BooleanBuffer::new_unset(rows);
let mut twice = BooleanBuffer::new_unset(rows);
for column_index in plan.columns() {
let column = batch.column(*column_index);
let present = column.nulls().map_or_else(
|| BooleanBuffer::new_set(rows),
|nulls| nulls.inner().clone(),
);
twice = &twice | &(&once & &present);
once = &once | &present;
}
match plan.mode() {
ExclusiveModeAst::AtMostOne => twice,
ExclusiveModeAst::ExactlyOne => &twice | &!&once,
}
}
pub(super) fn filled_columns(
plan: &MutuallyExclusivePlan,
batch: &RecordBatch,
row: usize,
) -> Vec<String> {
plan.columns()
.iter()
.zip(plan.names())
.filter(|(column_index, _)| !batch.column(**column_index).is_null(row))
.map(|(_, name)| name.to_string())
.collect()
}