chronix 0.7.0

//! Shaping a storage batch into the batch SQL and the archive both expect.
//!
//! The storage layer emits `timestamp: Int64` and its own canonical column
//! order; every consumer above it wants `_time: Timestamp(ns)` and the table
//! schema's order. That mapping lives here once.
//!
//! It lives here because it had **two** implementations, and they drifted: the
//! SQL scan converted the timestamp column and the cold archive did not, so
//! the same measurement was `_time` when queried hot and `_time` when
//! queried from the archive, and `time_bucket` worked on one and not the
//! other.

use std::collections::HashMap;
use std::sync::Arc;

use arrow::compute;
use arrow::datatypes::{DataType, SchemaRef, TimeUnit};
use arrow::record_batch::{RecordBatch, RecordBatchOptions};
use datafusion::common::DataFusionError;

/// Retype the time column from storage's `Int64` nanoseconds to SQL's
/// `Timestamp(Nanosecond)`.
///
/// A **retype**, not a rename: the column is `_time` on both sides. It used to
/// be a rename too — storage called it `timestamp` — which is why a reader who
/// printed the schema and then wrote `SELECT timestamp …` got a planning
/// error. The type still differs deliberately: storage stores an `i64` because
/// that is what it encodes, and SQL presents a `TIMESTAMP` so DataFusion's
/// date functions and interval arithmetic work on it.
pub(crate) fn convert_timestamp_column(batch: RecordBatch) -> Result<RecordBatch, DataFusionError> {
    if batch.num_rows() == 0 {
        // Return empty batch — avoid schema mismatches.
        return Ok(batch);
    }

    let schema = batch.schema();
    let mut new_fields = Vec::with_capacity(schema.fields().len());
    let mut new_columns = Vec::with_capacity(batch.num_columns());

    for (i, field) in schema.fields().iter().enumerate() {
        let col = batch.column(i);
        if field.name() == chronix_core::TIME_COLUMN && *field.data_type() == DataType::Int64 {
            let ts = compute::cast(col, &DataType::Timestamp(TimeUnit::Nanosecond, None))
                .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
            new_fields.push(
                arrow::datatypes::Field::new(
                    chronix_core::TIME_COLUMN,
                    DataType::Timestamp(TimeUnit::Nanosecond, None),
                    false,
                )
                .into(),
            );
            new_columns.push(ts);
        } else {
            new_fields.push(field.clone());
            new_columns.push(col.clone());
        }
    }

    let new_schema = Arc::new(arrow::datatypes::Schema::new(new_fields));
    RecordBatch::try_new(new_schema, new_columns)
        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
}

/// Reorder/select `batch`'s columns so the result matches `target` exactly,
/// resolving columns **by name**.
///
/// Columns present in `target` but absent from `batch` become all-null arrays
/// of the target type — this is the correct reading for a measurement whose
/// schema has gained a field that this particular segment predates.
///
/// This replaces an index-based `RecordBatch::project`, which assumed
/// the storage batch and the DataFusion table schema agreed on column order.
/// They do not: storage emits `timestamp, tags…, fields…` each sorted by name,
/// while the table schema follows schema-registration order.
pub(crate) fn align_batch_to_schema(
    batch: &RecordBatch,
    target: &SchemaRef,
) -> Result<RecordBatch, DataFusionError> {
    let src = batch.schema();

    // Fast path: already identical, no work to do.
    //
    // Identical means the *whole* schema, not just the column names. A scan
    // batch carries each column's role in its field metadata, and returning
    // it here left the caller with a batch whose schema is not the one it
    // asked to align to — which a Parquet writer built for `target` then
    // refuses, silently, because the cold archive counts an encode failure
    // and leaves the group hot.
    if src.as_ref() == target.as_ref() {
        return Ok(batch.clone());
    }

    let index: HashMap<&str, usize> = src
        .fields()
        .iter()
        .enumerate()
        .map(|(i, f)| (f.name().as_str(), i))
        .collect();

    let mut columns = Vec::with_capacity(target.fields().len());
    for field in target.fields() {
        match index.get(field.name().as_str()) {
            Some(&i) => {
                let col = batch.column(i);
                // Storage may hand back a narrower/wider numeric type than the
                // table schema advertises; cast rather than fail the query.
                if col.data_type() == field.data_type() {
                    columns.push(Arc::clone(col));
                } else {
                    columns.push(
                        compute::cast(col, field.data_type())
                            .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?,
                    );
                }
            }
            None => columns.push(arrow::array::new_null_array(
                field.data_type(),
                batch.num_rows(),
            )),
        }
    }

    // The row count is carried explicitly rather than inferred from the
    // columns, because an aggregate that reads no column — `count(*)` — pushes
    // an **empty** projection down, and Arrow refuses a zero-column batch
    // unless it is told how many rows it stands for. Inferring it made the
    // most ordinary SQL query there is fail with "must either specify a row
    // count or at least one column".
    let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
    RecordBatch::try_new_with_options(Arc::clone(target), columns, &options)
        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
}

/// Shape a read-path batch into an archive object's schema.
///
/// The two steps the SQL scan also takes, in the same order: rename and retype
/// the timestamp column, then resolve columns by name against the target.
///
/// # Errors
///
/// Returns an error if a column cannot be cast to its target type.
#[cfg(feature = "object-store")]
pub(crate) fn to_archive_batch(
    batch: &RecordBatch,
    target: &SchemaRef,
) -> Result<RecordBatch, DataFusionError> {
    let converted = convert_timestamp_column(batch.clone())?;
    align_batch_to_schema(&converted, target)
}

#[cfg(test)]
mod align_tests {
    #![allow(clippy::unwrap_used, clippy::expect_used)]

    use super::*;
    use arrow::array::Int64Array;
    use arrow::datatypes::{Field, Schema};

    /// Aligning a batch to a schema yields a batch **with that schema**.
    ///
    /// The fast path compared column *names* and returned the input, so a
    /// batch whose fields carry metadata the target does not — a scan batch
    /// carries each column's role — came back with the wrong schema. A
    /// Parquet writer built for the target then refused it, and the cold
    /// archive counts an encode failure and leaves the group hot: nothing was
    /// archived, and nothing said so louder than a warning.
    #[test]
    fn the_result_carries_the_target_schema_even_when_the_names_match() {
        let src = Arc::new(Schema::new(vec![
            Field::new("v", DataType::Int64, true)
                .with_metadata(crate::db::role_metadata(chronix_core::ColumnRole::Field)),
        ]));
        let batch =
            RecordBatch::try_new(src, vec![Arc::new(Int64Array::from(vec![1_i64, 2]))]).unwrap();

        let target: SchemaRef = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, true)]));
        let aligned = align_batch_to_schema(&batch, &target).unwrap();
        assert_eq!(aligned.schema().as_ref(), target.as_ref());
        assert_eq!(aligned.num_rows(), 2);
    }

    #[test]
    fn an_identical_schema_is_returned_unchanged() {
        let schema: SchemaRef = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, true)]));
        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(Int64Array::from(vec![1_i64]))],
        )
        .unwrap();
        let aligned = align_batch_to_schema(&batch, &schema).unwrap();
        assert_eq!(aligned.schema().as_ref(), schema.as_ref());
    }
}

// The archive conversion only exists behind `object-store`, and so do its
// tests — an unconditional test module would import a function that is not
// compiled.
#[cfg(all(test, feature = "object-store"))]
mod tests {
    #![allow(clippy::unwrap_used, clippy::expect_used)]

    use super::*;
    use arrow::array::{Float64Array, Int64Array, StringArray};
    use arrow::datatypes::{Field, Schema};

    fn storage_batch() -> RecordBatch {
        let schema = Arc::new(Schema::new(vec![
            Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false),
            Field::new("host", DataType::Utf8, true),
            Field::new("watts", DataType::Float64, true),
        ]));
        RecordBatch::try_new(
            schema,
            vec![
                Arc::new(Int64Array::from(vec![1_i64, 2])),
                Arc::new(StringArray::from(vec!["a", "b"])),
                Arc::new(Float64Array::from(vec![1.0, 2.0])),
            ],
        )
        .unwrap()
    }

    /// The archive must get `_time`, the target's column order, and nulls for
    /// a field this batch predates.
    #[test]
    fn archive_batch_matches_the_target_schema() {
        let target: SchemaRef = Arc::new(Schema::new(vec![
            Field::new(
                "_time",
                DataType::Timestamp(TimeUnit::Nanosecond, None),
                false,
            ),
            Field::new("host", DataType::Utf8, true),
            Field::new("amps", DataType::Float64, true),
            Field::new("watts", DataType::Float64, true),
        ]));

        let out = to_archive_batch(&storage_batch(), &target).unwrap();
        assert_eq!(out.schema(), target, "the archive schema is the contract");
        assert_eq!(out.num_rows(), 2);
        assert!(
            out.column(2).null_count() == 2,
            "a field the batch predates must be null, not absent"
        );
    }
}