use std::cmp::{Ordering, Reverse};
use std::collections::{BTreeMap, BTreeSet, BinaryHeap};
use std::sync::Arc;
use arrow_array::RecordBatch;
use arrow_row::{RowConverter, Rows, SortField};
use arrow_schema::{SchemaRef, SortOptions};
use arrow_select::concat::concat_batches;
use futures::{StreamExt as _, TryStreamExt as _, stream};
use super::RecordBatchStream;
use super::api::SortColumn;
use super::error::ScandalError;
use super::plan::{PLAN_METADATA_KEY, PlanDescription};
use super::profile::PROFILE_METADATA_KEY;
struct Cursor {
batch: RecordBatch,
rows: Arc<Rows>,
position: usize,
}
struct Head {
rows: Arc<Rows>,
position: usize,
stream: usize,
}
impl PartialEq for Head {
fn eq(&self, other: &Self) -> bool {
self.cmp(other) == Ordering::Equal
}
}
impl Eq for Head {}
impl PartialOrd for Head {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for Head {
fn cmp(&self, other: &Self) -> Ordering {
self.rows
.row(self.position)
.cmp(&other.rows.row(other.position))
.then_with(|| self.stream.cmp(&other.stream))
}
}
struct Merge {
schema: SchemaRef,
streams: Vec<RecordBatchStream>,
cursors: Vec<Option<Cursor>>,
heap: BinaryHeap<Reverse<Head>>,
converter: RowConverter,
columns: Vec<usize>,
batch_rows: usize,
initialized: bool,
diagnostics: BTreeMap<String, RecordBatch>,
diagnostic_omissions: u64,
diagnostic_bytes: usize,
failure: Option<ScandalError>,
}
impl Merge {
async fn advance(&mut self, index: usize) -> Result<(), ScandalError> {
loop {
let batch = match self.streams[index].try_next().await {
Ok(Some(batch)) => batch,
Ok(None) => break,
Err(error) => {
self.failure = Some(error);
break;
}
};
if batch
.schema_ref()
.metadata()
.contains_key(PROFILE_METADATA_KEY)
&& let Some(mut plan) = PlanDescription::from_schema(batch.schema_ref())?
&& plan.profile.is_some()
{
let id = plan
.execution_id
.clone()
.unwrap_or_else(|| plan.physical_id.clone());
let encoded_size = batch
.schema_ref()
.metadata()
.get(PLAN_METADATA_KEY)
.map_or(0, String::len)
.saturating_add(32);
if let Some(previous) = self.diagnostics.remove(&id) {
self.diagnostic_bytes = self.diagnostic_bytes.saturating_sub(
previous
.schema_ref()
.metadata()
.get(PLAN_METADATA_KEY)
.map_or(0, String::len)
.saturating_add(32),
);
}
while self.diagnostics.len() >= 64
|| self.diagnostic_bytes.saturating_add(encoded_size) > 16 * 1024 * 1024
{
let Some((_, previous)) = self.diagnostics.pop_first() else {
break;
};
self.diagnostic_bytes = self.diagnostic_bytes.saturating_sub(
previous
.schema_ref()
.metadata()
.get(PLAN_METADATA_KEY)
.map_or(0, String::len)
.saturating_add(32),
);
self.diagnostic_omissions = self.diagnostic_omissions.saturating_add(1);
}
if let Some(profile) = &mut plan.profile {
profile.delivery_omissions = self.diagnostic_omissions;
}
let summary =
RecordBatch::new_empty(Arc::new(plan.with_schema(batch.schema_ref())?));
self.diagnostic_bytes = self.diagnostic_bytes.saturating_add(
summary
.schema_ref()
.metadata()
.get(PLAN_METADATA_KEY)
.map_or(0, String::len)
.saturating_add(32),
);
self.diagnostics.insert(id, summary);
}
if batch.num_rows() == 0 {
continue;
}
if PlanDescription::row_schema(&batch.schema())
!= PlanDescription::row_schema(&self.schema)
{
return Err(invalid("ordered stream schema changed"));
}
let columns = self
.columns
.iter()
.map(|index| Arc::clone(batch.column(*index)))
.collect::<Vec<_>>();
let rows = Arc::new(self.converter.convert_columns(&columns).map_err(invalid)?);
self.heap.push(Reverse(Head {
rows: Arc::clone(&rows),
position: 0,
stream: index,
}));
self.cursors[index] = Some(Cursor {
batch,
rows,
position: 0,
});
return Ok(());
}
self.cursors[index] = None;
Ok(())
}
fn take_diagnostic(&mut self) -> Option<RecordBatch> {
let (_, batch) = self.diagnostics.pop_first()?;
self.diagnostic_bytes = self.diagnostic_bytes.saturating_sub(
batch
.schema_ref()
.metadata()
.get(PLAN_METADATA_KEY)
.map_or(0, String::len)
.saturating_add(32),
);
Some(batch)
}
async fn next(&mut self) -> Result<Option<RecordBatch>, ScandalError> {
if !self.initialized {
for index in 0..self.streams.len() {
self.advance(index).await?;
}
self.initialized = true;
}
if let Some(report) = self.take_diagnostic() {
return Ok(Some(report));
}
if let Some(error) = self.failure.take() {
return Err(error);
}
let mut pieces = Vec::new();
let mut count = 0;
while count < self.batch_rows {
let Some(Reverse(head)) = self.heap.pop() else {
break;
};
let cursor = self.cursors[head.stream]
.as_mut()
.expect("heap head has a cursor");
let end = (cursor.position + self.batch_rows - count).min(cursor.batch.num_rows());
let end = match self.heap.peek() {
Some(Reverse(other)) => {
let mut low = cursor.position + 1;
let mut high = end;
while low < high {
let middle = low + (high - low) / 2;
let comparison = cursor
.rows
.row(middle)
.cmp(&other.rows.row(other.position))
.then_with(|| head.stream.cmp(&other.stream));
if comparison.is_le() {
low = middle + 1;
} else {
high = middle;
}
}
low
}
None => end,
};
pieces.push(cursor.batch.slice(cursor.position, end - cursor.position));
count += end - cursor.position;
cursor.position = end;
if end == cursor.batch.num_rows() {
self.advance(head.stream).await?;
} else {
self.heap.push(Reverse(Head {
rows: Arc::clone(&cursor.rows),
position: end,
stream: head.stream,
}));
}
}
if pieces.is_empty() {
return Ok(self.take_diagnostic());
}
let schema = merged_description(&self.schema, &pieces)?;
concat_batches(&schema, &pieces).map(Some).map_err(invalid)
}
}
fn merged_description(
schema: &SchemaRef,
pieces: &[RecordBatch],
) -> Result<SchemaRef, ScandalError> {
let mut encoded = BTreeSet::new();
let mut plans = BTreeMap::new();
let mut executions = BTreeSet::new();
let mut semantic = None;
for piece in pieces {
let Some(value) = piece.schema_ref().metadata().get(PLAN_METADATA_KEY) else {
return Ok(Arc::clone(schema));
};
if !encoded.insert(value) {
continue;
}
let Some(mut plan) = PlanDescription::from_schema(piece.schema_ref())? else {
return Ok(Arc::clone(schema));
};
let Some(execution) = plan.execution_id.take() else {
return Ok(Arc::clone(schema));
};
if semantic
.as_ref()
.is_some_and(|semantic| semantic != &plan.semantic_id)
{
return Err(invalid(
"ordered stream projection changed between executions",
));
}
semantic = Some(plan.semantic_id.clone());
executions.insert(execution);
plans.entry(plan.physical_id.clone()).or_insert(plan);
}
let Some(mut result) = plans.values().next().cloned() else {
return Ok(Arc::clone(schema));
};
if plans.len() > 1 {
result.physical_id = format!(
"variants[{}]",
plans.keys().cloned().collect::<Vec<_>>().join(",")
);
result.physical.clear();
result.decisions.clear();
for (index, plan) in plans.into_values().enumerate() {
for mut node in plan.physical {
node.id = format!("variant.{index}.{}", node.id);
node.inputs = node
.inputs
.into_iter()
.map(|id| format!("variant.{index}.{id}"))
.collect();
result.physical.push(node);
}
result.decisions.extend(plan.decisions);
}
}
result.profile = None;
result.execution_id = Some(if executions.len() == 1 {
executions.into_iter().next().expect("one execution")
} else {
format!(
"merged[{}]",
executions.into_iter().collect::<Vec<_>>().join(",")
)
});
result.with_schema(schema).map(Arc::new)
}
pub fn merge_ordered(
schema: SchemaRef,
streams: Vec<RecordBatchStream>,
ordering: &[SortColumn],
batch_rows: usize,
) -> Result<RecordBatchStream, ScandalError> {
if batch_rows == 0 || ordering.is_empty() {
return Err(invalid(
"ordered merge needs keys and a positive batch size",
));
}
let columns = ordering
.iter()
.map(|key| {
if key.path.len() != 1
|| key
.transform
.as_deref()
.is_some_and(|transform| transform != "identity")
{
return Err(invalid("ordered merge requires identity output columns"));
}
schema.index_of(&key.path[0]).map_err(invalid)
})
.collect::<Result<Vec<_>, _>>()?;
let fields = columns
.iter()
.zip(ordering)
.map(|(index, key)| {
SortField::new_with_options(
schema.field(*index).data_type().clone(),
SortOptions {
descending: key.descending,
nulls_first: key.nulls_first,
},
)
})
.collect();
let merge = Merge {
cursors: (0..streams.len()).map(|_| None).collect(),
streams,
schema,
heap: BinaryHeap::new(),
converter: RowConverter::new(fields).map_err(invalid)?,
columns,
batch_rows,
initialized: false,
diagnostics: BTreeMap::new(),
diagnostic_omissions: 0,
diagnostic_bytes: 0,
failure: None,
};
Ok(stream::try_unfold(merge, |mut merge| async move {
Ok(merge.next().await?.map(|batch| (batch, merge)))
})
.boxed())
}
fn invalid(error: impl ToString) -> ScandalError {
ScandalError::InvalidArgument {
message: error.to_string(),
}
}
#[cfg(test)]
mod tests {
use arrow_array::{Array as _, Int64Array};
use arrow_schema::{DataType, Field, Schema};
use super::super::profile::ExecutionProfile;
use super::*;
#[tokio::test]
async fn ordered_delivery_preserves_terminal_profiles_before_errors() {
let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int64, false)]));
for failed in [false, true] {
let streams = (0..2)
.map(|index| {
let mut plan = PlanDescription {
version: 1,
semantic_id: "semantic".into(),
physical_id: "physical".into(),
execution_id: Some(index.to_string()),
logical: vec![],
physical: vec![],
decisions: vec![],
profile: Some(ExecutionProfile {
status: "running".into(),
..Default::default()
}),
};
let batch = RecordBatch::try_new(
Arc::new(plan.with_schema(&schema).unwrap()),
vec![Arc::new(Int64Array::from(vec![index]))],
)
.unwrap();
plan.profile.as_mut().unwrap().status = if failed && index == 0 {
"failed"
} else {
"completed"
}
.into();
let terminal =
RecordBatch::new_empty(Arc::new(plan.with_schema(&schema).unwrap()));
let mut batches = vec![Ok(batch), Ok(terminal)];
if failed && index == 0 {
batches.push(Err(invalid("failed execution")));
}
stream::iter(batches).boxed()
})
.collect();
let mut merged = merge_ordered(
Arc::clone(&schema),
streams,
&[SortColumn {
path: vec!["key".into()],
transform: None,
descending: false,
nulls_first: false,
}],
2,
)
.unwrap();
let mut observed = BTreeMap::new();
let mut rows = 0;
let mut error = false;
while let Some(next) = merged.next().await {
let Ok(batch) = next else {
error = true;
break;
};
rows += batch.num_rows();
if let Some(plan) = PlanDescription::from_schema(batch.schema_ref()).unwrap()
&& let Some(profile) = plan.profile
{
observed.insert(plan.execution_id.unwrap(), profile.status);
}
}
assert_eq!(error, failed);
assert_eq!(rows, 2);
assert_eq!(observed["0"], if failed { "failed" } else { "completed" });
assert_eq!(observed["1"], "completed");
}
}
#[tokio::test]
async fn merges_overlapping_runs_across_batches_with_nulls_and_ties() {
let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int64, true)]));
for descending in [false, true] {
for nulls_first in [false, true] {
let mut expected: Vec<Option<i64>> =
vec![None, Some(-5), Some(0), Some(0), Some(3), Some(9), None];
expected.sort_by(|left, right| match (left, right) {
(None, None) => Ordering::Equal,
(None, _) => {
if nulls_first {
Ordering::Less
} else {
Ordering::Greater
}
}
(_, None) => {
if nulls_first {
Ordering::Greater
} else {
Ordering::Less
}
}
(Some(left), Some(right)) => {
if descending {
right.cmp(left)
} else {
left.cmp(right)
}
}
});
let streams = (0..3)
.map(|index| {
let batches = expected
.iter()
.skip(index)
.step_by(3)
.map(|value| {
RecordBatch::try_new(
Arc::clone(&schema),
vec![Arc::new(Int64Array::from(vec![*value]))],
)
.map_err(invalid)
})
.collect::<Vec<_>>();
stream::iter(batches).boxed()
})
.collect();
let merged = merge_ordered(
Arc::clone(&schema),
streams,
&[SortColumn {
path: vec!["key".into()],
transform: None,
descending,
nulls_first,
}],
2,
)
.unwrap();
let batches = merged.try_collect::<Vec<_>>().await.unwrap();
assert!(batches.iter().all(|batch| batch.num_rows() <= 2));
let actual = batches
.iter()
.flat_map(|batch| {
let array = batch
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
(0..array.len())
.map(move |index| (!array.is_null(index)).then(|| array.value(index)))
})
.collect::<Vec<_>>();
assert_eq!(actual, expected);
}
}
}
}