use std::cmp::{Ordering, Reverse};
use std::collections::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;
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,
}
impl Merge {
async fn advance(&mut self, index: usize) -> Result<(), ScandalError> {
while let Some(batch) = self.streams[index].try_next().await? {
if batch.num_rows() == 0 {
continue;
}
if batch.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(())
}
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;
}
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(None);
}
concat_batches(&self.schema, &pieces)
.map(Some)
.map_err(invalid)
}
}
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,
};
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::*;
#[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);
}
}
}
}