use std::collections::BTreeMap;
use linkrs_core::error::QueryError;
use linkrs_core::types::expr::Expression;
use linkrs_core::value::NullType;
use linkrs_core::Value;
use crate::executor::base::{MemoryBudget, MemoryTracker};
use crate::executor::expression::evaluator::ExpressionEvaluator;
use crate::executor::streaming::chunk::DataChunk;
use crate::executor::streaming::executor::{SortDirection, StreamingExecutor, ValueRowContext};
use crate::executor::streaming::spill::{
finalize_partitions_with_runtime, HashPartitionConfig, HashPartitionSpiller, SpillManager,
};
use super::helpers::{emit_batch_slice, BlockingContext};
use super::window::{
compute_window_partition_result, sort_partition_rows, WindowFunctionState, WindowState,
};
pub(super) fn open_window_function(state: &mut Option<WindowFunctionState>) {
*state = Some(WindowFunctionState {
batch: linkrs_core::columnar::MaterializedBatch::new(0, 0),
col_names: vec![],
result_batch: linkrs_core::columnar::MaterializedBatch::new(0, 0),
emitted_offset: 0,
partition_spiller: None,
spilled_runs: vec![],
current_partition: 0,
has_spilled: false,
output_complete: false,
});
}
pub(super) fn open_window(state: &mut Option<WindowState>) {
*state = Some(WindowState {
batch: linkrs_core::columnar::MaterializedBatch::new(0, 0),
col_names: vec![],
result_batch: linkrs_core::columnar::MaterializedBatch::new(0, 0),
emitted_offset: 0,
partition_spiller: None,
spilled_runs: vec![],
current_partition: 0,
has_spilled: false,
output_complete: false,
});
}
#[allow(clippy::too_many_arguments)]
pub(super) fn next_window_function(
window_exprs: &[Expression],
partition_by_exprs: &[Expression],
order_by_exprs: &[Expression],
order_by_directions: &[SortDirection],
memory_tracker: &mut MemoryTracker,
state: &mut WindowFunctionState,
ctx: &BlockingContext<'_>,
input: &mut StreamingExecutor,
) -> Result<Option<DataChunk>, QueryError> {
use crate::executor::streaming::spill::RunReader;
let eval_partition_key = |row: &[Value], col_names: &[String]| -> Vec<Value> {
if partition_by_exprs.is_empty() {
return vec![Value::Null(NullType::Null)];
}
let mut key = Vec::with_capacity(partition_by_exprs.len());
for expr in partition_by_exprs.iter() {
let mut ctx = ValueRowContext::from_names(row.to_vec(), col_names.to_vec());
key.push(
ExpressionEvaluator::evaluate(expr, &mut ctx)
.unwrap_or(Value::Null(NullType::Null)),
);
}
key
};
loop {
if state.output_complete {
return Ok(None);
}
if state.emitted_offset < state.result_batch.num_rows() {
return Ok(Some(emit_batch_slice(
&state.result_batch,
&mut state.emitted_offset,
ctx,
)));
}
if state.emitted_offset > 0 {
state.emitted_offset = 0;
state.result_batch.clear();
if !state.has_spilled || state.current_partition >= state.spilled_runs.len() {
state.output_complete = true;
return Ok(None);
}
}
if state.has_spilled && state.partition_spiller.is_none() {
while state.current_partition < state.spilled_runs.len() {
if let Some(rt) = ctx.runtime.as_ref() {
rt.ensure_not_cancelled()?;
}
let run = match &state.spilled_runs[state.current_partition] {
Some(r) => r,
None => {
state.current_partition += 1;
continue;
}
};
let mut reader = RunReader::open(run)?;
let mut partitions: BTreeMap<Vec<Value>, Vec<(usize, Vec<Value>)>> =
BTreeMap::new();
let mut run_row_count: u64 = 0;
while let Some(row) = reader.read_row()? {
memory_tracker.try_reserve_row(&row)?;
let partition_key = eval_partition_key(&row, &state.col_names);
partitions
.entry(partition_key)
.or_default()
.push((run_row_count as usize, row));
run_row_count += 1;
}
let mut result_rows = Vec::with_capacity(run_row_count as usize);
for (_key, mut partition_rows) in partitions {
sort_partition_rows(
&mut partition_rows,
&state.col_names,
order_by_exprs,
order_by_directions,
);
result_rows.extend(compute_window_partition_result(
&partition_rows,
&state.col_names,
window_exprs,
));
}
let _ = std::fs::remove_file(&run.path);
state.current_partition += 1;
memory_tracker.reset();
if !result_rows.is_empty() {
state.result_batch.clear();
state.emitted_offset = 0;
for result_row in result_rows {
state.result_batch.append_row(result_row);
}
return Ok(Some(emit_batch_slice(
&state.result_batch,
&mut state.emitted_offset,
ctx,
)));
}
}
state.output_complete = true;
return Ok(None);
}
let mut accumulating = true;
while accumulating {
match input.advance()? {
Some(mut chunk) => {
chunk.normalize_for_opaque("WindowFunction");
if let Some(rt) = ctx.runtime.as_ref() {
rt.ensure_not_cancelled()?;
}
if state.col_names.is_empty() {
state.col_names = chunk.col_names();
}
for row in chunk.rows {
if let Some(ref mut spiller) = state.partition_spiller {
let manager = ctx
.runtime
.as_ref()
.and_then(|rt| rt.get_spill_manager())
.ok_or_else(|| {
QueryError::execution("Spill manager not available".to_string())
})?;
let partition_key = eval_partition_key(&row, &state.col_names);
let p = crate::executor::streaming::spill::hash_row_partition(
&partition_key,
spiller.num_partitions(),
) as usize;
spiller.insert_row_to_partition(&row, p, &manager)?;
continue;
}
if let Err(e) = memory_tracker.try_reserve_row(&row) {
if let Some(sm) =
ctx.runtime.as_ref().and_then(|rt| rt.get_spill_manager())
{
let config = HashPartitionConfig::default();
let num_partitions = config.num_partitions;
let mut spiller = HashPartitionSpiller::new(config, &sm, 0)?;
for pending in state.batch.into_rows() {
let partition_key =
eval_partition_key(&pending, &state.col_names);
let p = crate::executor::streaming::spill::hash_row_partition(
&partition_key,
num_partitions,
) as usize;
spiller.insert_row_to_partition(&pending, p, &sm)?;
memory_tracker
.release(MemoryBudget::estimate_row_memory(&pending));
}
let partition_key = eval_partition_key(&row, &state.col_names);
let p = crate::executor::streaming::spill::hash_row_partition(
&partition_key,
num_partitions,
) as usize;
spiller.insert_row_to_partition(&row, p, &sm)?;
state.partition_spiller = Some(spiller);
state.has_spilled = true;
continue;
} else {
return Err(e);
}
}
state.batch.append_row(row);
}
}
None => {
accumulating = false;
}
}
}
if state.partition_spiller.is_some() {
let runs = finalize_partitions_with_runtime(
state.partition_spiller.take().unwrap(),
ctx.runtime.as_ref(),
)?;
state.spilled_runs = runs;
state.current_partition = 0;
continue;
}
if state.batch.is_empty() {
state.output_complete = true;
return Ok(None);
}
let mut partitions: BTreeMap<Vec<Value>, Vec<(usize, Vec<Value>)>> = BTreeMap::new();
for (idx, row) in state.batch.into_rows().into_iter().enumerate() {
let partition_key = eval_partition_key(&row, &state.col_names);
partitions
.entry(partition_key)
.or_default()
.push((idx, row));
}
let mut result_rows = Vec::new();
for (_key, mut partition_rows) in partitions {
sort_partition_rows(
&mut partition_rows,
&state.col_names,
order_by_exprs,
order_by_directions,
);
result_rows.extend(compute_window_partition_result(
&partition_rows,
&state.col_names,
window_exprs,
));
}
state.result_batch.clear();
state.emitted_offset = 0;
for result_row in result_rows {
state.result_batch.append_row(result_row);
}
if state.result_batch.is_empty() {
state.output_complete = true;
return Ok(None);
}
return Ok(Some(emit_batch_slice(
&state.result_batch,
&mut state.emitted_offset,
ctx,
)));
}
}
#[allow(clippy::too_many_arguments)]
pub(super) fn next_window(
window_exprs: &[Expression],
partition_by_exprs: &[Expression],
order_by_exprs: &[Expression],
order_by_directions: &[SortDirection],
memory_tracker: &mut MemoryTracker,
state: &mut WindowState,
ctx: &BlockingContext<'_>,
input: &mut StreamingExecutor,
) -> Result<Option<DataChunk>, QueryError> {
use crate::executor::streaming::spill::RunReader;
let eval_partition_key = |row: &[Value], col_names: &[String]| -> Vec<Value> {
if partition_by_exprs.is_empty() {
return vec![Value::Null(NullType::Null)];
}
let mut key = Vec::with_capacity(partition_by_exprs.len());
for expr in partition_by_exprs.iter() {
let mut ctx = ValueRowContext::from_names(row.to_vec(), col_names.to_vec());
key.push(
ExpressionEvaluator::evaluate(expr, &mut ctx)
.unwrap_or(Value::Null(NullType::Null)),
);
}
key
};
loop {
if state.output_complete {
return Ok(None);
}
if state.emitted_offset < state.result_batch.num_rows() {
return Ok(Some(emit_batch_slice(
&state.result_batch,
&mut state.emitted_offset,
ctx,
)));
}
if state.emitted_offset > 0 {
state.emitted_offset = 0;
state.result_batch.clear();
if !state.has_spilled || state.current_partition >= state.spilled_runs.len() {
state.output_complete = true;
return Ok(None);
}
}
if state.has_spilled && state.partition_spiller.is_none() {
while state.current_partition < state.spilled_runs.len() {
if let Some(rt) = ctx.runtime.as_ref() {
rt.ensure_not_cancelled()?;
}
let run = match &state.spilled_runs[state.current_partition] {
Some(r) => r,
None => {
state.current_partition += 1;
continue;
}
};
let mut reader = RunReader::open(run)?;
let mut partitions: BTreeMap<Vec<Value>, Vec<(usize, Vec<Value>)>> =
BTreeMap::new();
let mut run_row_count: u64 = 0;
while let Some(row) = reader.read_row()? {
memory_tracker.try_reserve_row(&row)?;
let partition_key = eval_partition_key(&row, &state.col_names);
partitions
.entry(partition_key)
.or_default()
.push((run_row_count as usize, row));
run_row_count += 1;
}
let mut result_rows = Vec::with_capacity(run_row_count as usize);
for (_key, mut partition_rows) in partitions {
sort_partition_rows(
&mut partition_rows,
&state.col_names,
order_by_exprs,
order_by_directions,
);
result_rows.extend(compute_window_partition_result(
&partition_rows,
&state.col_names,
window_exprs,
));
}
let _ = std::fs::remove_file(&run.path);
state.current_partition += 1;
memory_tracker.reset();
if !result_rows.is_empty() {
state.result_batch.clear();
state.emitted_offset = 0;
for result_row in result_rows {
state.result_batch.append_row(result_row);
}
return Ok(Some(emit_batch_slice(
&state.result_batch,
&mut state.emitted_offset,
ctx,
)));
}
}
state.output_complete = true;
return Ok(None);
}
let mut accumulating = true;
while accumulating {
match input.advance()? {
Some(mut chunk) => {
chunk.normalize_for_opaque("Window");
if let Some(rt) = ctx.runtime.as_ref() {
rt.ensure_not_cancelled()?;
}
if state.col_names.is_empty() {
state.col_names = chunk.col_names();
}
for row in chunk.rows {
if let Some(ref mut spiller) = state.partition_spiller {
let manager = ctx
.runtime
.as_ref()
.and_then(|rt| rt.get_spill_manager())
.ok_or_else(|| {
QueryError::execution("Spill manager not available".to_string())
})?;
let partition_key = eval_partition_key(&row, &state.col_names);
let p = crate::executor::streaming::spill::hash_row_partition(
&partition_key,
spiller.num_partitions(),
) as usize;
spiller.insert_row_to_partition(&row, p, &manager)?;
continue;
}
if let Err(e) = memory_tracker.try_reserve_row(&row) {
if let Some(sm) =
ctx.runtime.as_ref().and_then(|rt| rt.get_spill_manager())
{
let config = HashPartitionConfig::default();
let num_partitions = config.num_partitions;
let mut spiller = HashPartitionSpiller::new(config, &sm, 0)?;
for pending in state.batch.into_rows() {
let partition_key =
eval_partition_key(&pending, &state.col_names);
let p = crate::executor::streaming::spill::hash_row_partition(
&partition_key,
num_partitions,
) as usize;
spiller.insert_row_to_partition(&pending, p, &sm)?;
memory_tracker
.release(MemoryBudget::estimate_row_memory(&pending));
}
let partition_key = eval_partition_key(&row, &state.col_names);
let p = crate::executor::streaming::spill::hash_row_partition(
&partition_key,
num_partitions,
) as usize;
spiller.insert_row_to_partition(&row, p, &sm)?;
state.partition_spiller = Some(spiller);
state.has_spilled = true;
continue;
} else {
return Err(e);
}
}
state.batch.append_row(row);
}
}
None => {
accumulating = false;
}
}
}
if state.partition_spiller.is_some() {
let runs = finalize_partitions_with_runtime(
state.partition_spiller.take().unwrap(),
ctx.runtime.as_ref(),
)?;
state.spilled_runs = runs;
state.current_partition = 0;
continue;
}
if state.batch.is_empty() {
state.output_complete = true;
return Ok(None);
}
let mut partitions: BTreeMap<Vec<Value>, Vec<(usize, Vec<Value>)>> = BTreeMap::new();
for (idx, row) in state.batch.into_rows().into_iter().enumerate() {
let partition_key = eval_partition_key(&row, &state.col_names);
partitions
.entry(partition_key)
.or_default()
.push((idx, row));
}
let mut result_rows = Vec::new();
for (_key, mut partition_rows) in partitions {
sort_partition_rows(
&mut partition_rows,
&state.col_names,
order_by_exprs,
order_by_directions,
);
result_rows.extend(compute_window_partition_result(
&partition_rows,
&state.col_names,
window_exprs,
));
}
state.result_batch.clear();
state.emitted_offset = 0;
for result_row in result_rows {
state.result_batch.append_row(result_row);
}
if state.result_batch.is_empty() {
state.output_complete = true;
return Ok(None);
}
return Ok(Some(emit_batch_slice(
&state.result_batch,
&mut state.emitted_offset,
ctx,
)));
}
}
pub(super) fn close_window_function(state: &mut Option<WindowFunctionState>) {
if let Some(ref s) = state {
for r in s.spilled_runs.iter().flatten() {
let _ = std::fs::remove_file(&r.path);
}
}
*state = None;
}
pub(super) fn close_window(state: &mut Option<WindowState>) {
if let Some(ref s) = state {
for r in s.spilled_runs.iter().flatten() {
let _ = std::fs::remove_file(&r.path);
}
}
*state = None;
}
pub(super) fn spill_window_function(
state: &mut WindowFunctionState,
memory_tracker: &mut MemoryTracker,
partition_by_exprs: &[Expression],
sm: &SpillManager,
) -> Result<(), QueryError> {
if state.partition_spiller.is_none() && !state.batch.is_empty() {
let config = HashPartitionConfig::default();
let num_partitions = config.num_partitions;
let mut spiller = HashPartitionSpiller::new(config, sm, 0)?;
for row in state.batch.into_rows() {
let mut partition_key = Vec::new();
if partition_by_exprs.is_empty() {
partition_key.push(Value::Null(NullType::Null));
} else {
for expr in partition_by_exprs.iter() {
let mut ctx = ValueRowContext::from_names(row.clone(), state.col_names.clone());
partition_key.push(
ExpressionEvaluator::evaluate(expr, &mut ctx)
.unwrap_or(Value::Null(NullType::Null)),
);
}
}
let p = crate::executor::streaming::spill::hash_row_partition(
&partition_key,
num_partitions,
) as usize;
spiller.insert_row_to_partition(&row, p, sm)?;
memory_tracker.release(MemoryBudget::estimate_row_memory(&row));
}
state.partition_spiller = Some(spiller);
state.has_spilled = true;
}
Ok(())
}
pub(super) fn spill_window(
state: &mut WindowState,
memory_tracker: &mut MemoryTracker,
partition_by_exprs: &[Expression],
sm: &SpillManager,
) -> Result<(), QueryError> {
if state.partition_spiller.is_none() && !state.batch.is_empty() {
let config = HashPartitionConfig::default();
let num_partitions = config.num_partitions;
let mut spiller = HashPartitionSpiller::new(config, sm, 0)?;
for row in state.batch.into_rows() {
let mut partition_key = Vec::new();
if partition_by_exprs.is_empty() {
partition_key.push(Value::Null(NullType::Null));
} else {
for expr in partition_by_exprs.iter() {
let mut ctx = ValueRowContext::from_names(row.clone(), state.col_names.clone());
partition_key.push(
ExpressionEvaluator::evaluate(expr, &mut ctx)
.unwrap_or(Value::Null(NullType::Null)),
);
}
}
let p = crate::executor::streaming::spill::hash_row_partition(
&partition_key,
num_partitions,
) as usize;
spiller.insert_row_to_partition(&row, p, sm)?;
memory_tracker.release(MemoryBudget::estimate_row_memory(&row));
}
state.partition_spiller = Some(spiller);
state.has_spilled = true;
}
Ok(())
}