use super::*;
use crate::expression::evaluate_projection_in_graph;
use crate::pipeline::AccountedBindingBatch;
use hawdb_core::{HawDBError, Value};
pub struct UnwindRows<'a> {
pub expression: &'a ProjectionExpression,
pub variable: &'a str,
pub consumes_source: bool,
}
pub fn stream_unwind_rows_batches(
unwind: UnwindRows<'_>,
input: &PhysicalPlan,
source: &mut dyn BindingBatchSource,
context: BatchExecutionContext<'_>,
execution_limit: ExecutionLimit,
emit: &mut dyn FnMut(BindingBatch) -> Result<BatchControl>,
) -> Result<BatchControl> {
runtime_checkpoint(context.task_context)?;
if execution_limit.is_reached(0) {
return Ok(BatchControl::Stop);
}
let working = context.source_account("UnwindRowsExec expressions");
let graph = context.expression_graph();
let mut output = AccountedBindingBatch::with_ledger(
"UnwindRowsExec",
context.memory.batch_rows.get(),
context.memory.batch_payload_bytes,
context.memory_ledger,
);
let mut produced = 0usize;
let UnwindRows {
expression,
variable,
consumes_source,
} = unwind;
let consumed = match expression {
ProjectionExpression::Column(column) if consumes_source => Some(column.as_str()),
_ => None,
};
let control = source.execute(input, ExecutionLimit::unlimited(), &mut |batch| {
for mut binding in batch {
runtime_checkpoint(context.task_context)?;
if execution_limit.is_reached(produced) {
return Ok(BatchControl::Stop);
}
if binding.values.contains_key(variable) {
return Err(HawDBError::Semantic(format!(
"UNWIND variable '{variable}' is already bound"
)));
}
let taken;
let evaluated;
let value = match consumed.and_then(|column| binding.values.remove(column)) {
Some(list) => {
taken = list;
&taken
}
None => {
evaluated = evaluate_projection_in_graph(
expression,
context.catalog,
&binding,
Some(&working),
graph.as_ref(),
)?;
evaluated.as_ref()
}
};
let values = match value {
Value::List(values) => values,
Value::Null => continue,
_ => {
return Err(HawDBError::Execution(
"UNWIND source must resolve to a list value or NULL".into(),
))
}
};
for value in values {
runtime_checkpoint(context.task_context)?;
if execution_limit.is_reached(produced) {
return Ok(BatchControl::Stop);
}
let bytes = binding_memory_bytes_with_values(
&binding,
binding
.values
.iter()
.map(|(name, value)| (name.as_str(), value))
.chain(std::iter::once((variable, value))),
);
if output.push_generated(
bytes,
|| {
let mut row = binding.clone();
row.values.insert(variable.to_owned(), value.clone());
row
},
emit,
)? == BatchControl::Stop
{
return Ok(BatchControl::Stop);
}
produced = produced.saturating_add(1);
if output.is_full() && output.emit(emit)? == BatchControl::Stop {
return Ok(BatchControl::Stop);
}
}
}
Ok(BatchControl::Continue)
})?;
if output.emit(emit)? == BatchControl::Stop {
return Ok(BatchControl::Stop);
}
Ok(control)
}