spiral-sdk 0.1.7

//! The top-k emit: each split's winners with every output column, merged by their keys.

use std::sync::Arc;

use arrow_array::{ArrayRef, RecordBatch, UInt32Array};
use arrow_row::{RowConverter, SortField};
use arrow_schema::{SchemaRef, SortOptions};
use arrow_select::concat::concat_batches;
use arrow_select::take::take_record_batch;

use super::super::error::ScandalError;
use super::{MAX_TOP_K, TopK, column, internal, invalid, require_orderable};

/// A top-k checked against the scan's output rows.
#[derive(Clone, Debug)]
pub struct BoundTopK {
    rows: SchemaRef,
    /// Each key's output column and its order, in key order.
    keys: Vec<(usize, SortOptions)>,
    k: usize,
}

impl BoundTopK {
    /// Check a top-k against the schema of its candidates.
    pub fn bind(top_k: &TopK, rows: &SchemaRef) -> Result<Self, ScandalError> {
        if !(1..=MAX_TOP_K).contains(&top_k.k) {
            return Err(invalid(format!(
                "a top-k keeps between 1 and {MAX_TOP_K} rows; order more rows in your own engine"
            )));
        }
        if top_k.keys.is_empty() {
            return Err(invalid("a top-k orders by at least one key"));
        }
        let mut keys = Vec::with_capacity(top_k.keys.len());
        for key in &top_k.keys {
            let (index, field) = column(rows, &key.column, "top-k key")?;
            require_orderable(field, "top-k key")?;
            if keys.iter().any(|(existing, _)| *existing == index) {
                return Err(invalid(format!("top-k key `{}` repeats", key.column)));
            }
            let descending = key.descending.unwrap_or(false);
            keys.push((
                index,
                SortOptions {
                    descending,
                    nulls_first: key.nulls_first.unwrap_or(descending),
                },
            ));
        }
        Ok(Self {
            rows: Arc::clone(rows),
            keys,
            k: top_k.k as usize,
        })
    }

    /// The output rows, which are also the schema of a split's winners and of their merge.
    pub fn schema(&self) -> &SchemaRef {
        &self.rows
    }

    /// The most rows a split returns, and the merge keeps.
    pub fn k(&self) -> usize {
        self.k
    }

    /// The key columns' names, in key order.
    pub fn key_columns(&self) -> impl Iterator<Item = &str> {
        self.keys
            .iter()
            .map(|(index, _)| self.rows.field(*index).name().as_str())
    }

    /// The first key's output column and its order.
    pub fn first_key(&self) -> Option<(&str, SortOptions)> {
        self.keys
            .first()
            .map(|(index, options)| (self.rows.field(*index).name().as_str(), *options))
    }

    /// Whether the output column `name` is a key.
    pub fn is_key(&self, name: &str) -> bool {
        self.key_columns().any(|key| key == name)
    }

    /// A converter over the key columns, in key order, whose rows compare as the top-k orders.
    pub fn converter(&self) -> Result<RowConverter, ScandalError> {
        RowConverter::new(
            self.keys
                .iter()
                .map(|(index, options)| {
                    SortField::new_with_options(
                        self.rows.field(*index).data_type().clone(),
                        *options,
                    )
                })
                .collect(),
        )
        .map_err(internal)
    }

    /// Merge the splits' winners, in split order: the `k` that sort first, ties in split order.
    pub fn merge(&self, results: &[RecordBatch]) -> Result<RecordBatch, ScandalError> {
        let batch = concat_batches(&self.rows, results).map_err(internal)?;
        let keys = self
            .keys
            .iter()
            .map(|(index, _)| Arc::clone(batch.column(*index)))
            .collect::<Vec<ArrayRef>>();
        let rows = self.converter()?.convert_columns(&keys).map_err(internal)?;
        let mut order = (0..batch.num_rows()).collect::<Vec<_>>();
        // A stable sort keeps each split's own order among ties, and splits in order.
        order.sort_by(|left, right| rows.row(*left).cmp(&rows.row(*right)));
        order.truncate(self.k);
        let indices = order
            .into_iter()
            .map(|index| u32::try_from(index).map_err(internal))
            .collect::<Result<UInt32Array, _>>()?;
        take_record_batch(&batch, &indices).map_err(internal)
    }
}

#[cfg(test)]
mod tests {
    use arrow_array::cast::AsArray as _;
    use arrow_array::{Int64Array, StringArray};
    use arrow_schema::{DataType, Field, Schema};

    use super::super::super::emit::{BoundEmit, Emit, SortKey};
    use super::*;

    fn batch(schema: &SchemaRef, scores: Vec<Option<i64>>, labels: Vec<&str>) -> RecordBatch {
        RecordBatch::try_new(
            Arc::clone(schema),
            vec![
                Arc::new(Int64Array::from(scores)),
                Arc::new(StringArray::from(labels)),
            ],
        )
        .unwrap()
    }

    /// Winners merge by key, descending with nulls first, and ties keep split order.
    #[test]
    fn merge_orders_by_key_and_breaks_ties_in_split_order() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("score", DataType::Int64, true),
            Field::new("label", DataType::Utf8, false),
        ]));
        let emit = Emit::TopK(TopK {
            keys: vec![SortKey {
                column: "score".into(),
                descending: Some(true),
                nulls_first: None,
            }],
            k: 4,
        });
        let BoundEmit::TopK(top_k) = emit.bind_results(&schema).unwrap() else {
            unreachable!()
        };
        let merged = top_k
            .merge(&[
                batch(
                    &schema,
                    vec![Some(9), Some(5), Some(1)],
                    vec!["a", "b", "c"],
                ),
                batch(&schema, vec![None, Some(5), Some(2)], vec!["d", "e", "f"]),
            ])
            .unwrap();
        let labels = merged.column(1).as_string::<i32>();
        assert_eq!(
            labels.iter().flatten().collect::<Vec<_>>(),
            ["d", "a", "b", "e"]
        );
        assert!(merged.column(0).is_null(0));
    }
}