timelake-format 0.7.0

//! Preparing rows for a table, shared by the engine and the WASM writer: aligning batches to the
//! table schema, stamping the insertion time, sorting by `(partition, series…, time, insertion
//! time)`, splitting by partition, and keeping the latest write of each `(series…, time)` — every
//! file is then sorted and free of duplicates on its own, which dedup on read and streaming
//! compaction rely on.

use std::sync::Arc;

use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, TimestampMicrosecondArray, UInt32Array, new_null_array};
use arrow_cast::cast;
use arrow_ord::partition::partition;
use arrow_ord::sort::{SortColumn, lexsort_to_indices};
use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions, TimeUnit};
use arrow_select::concat::concat_batches;
use arrow_select::take::take_record_batch;
use arrow_select::zip::zip;

use crate::Contract;
use crate::error::{WriteError, WriteResult};
use crate::merge::{Merge, Step};

/// Ascending with nulls first: the order Timelake writes and declares for `(series…, time)`.
pub const ASCENDING: SortOptions = SortOptions { descending: false, nulls_first: true };

/// The time zone timestamps with any zone are stored in.
pub const UTC: &str = "+00:00";

/// Rows of one time partition (`None` for unpartitioned tables), sorted by the table's write order,
/// in one or more batches.
pub type PartitionRows = (Option<i32>, Vec<RecordBatch>);

/// Casts and reorders a batch to the table schema by column name. Missing nullable columns are
/// filled with nulls; extra columns are an error.
pub fn align_batch(batch: &RecordBatch, schema: &SchemaRef) -> WriteResult<RecordBatch> {
    let input = batch.schema();
    if let Some(extra) = input.fields().iter().find(|f| schema.field_with_name(f.name()).is_err()) {
        return Err(WriteError::invalid(format!("column {:?} is not in the table schema", extra.name())));
    }
    let columns = schema
        .fields()
        .iter()
        .map(|f| match input.index_of(f.name()) {
            Ok(i) if batch.column(i).data_type() == f.data_type() => Ok(batch.column(i).clone()),
            Ok(i) => Ok(cast(batch.column(i), f.data_type())?),
            Err(_) if f.is_nullable() => Ok(new_null_array(f.data_type(), batch.num_rows())),
            Err(_) => Err(WriteError::invalid(format!("required column {:?} is missing", f.name()))),
        })
        .collect::<WriteResult<Vec<ArrayRef>>>()?;
    RecordBatch::try_new(schema.clone(), columns)
        .map_err(|e| WriteError::invalid(format!("batch does not fit the table schema: {e}")))
}

/// `batch` with its insertion-time `column` set to `now` (µs since the epoch) where it is missing or
/// null; values a caller supplied (an import keeping arrival times) stay.
pub fn with_ingest_time(batch: &RecordBatch, column: &str, now: i64) -> WriteResult<RecordBatch> {
    let n = batch.num_rows();
    let fill = || Arc::new(TimestampMicrosecondArray::from_value(now, n).with_timezone(UTC)) as ArrayRef;
    let ty = DataType::Timestamp(TimeUnit::Microsecond, Some(UTC.into()));
    let schema = batch.schema();
    let Ok(i) = schema.index_of(column) else {
        let mut fields: Vec<_> = schema.fields().iter().cloned().collect();
        fields.push(Arc::new(Field::new(column, ty, false)));
        let mut columns = batch.columns().to_vec();
        columns.push(fill());
        return Ok(RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?);
    };
    if batch.column(i).null_count() == 0 {
        return Ok(batch.clone());
    }
    let given = cast(batch.column(i), &ty)?;
    let filled = zip(&arrow_arith::boolean::is_not_null(&given)?, &given, &fill())?;
    let mut fields: Vec<_> = schema.fields().iter().cloned().collect();
    fields[i] = Arc::new(Field::new(column, filled.data_type().clone(), false));
    let mut columns = batch.columns().to_vec();
    columns[i] = filled;
    Ok(RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?)
}

/// Input is sorted in groups of at most about this many bytes: a group is concatenated and sorted
/// at once, and an append of several groups merges its sorted groups. An Arrow string or binary
/// array holds at most 2 GiB, so no single sort may hold more.
const SORT_GROUP_BYTES: usize = 1 << 30;

/// The partition of each row, as a key column while sorted groups are merged.
const PARTITION: &str = "_timelake_partition";

/// Aligns `batches` to `schema`, stamps the insertion time `now`, sorts by `(partition, series…,
/// time, insertion time)` with input order as the tiebreak, splits by partition and, when dedup
/// is on, keeps the last of rows with the same `(series…, time)`.
pub fn sort_and_split(
    schema: &SchemaRef,
    contract: &Contract,
    batches: &[RecordBatch],
    now: i64,
) -> WriteResult<Vec<PartitionRows>> {
    sort_and_split_in_groups(schema, contract, batches, now, SORT_GROUP_BYTES)
}

fn sort_and_split_in_groups(
    schema: &SchemaRef,
    contract: &Contract,
    batches: &[RecordBatch],
    now: i64,
    group_bytes: usize,
) -> WriteResult<Vec<PartitionRows>> {
    let ingest = contract.ingest_time_column.as_deref();
    let aligned = batches
        .iter()
        .filter(|b| b.num_rows() > 0)
        .map(|b| align_batch(&ingest.map_or_else(|| Ok(b.clone()), |c| with_ingest_time(b, c, now))?, schema))
        .collect::<WriteResult<Vec<_>>>()?;
    let groups = groups_of(aligned, group_bytes);
    if let [group] = groups.as_slice() {
        return sort_group(schema, contract, group);
    }
    merge_groups(schema, contract, &groups)
}

/// `batches` in consecutive groups of at most `bytes` each (a larger batch is a group of its own).
fn groups_of(batches: Vec<RecordBatch>, bytes: usize) -> Vec<Vec<RecordBatch>> {
    let mut groups: Vec<Vec<RecordBatch>> = vec![];
    let mut size = 0;
    for batch in batches {
        let n = batch.get_array_memory_size();
        match groups.last_mut() {
            Some(group) if size + n <= bytes => {
                size += n;
                group.push(batch);
            }
            _ => {
                size = n;
                groups.push(vec![batch]);
            }
        }
    }
    groups
}

/// One group's rows, sorted, split by partition and deduplicated.
fn sort_group(schema: &SchemaRef, contract: &Contract, group: &[RecordBatch]) -> WriteResult<Vec<PartitionRows>> {
    let all = concat_batches(schema, group)?;
    let partitions = partition_values(&all, contract)?;
    let mut keys: Vec<ArrayRef> =
        partitions.iter().map(|p| Arc::new(Int32Array::from(p.clone())) as ArrayRef).collect();
    keys.extend(columns(&all, &contract.write_order())?);
    let (sorted, runs) = if keys.is_empty() {
        let rows = all.num_rows();
        (all, vec![(None, rows)])
    } else {
        let order = stable_sort_indices(&keys, all.num_rows())?;
        let sorted = take_record_batch(&all, &order)?;
        let runs = match partitions {
            None => vec![(None, sorted.num_rows())],
            // Rows are now grouped by partition: one run per partition.
            Some(p) => {
                let sorted_partitions: Vec<i32> = order.values().iter().map(|&i| p[i as usize]).collect();
                sorted_partitions.chunk_by(|a, b| a == b).map(|run| (Some(run[0]), run.len())).collect()
            }
        };
        (sorted, runs)
    };
    let dedup = contract.dedup && !contract.sort_columns().is_empty();
    let mut out = vec![];
    let mut offset = 0;
    for (partition, len) in runs {
        let rows = sorted.slice(offset, len);
        let rows = if dedup { last_per_group(&rows, &contract.sort_columns())? } else { rows };
        out.push((partition, vec![rows]));
        offset += len;
    }
    Ok(out)
}

/// Sorts each group with its partition as a leading key column, then merges the groups in that
/// order (with dedup: the last write of each key kept), split by partition into batches of
/// `row_group_rows`.
fn merge_groups(
    schema: &SchemaRef,
    contract: &Contract,
    groups: &[Vec<RecordBatch>],
) -> WriteResult<Vec<PartitionRows>> {
    let partitioned = contract.time_column.is_some() && contract.chunk.is_some();
    let mut fields: Vec<_> = schema.fields().iter().cloned().collect();
    if partitioned {
        fields.push(Arc::new(Field::new(PARTITION, DataType::Int32, false)));
    }
    let keyed = Arc::new(Schema::new(fields));
    let mut keys: Vec<String> = partitioned.then(|| PARTITION.to_string()).into_iter().collect();
    keys.extend(contract.write_order());
    let mut runs = vec![];
    for group in groups {
        let all = concat_batches(schema, group)?;
        let mut columns = all.columns().to_vec();
        if let Some(p) = partition_values(&all, contract)? {
            columns.push(Arc::new(Int32Array::from(p)));
        }
        let all = RecordBatch::try_new(keyed.clone(), columns)?;
        runs.push(Some(stable_sort(&all, &keys)?));
    }
    let sort_columns = contract.sort_columns();
    let unique =
        if contract.dedup && !sort_columns.is_empty() { usize::from(partitioned) + sort_columns.len() } else { 0 };
    let names: Vec<&str> = keys.iter().map(String::as_str).collect();
    let mut merge = Merge::new(keyed, &names, unique, runs.len(), contract.row_group_rows)?;
    let mut out: Vec<PartitionRows> = vec![];
    loop {
        match merge.step()? {
            Step::Need(i) => merge.push(i, runs[i].take())?,
            Step::Output(batch) => {
                for (partition, rows) in split_partitions(&batch, partitioned)? {
                    match out.last_mut() {
                        Some((p, batches)) if *p == partition => batches.push(rows),
                        _ => out.push((partition, vec![rows])),
                    }
                }
            }
            Step::Done => break,
        }
    }
    Ok(out)
}

/// A merged batch's rows by partition, without the partition column.
fn split_partitions(batch: &RecordBatch, partitioned: bool) -> WriteResult<Vec<(Option<i32>, RecordBatch)>> {
    if !partitioned {
        return Ok(vec![(None, batch.clone())]);
    }
    let last = batch.num_columns() - 1;
    let values = batch.column(last).as_any().downcast_ref::<Int32Array>().expect("the partition column");
    let rows = batch.project(&(0..last).collect::<Vec<_>>())?;
    let mut out = vec![];
    let mut offset = 0;
    for run in values.values().chunk_by(|a, b| a == b) {
        out.push((Some(run[0]), rows.slice(offset, run.len())));
        offset += run.len();
    }
    Ok(out)
}

/// Partition value of every row (`None` for unpartitioned tables).
fn partition_values(batch: &RecordBatch, contract: &Contract) -> WriteResult<Option<Vec<i32>>> {
    let (Some(time_column), Some(chunk)) = (&contract.time_column, contract.chunk) else { return Ok(None) };
    let times =
        cast(batch.column(batch.schema().index_of(time_column)?), &DataType::Timestamp(TimeUnit::Microsecond, None))?;
    let Some(times) = times.as_any().downcast_ref::<TimestampMicrosecondArray>() else {
        return Err(WriteError::invalid(format!("time column {time_column:?} is not a timestamp")));
    };
    if times.null_count() > 0 {
        return Err(WriteError::invalid(format!("time column {time_column:?} contains nulls")));
    }
    let values = times.values().iter().map(|&t| chunk.value(t)).collect::<crate::Result<_>>()?;
    Ok(Some(values))
}

/// The columns of `batch` named `names`, in that order.
pub fn columns(batch: &RecordBatch, names: &[String]) -> WriteResult<Vec<ArrayRef>> {
    let schema = batch.schema();
    names.iter().map(|n| Ok(batch.column(schema.index_of(n)?).clone())).collect()
}

/// Row indices that sort by `keys` (ascending, nulls first), keeping input order among equal keys.
pub fn stable_sort_indices(keys: &[ArrayRef], rows: usize) -> WriteResult<UInt32Array> {
    let rows = u32::try_from(rows).map_err(|_| too_many_rows(rows))?;
    let mut columns: Vec<SortColumn> =
        keys.iter().map(|k| SortColumn { values: k.clone(), options: Some(ASCENDING) }).collect();
    columns.push(SortColumn { values: Arc::new(UInt32Array::from_iter_values(0..rows)), options: None });
    Ok(lexsort_to_indices(&columns, None)?)
}

/// `batch` sorted by its columns `names`, keeping input order among equal keys.
pub fn stable_sort(batch: &RecordBatch, names: &[String]) -> WriteResult<RecordBatch> {
    if names.is_empty() || batch.num_rows() == 0 {
        return Ok(batch.clone());
    }
    let indices = stable_sort_indices(&columns(batch, names)?, batch.num_rows())?;
    Ok(take_record_batch(batch, &indices)?)
}

/// The last row of each run of equal values in the columns `names` of a sorted batch (the last
/// row overall if `names` is empty).
pub fn last_per_group(sorted: &RecordBatch, names: &[String]) -> WriteResult<RecordBatch> {
    let rows = sorted.num_rows();
    if rows == 0 {
        return Ok(sorted.clone());
    }
    if names.is_empty() {
        return Ok(sorted.slice(rows - 1, 1));
    }
    let groups = partition(&columns(sorted, names)?)?.ranges();
    if groups.len() == rows {
        return Ok(sorted.clone());
    }
    let last: UInt32Array =
        groups.iter().map(|g| u32::try_from(g.end - 1)).collect::<Result<_, _>>().map_err(|_| too_many_rows(rows))?;
    Ok(take_record_batch(sorted, &last)?)
}

/// Newest row per series among `batches`, or `None` without a time column or rows. On equal
/// times the later insertion wins, then the later row (later input position), matching dedup's
/// "latest write wins".
pub fn latest_per_series(
    schema: &SchemaRef,
    contract: &Contract,
    batches: &[RecordBatch],
) -> WriteResult<Option<RecordBatch>> {
    if contract.time_column.is_none() || batches.iter().all(|b| b.num_rows() == 0) {
        return Ok(None);
    }
    let sorted = stable_sort(&concat_batches(schema, batches)?, &contract.write_order())?;
    Ok(Some(last_per_group(&sorted, &contract.series_key)?))
}

/// Newest row per series among partitions that [`sort_and_split`] returned: each is sorted by the
/// write order, so its last row per series is its newest, and only those few rows are sorted to
/// pick the newest across partitions. The same rows as [`latest_per_series`] of the partitions,
/// without sorting the whole append again.
pub fn latest_of_sorted(
    schema: &SchemaRef,
    contract: &Contract,
    partitions: &[PartitionRows],
) -> WriteResult<Option<RecordBatch>> {
    if contract.time_column.is_none() {
        return Ok(None);
    }
    let lasts = partitions
        .iter()
        .flat_map(|(_, batches)| batches)
        .filter(|rows| rows.num_rows() > 0)
        .map(|rows| last_per_group(rows, &contract.series_key))
        .collect::<WriteResult<Vec<_>>>()?;
    latest_per_series(schema, contract, &lasts)
}

fn too_many_rows(rows: usize) -> WriteError {
    WriteError::invalid(format!("{rows} rows are too many to sort at once"))
}

#[cfg(test)]
mod tests {
    use arrow_array::{Int32Array, StringArray};
    use arrow_schema::{DataType, Field, Schema};

    use super::*;

    fn batch(keys: &[Option<i32>], tags: &[&str]) -> RecordBatch {
        let schema =
            Schema::new(vec![Field::new("k", DataType::Int32, true), Field::new("tag", DataType::Utf8, false)]);
        RecordBatch::try_new(
            Arc::new(schema),
            vec![Arc::new(Int32Array::from(keys.to_vec())), Arc::new(StringArray::from(tags.to_vec()))],
        )
        .unwrap()
    }

    fn tags(b: &RecordBatch) -> Vec<&str> {
        let tags = b.column(1).as_any().downcast_ref::<StringArray>().unwrap();
        (0..tags.len()).map(|i| tags.value(i)).collect()
    }

    #[test]
    fn sorts_nulls_first_and_keeps_input_order_on_ties() {
        let b = batch(&[Some(2), None, Some(1), Some(2), Some(1)], &["a", "b", "c", "d", "e"]);
        let sorted = stable_sort(&b, &["k".into()]).unwrap();
        assert_eq!(tags(&sorted), ["b", "c", "e", "a", "d"]);
        assert_eq!(tags(&last_per_group(&sorted, &["k".into()]).unwrap()), ["b", "e", "d"]);
    }

    #[test]
    fn without_keys_keeps_the_last_row() {
        let b = batch(&[Some(1), Some(2)], &["a", "b"]);
        assert_eq!(tags(&stable_sort(&b, &[]).unwrap()), ["a", "b"]);
        assert_eq!(tags(&last_per_group(&b, &[]).unwrap()), ["b"]);
    }

    /// The newest row per series from sorted partitions is what sorting everything gives.
    #[test]
    fn latest_of_sorted_partitions_matches_sorting_everything() {
        use std::collections::HashMap;

        use arrow_array::{Float64Array, TimestampMicrosecondArray};

        const DAY: i64 = 86_400_000_000;
        let schema = Arc::new(Schema::new(vec![
            Field::new("time", DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some(UTC.into())), false),
            Field::new("device", DataType::Int32, true),
            Field::new("value", DataType::Float64, true),
        ]));
        let contract = Contract::from_properties(&HashMap::from([
            (crate::property::TIME_COLUMN.to_string(), "time".to_string()),
            (crate::property::SERIES_KEY.to_string(), "device".to_string()),
            (crate::property::CHUNK.to_string(), "day".to_string()),
        ]))
        .unwrap();
        let rows: Vec<(i64, i32, f64)> = (0..200)
            .map(|i| {
                ((i * 7919) % 5 * DAY + i % 13, i32::try_from(i % 6).unwrap(), f64::from(i32::try_from(i).unwrap()))
            })
            .collect();
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(TimestampMicrosecondArray::from_iter_values(rows.iter().map(|r| r.0)).with_timezone(UTC)),
                Arc::new(Int32Array::from_iter_values(rows.iter().map(|r| r.1))),
                Arc::new(Float64Array::from_iter_values(rows.iter().map(|r| r.2))),
            ],
        )
        .unwrap();
        let partitions = sort_and_split(&schema, &contract, &[batch], 0).unwrap();
        assert!(partitions.len() > 1);
        let all: Vec<RecordBatch> = partitions.iter().flat_map(|(_, b)| b.clone()).collect();
        let expected = latest_per_series(&schema, &contract, &all).unwrap();
        assert_eq!(latest_of_sorted(&schema, &contract, &partitions).unwrap(), expected);
    }

    /// An append sorted in several groups and merged gives the rows sorting it at once gives, with
    /// dedup across groups (the later write wins) and with dedup off.
    #[test]
    fn merged_groups_match_one_sort() {
        use std::collections::HashMap;

        use arrow_array::{Float64Array, TimestampMicrosecondArray};

        const DAY: i64 = 86_400_000_000;
        let schema = Arc::new(Schema::new(vec![
            Field::new("time", DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some(UTC.into())), false),
            Field::new("device", DataType::Utf8, true),
            Field::new("value", DataType::Float64, true),
        ]));
        let batch = |from: i64| {
            // Times repeat across batches (every batch writes the same keys), over three days.
            let rows: Vec<i64> = (0..300).collect();
            RecordBatch::try_new(
                schema.clone(),
                vec![
                    Arc::new(
                        TimestampMicrosecondArray::from_iter_values(rows.iter().map(|i| i % 3 * DAY + i % 50))
                            .with_timezone(UTC),
                    ),
                    Arc::new(StringArray::from_iter_values(rows.iter().map(|i| format!("d{}", (i * 7) % 11)))),
                    Arc::new(Float64Array::from_iter_values(rows.iter().map(|i| (from + i) as f64))),
                ],
            )
            .unwrap()
        };
        let batches: Vec<RecordBatch> = (0..4).map(|b| batch(b * 1000)).collect();
        for dedup in ["true", "false"] {
            let contract = Contract::from_properties(&HashMap::from([
                (crate::property::TIME_COLUMN.to_string(), "time".to_string()),
                (crate::property::SERIES_KEY.to_string(), "device".to_string()),
                (crate::property::CHUNK.to_string(), "day".to_string()),
                (crate::property::DEDUP.to_string(), dedup.to_string()),
                (crate::property::ROW_GROUP_ROWS.to_string(), "64".to_string()),
            ]))
            .unwrap();
            let flat = |p: Vec<PartitionRows>| -> Vec<(Option<i32>, RecordBatch)> {
                p.into_iter().map(|(p, b)| (p, concat_batches(&schema, &b).unwrap())).collect()
            };
            let once = flat(sort_and_split_in_groups(&schema, &contract, &batches, 0, usize::MAX).unwrap());
            let merged = sort_and_split_in_groups(&schema, &contract, &batches, 0, 1).unwrap();
            assert!(merged.iter().any(|(_, b)| b.len() > 1), "batches of row_group_rows");
            assert_eq!(flat(merged), once, "dedup {dedup}");
            assert_eq!(once.len(), 3);
        }
    }

    #[test]
    fn empty_input() {
        let b = batch(&[], &[]);
        assert_eq!(last_per_group(&stable_sort(&b, &["k".into()]).unwrap(), &["k".into()]).unwrap().num_rows(), 0);
    }
}