promql/extension_plan/

series_divide.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use datafusion::arrow::array::{
    Array, ArrayRef, DictionaryArray, LargeStringArray, PrimitiveArray, StringArray,
    StringViewArray, UInt64Array,
};
use datafusion::arrow::buffer::NullBuffer;
use datafusion::arrow::datatypes::{ArrowDictionaryKeyType, DataType, SchemaRef};
use datafusion::arrow::downcast_dictionary_array;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::common::{DFSchema, DFSchemaRef};
use datafusion::error::Result as DataFusionResult;
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::{EmptyRelation, Expr, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::{LexRequirement, OrderingRequirements, PhysicalSortRequirement};
use datafusion::physical_plan::expressions::Column as ColumnExpr;
use datafusion::physical_plan::metrics::{
    BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricValue, MetricsSet,
};
use datafusion::physical_plan::{
    DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements,
    PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
};
use datafusion_expr::ident;
use datatypes::arrow::compute;
use datatypes::arrow_array::string_array_value_at_index;
use datatypes::compute::SortOptions;
use futures::{Stream, StreamExt, ready};
use greptime_proto::substrait_extension as pb;
use prost::Message;
use snafu::ResultExt;

use crate::error::{DeserializeSnafu, Result};
use crate::extension_plan::{METRIC_NUM_SERIES, resolve_column_name, serialize_column_index};
use crate::metrics::PROMQL_SERIES_COUNT;

enum TagIdentifier<'a> {
    /// A group of raw string tag columns.
    Raw(Vec<RawTagColumn<'a>>),
    /// A single UInt64 identifier (tsid).
    Id(&'a UInt64Array),
}

impl<'a> TagIdentifier<'a> {
    fn try_new(batch: &'a RecordBatch, tag_indices: &[usize]) -> DataFusionResult<Self> {
        match tag_indices {
            [] => Ok(Self::Raw(Vec::new())),
            [index] => {
                let array = batch.column(*index);
                if array.data_type() == &DataType::UInt64 {
                    let array = array
                        .as_any()
                        .downcast_ref::<UInt64Array>()
                        .ok_or_else(|| {
                            datafusion::error::DataFusionError::Internal(
                                "Failed to downcast tag column to UInt64Array".to_string(),
                            )
                        })?;
                    Ok(Self::Id(array))
                } else {
                    Ok(Self::Raw(vec![RawTagColumn::try_new(array)?]))
                }
            }
            indices => Ok(Self::Raw(
                indices
                    .iter()
                    .map(|index| RawTagColumn::try_new(batch.column(*index)))
                    .collect::<DataFusionResult<Vec<_>>>()?,
            )),
        }
    }

    fn equal_at(&self, left_row: usize, other: &Self, right_row: usize) -> DataFusionResult<bool> {
        match (self, other) {
            (Self::Id(left), Self::Id(right)) => {
                if left.is_null(left_row) || right.is_null(right_row) {
                    return Ok(left.is_null(left_row) && right.is_null(right_row));
                }
                Ok(left.value(left_row) == right.value(right_row))
            }
            (Self::Raw(left), Self::Raw(right)) => {
                if left.len() != right.len() {
                    return Err(datafusion::error::DataFusionError::Internal(format!(
                        "Mismatched tag column count: left={}, right={}",
                        left.len(),
                        right.len()
                    )));
                }

                for (left_column, right_column) in left.iter().zip(right.iter()) {
                    if !left_column.equal_at(left_row, right_column, right_row) {
                        return Ok(false);
                    }
                }
                Ok(true)
            }
            _ => Err(datafusion::error::DataFusionError::Internal(format!(
                "Mismatched tag identifier types: left={:?}, right={:?}",
                self.data_type(),
                other.data_type()
            ))),
        }
    }

    fn data_type(&self) -> &'static str {
        match self {
            Self::Raw(_) => "Raw",
            Self::Id(_) => "Id",
        }
    }
}

struct RawTagColumn<'a>(&'a ArrayRef);

impl<'a> RawTagColumn<'a> {
    fn try_new(array: &'a ArrayRef) -> DataFusionResult<Self> {
        match array.data_type() {
            DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Ok(Self(array)),
            DataType::Dictionary(key, value) if key.is_integer() && value.is_string() => {
                Ok(Self(array))
            }
            other => Err(datafusion::error::DataFusionError::Internal(format!(
                "Unsupported tag column type: {other:?}"
            ))),
        }
    }

    fn equal_at(&self, left_row: usize, other: &Self, right_row: usize) -> bool {
        string_array_value_at_index(self.0, left_row)
            == string_array_value_at_index(other.0, right_row)
    }
}

#[derive(Debug, PartialEq, Eq, Hash, PartialOrd)]
pub struct SeriesDivide {
    tag_columns: Vec<String>,
    /// `SeriesDivide` requires `time_index` column's name to generate ordering requirement
    /// for input data. But this plan itself doesn't depend on the ordering of time index
    /// column. This is for follow on plans like `RangeManipulate`. Because requiring ordering
    /// here can avoid unnecessary sort in follow on plans.
    time_index_column: String,
    input: LogicalPlan,
    unfix: Option<UnfixIndices>,
}

#[derive(Debug, PartialEq, Eq, Hash, PartialOrd)]
struct UnfixIndices {
    pub tag_column_indices: Vec<u64>,
    pub time_index_column_idx: u64,
}

impl UserDefinedLogicalNodeCore for SeriesDivide {
    fn name(&self) -> &str {
        Self::name()
    }

    fn inputs(&self) -> Vec<&LogicalPlan> {
        vec![&self.input]
    }

    fn schema(&self) -> &DFSchemaRef {
        self.input.schema()
    }

    fn expressions(&self) -> Vec<Expr> {
        if self.unfix.is_some() {
            return vec![];
        }

        self.tag_columns
            .iter()
            .map(ident)
            .chain(std::iter::once(ident(&self.time_index_column)))
            .collect()
    }

    fn necessary_children_exprs(&self, output_columns: &[usize]) -> Option<Vec<Vec<usize>>> {
        if self.unfix.is_some() {
            return None;
        }

        let input_schema = self.input.schema();
        if output_columns.is_empty() {
            let indices = (0..input_schema.fields().len()).collect::<Vec<_>>();
            return Some(vec![indices]);
        }

        let mut required = Vec::with_capacity(output_columns.len() + 1 + self.tag_columns.len());
        required.extend_from_slice(output_columns);
        for tag in &self.tag_columns {
            required.push(input_schema.index_of_column_by_name(None, tag)?);
        }
        required.push(input_schema.index_of_column_by_name(None, &self.time_index_column)?);

        required.sort_unstable();
        required.dedup();
        Some(vec![required])
    }

    fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "PromSeriesDivide: tags={:?}", self.tag_columns)
    }

    fn with_exprs_and_inputs(
        &self,
        _exprs: Vec<Expr>,
        inputs: Vec<LogicalPlan>,
    ) -> DataFusionResult<Self> {
        if inputs.is_empty() {
            return Err(datafusion::error::DataFusionError::Internal(
                "SeriesDivide must have at least one input".to_string(),
            ));
        }

        let input: LogicalPlan = inputs[0].clone();
        let input_schema = input.schema();

        if let Some(unfix) = &self.unfix {
            // transform indices to names
            let tag_columns = unfix
                .tag_column_indices
                .iter()
                .map(|idx| resolve_column_name(*idx, input_schema, "SeriesDivide", "tag"))
                .collect::<DataFusionResult<Vec<String>>>()?;

            let time_index_column = resolve_column_name(
                unfix.time_index_column_idx,
                input_schema,
                "SeriesDivide",
                "time index",
            )?;

            Ok(Self {
                tag_columns,
                time_index_column,
                input,
                unfix: None,
            })
        } else {
            Ok(Self {
                tag_columns: self.tag_columns.clone(),
                time_index_column: self.time_index_column.clone(),
                input,
                unfix: None,
            })
        }
    }
}

impl SeriesDivide {
    pub fn new(tag_columns: Vec<String>, time_index_column: String, input: LogicalPlan) -> Self {
        Self {
            tag_columns,
            time_index_column,
            input,
            unfix: None,
        }
    }

    pub const fn name() -> &'static str {
        "SeriesDivide"
    }

    pub fn to_execution_plan(&self, exec_input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
        Arc::new(SeriesDivideExec {
            tag_columns: self.tag_columns.clone(),
            time_index_column: self.time_index_column.clone(),
            input: exec_input,
            metric: ExecutionPlanMetricsSet::new(),
        })
    }

    pub fn tags(&self) -> &[String] {
        &self.tag_columns
    }

    pub fn serialize(&self) -> Vec<u8> {
        let tag_column_indices = self
            .tag_columns
            .iter()
            .map(|name| serialize_column_index(self.input.schema(), name))
            .collect::<Vec<u64>>();

        let time_index_column_idx =
            serialize_column_index(self.input.schema(), &self.time_index_column);

        pb::SeriesDivide {
            tag_column_indices,
            time_index_column_idx,
            ..Default::default()
        }
        .encode_to_vec()
    }

    pub fn deserialize(bytes: &[u8]) -> Result<Self> {
        let pb_series_divide = pb::SeriesDivide::decode(bytes).context(DeserializeSnafu)?;
        let placeholder_plan = LogicalPlan::EmptyRelation(EmptyRelation {
            produce_one_row: false,
            schema: Arc::new(DFSchema::empty()),
        });

        let unfix = UnfixIndices {
            tag_column_indices: pb_series_divide.tag_column_indices.clone(),
            time_index_column_idx: pb_series_divide.time_index_column_idx,
        };

        Ok(Self {
            tag_columns: Vec::new(),
            time_index_column: String::new(),
            input: placeholder_plan,
            unfix: Some(unfix),
        })
    }
}

#[derive(Debug)]
pub struct SeriesDivideExec {
    tag_columns: Vec<String>,
    time_index_column: String,
    input: Arc<dyn ExecutionPlan>,
    metric: ExecutionPlanMetricsSet,
}

impl ExecutionPlan for SeriesDivideExec {
    fn apply_expressions(
        &self,
        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
    ) -> DataFusionResult<TreeNodeRecursion> {
        Ok(TreeNodeRecursion::Continue)
    }

    fn schema(&self) -> SchemaRef {
        self.input.schema()
    }

    fn properties(&self) -> &Arc<PlanProperties> {
        self.input.properties()
    }

    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
        if self.tag_columns.is_empty() {
            return InputDistributionRequirements::new(vec![Distribution::SinglePartition]);
        }
        let schema = self.input.schema();
        InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
            self.tag_columns
                .iter()
                // Safety: the tag column names is verified in the planning phase
                .map(|tag| Arc::new(ColumnExpr::new_with_schema(tag, &schema).unwrap()) as _)
                .collect(),
        )])
    }

    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
        let input_schema = self.input.schema();
        let mut exprs: Vec<PhysicalSortRequirement> = self
            .tag_columns
            .iter()
            .map(|tag| PhysicalSortRequirement {
                // Safety: the tag column names is verified in the planning phase
                expr: Arc::new(ColumnExpr::new_with_schema(tag, &input_schema).unwrap()),
                options: Some(SortOptions {
                    descending: false,
                    nulls_first: true,
                }),
            })
            .collect();

        exprs.push(PhysicalSortRequirement {
            expr: Arc::new(
                ColumnExpr::new_with_schema(&self.time_index_column, &input_schema).unwrap(),
            ),
            options: Some(SortOptions {
                descending: false,
                nulls_first: true,
            }),
        });

        // Safety: `exprs` is not empty
        let requirement = LexRequirement::new(exprs).unwrap();

        vec![Some(OrderingRequirements::Hard(vec![requirement]))]
    }

    fn maintains_input_order(&self) -> Vec<bool> {
        vec![true; self.children().len()]
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        vec![&self.input]
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
        assert!(!children.is_empty());
        Ok(Arc::new(Self {
            tag_columns: self.tag_columns.clone(),
            time_index_column: self.time_index_column.clone(),
            input: children[0].clone(),
            metric: self.metric.clone(),
        }))
    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> DataFusionResult<SendableRecordBatchStream> {
        let baseline_metric = BaselineMetrics::new(&self.metric, partition);
        let metrics_builder = MetricBuilder::new(&self.metric);
        let num_series = Count::new();
        metrics_builder
            .with_partition(partition)
            .build(MetricValue::Count {
                name: METRIC_NUM_SERIES.into(),
                count: num_series.clone(),
            });

        let input = self.input.execute(partition, context)?;
        let schema = input.schema();
        let tag_indices = self
            .tag_columns
            .iter()
            .map(|tag| {
                schema
                    .column_with_name(tag)
                    .unwrap_or_else(|| panic!("tag column not found {tag}"))
                    .0
            })
            .collect();
        Ok(Box::pin(SeriesDivideStream {
            tag_indices,
            buffer: vec![],
            schema,
            input,
            metric: baseline_metric,
            num_series,
            inspect_start: 0,
        }))
    }

    fn metrics(&self) -> Option<MetricsSet> {
        Some(self.metric.clone_inner())
    }

    fn name(&self) -> &str {
        "SeriesDivideExec"
    }
}

impl DisplayAs for SeriesDivideExec {
    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match t {
            DisplayFormatType::Default
            | DisplayFormatType::Verbose
            | DisplayFormatType::TreeRender => {
                write!(f, "PromSeriesDivideExec: tags={:?}", self.tag_columns)
            }
        }
    }
}

/// Assume the input stream is ordered on the tag columns.
pub struct SeriesDivideStream {
    tag_indices: Vec<usize>,
    buffer: Vec<RecordBatch>,
    schema: SchemaRef,
    input: SendableRecordBatchStream,
    metric: BaselineMetrics,
    /// Index of buffered batches to start inspect next time.
    inspect_start: usize,
    /// Number of series processed.
    num_series: Count,
}

impl RecordBatchStream for SeriesDivideStream {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

fn constant_string_dictionary<K: ArrowDictionaryKeyType>(
    dictionary: &DictionaryArray<K>,
    total_rows: usize,
) -> DataFusionResult<ArrayRef> {
    let key = dictionary.key(0);
    let value = key.and_then(|key| string_array_value_at_index(dictionary.values(), key));
    let values: ArrayRef = match dictionary.values().data_type() {
        DataType::Utf8 => Arc::new(StringArray::from(vec![value])),
        DataType::LargeUtf8 => Arc::new(LargeStringArray::from(vec![value])),
        DataType::Utf8View => Arc::new(StringViewArray::from(vec![value])),
        _ => unreachable!("dictionary values must be strings"),
    };
    let keys = PrimitiveArray::<K>::new(
        vec![K::Native::default(); total_rows].into(),
        key.is_none().then(|| NullBuffer::new_null(total_rows)),
    );

    Ok(Arc::new(DictionaryArray::try_new(keys, values)?))
}

/// Concatenates batches from one isolated series; every designated tag is constant.
fn concat_series_batches(
    schema: &SchemaRef,
    batches: &[RecordBatch],
    tag_indices: &[usize],
) -> DataFusionResult<RecordBatch> {
    if batches.len() <= 1 || tag_indices.is_empty() {
        return Ok(compute::concat_batches(schema, batches)?);
    }

    let Some(first_batch) = batches.iter().find(|batch| batch.num_rows() > 0) else {
        return Ok(compute::concat_batches(schema, batches)?);
    };

    // This endpoint-only sanity check does not validate tags in interior rows.
    #[cfg(debug_assertions)]
    {
        let last_batch = batches
            .iter()
            .rfind(|batch| batch.num_rows() > 0)
            .expect("first non-empty batch implies a last non-empty batch");
        if let (Ok(first_tags), Ok(last_tags)) = (
            TagIdentifier::try_new(first_batch, tag_indices),
            TagIdentifier::try_new(last_batch, tag_indices),
        ) {
            debug_assert!(
                first_tags.equal_at(0, &last_tags, last_batch.num_rows() - 1)?,
                "series batch tag endpoints must match"
            );
        }
    }

    let total_rows: usize = batches.iter().map(RecordBatch::num_rows).sum();
    let columns = schema
        .fields()
        .iter()
        .enumerate()
        .map(|(index, field)| -> DataFusionResult<ArrayRef> {
            if tag_indices.contains(&index)
                && matches!(field.data_type(), DataType::Dictionary(_, value_type) if value_type.is_string())
            {
                let array = first_batch.column(index);
                downcast_dictionary_array! {
                    array => constant_string_dictionary(array, total_rows),
                    _ => unreachable!("dictionary keys must be integers"),
                }
            } else {
                compute::concat(
                    &batches
                        .iter()
                        .map(|batch| batch.column(index).as_ref())
                        .collect::<Vec<_>>(),
                )
                .map_err(Into::into)
            }
        })
        .collect::<DataFusionResult<Vec<_>>>()?;

    RecordBatch::try_new(schema.clone(), columns).map_err(Into::into)
}

impl Stream for SeriesDivideStream {
    type Item = DataFusionResult<RecordBatch>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        loop {
            if !self.buffer.is_empty() {
                let timer = std::time::Instant::now();
                let cut_at = match self.find_first_diff_row() {
                    Ok(cut_at) => cut_at,
                    Err(e) => return Poll::Ready(Some(Err(e))),
                };
                if let Some((batch_index, row_index)) = cut_at {
                    // slice out the first time series and return it.
                    let half_batch_of_first_series =
                        self.buffer[batch_index].slice(0, row_index + 1);
                    let half_batch_of_second_series = self.buffer[batch_index].slice(
                        row_index + 1,
                        self.buffer[batch_index].num_rows() - row_index - 1,
                    );
                    let result_batches = self
                        .buffer
                        .drain(0..batch_index)
                        .chain([half_batch_of_first_series])
                        .collect::<Vec<_>>();
                    if half_batch_of_second_series.num_rows() > 0 {
                        self.buffer[0] = half_batch_of_second_series;
                    } else {
                        self.buffer.remove(0);
                    }
                    let result_batch =
                        concat_series_batches(&self.schema, &result_batches, &self.tag_indices)?;

                    self.inspect_start = 0;
                    self.num_series.add(1);
                    self.metric.elapsed_compute().add_elapsed(timer);
                    return Poll::Ready(Some(Ok(result_batch)));
                } else {
                    self.metric.elapsed_compute().add_elapsed(timer);
                    // continue to fetch next batch as the current buffer only contains one time series.
                    let next_batch = ready!(self.as_mut().fetch_next_batch(cx)).transpose()?;
                    let timer = std::time::Instant::now();
                    if let Some(next_batch) = next_batch {
                        if next_batch.num_rows() != 0 {
                            self.buffer.push(next_batch);
                        }
                        continue;
                    } else {
                        // input stream is ended
                        let result =
                            concat_series_batches(&self.schema, &self.buffer, &self.tag_indices)?;
                        self.buffer.clear();
                        self.inspect_start = 0;
                        self.num_series.add(1);
                        self.metric.elapsed_compute().add_elapsed(timer);
                        return Poll::Ready(Some(Ok(result)));
                    }
                }
            } else {
                let batch = match ready!(self.as_mut().fetch_next_batch(cx)) {
                    Some(Ok(batch)) => batch,
                    None => {
                        PROMQL_SERIES_COUNT.observe(self.num_series.value() as f64);
                        return Poll::Ready(None);
                    }
                    error => return Poll::Ready(error),
                };
                if batch.num_rows() != 0 {
                    self.buffer.push(batch);
                }
                continue;
            }
        }
    }
}

impl SeriesDivideStream {
    fn fetch_next_batch(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<DataFusionResult<RecordBatch>>> {
        let poll = self.input.poll_next_unpin(cx);
        self.metric.record_poll(poll)
    }

    /// Return the position to cut buffer.
    /// None implies the current buffer only contains one time series.
    fn find_first_diff_row(&mut self) -> DataFusionResult<Option<(usize, usize)>> {
        // fast path: no tag columns means all data belongs to the same series.
        if self.tag_indices.is_empty() {
            return Ok(None);
        }

        let mut resumed_batch_index = self.inspect_start;

        for batch in &self.buffer[resumed_batch_index..] {
            let num_rows = batch.num_rows();
            let tags = TagIdentifier::try_new(batch, &self.tag_indices)?;

            // check if the first row is the same with last batch's last row
            if resumed_batch_index > self.inspect_start.saturating_sub(1) {
                let last_batch = &self.buffer[resumed_batch_index - 1];
                let last_row = last_batch.num_rows() - 1;
                let last_tags = TagIdentifier::try_new(last_batch, &self.tag_indices)?;
                if !tags.equal_at(0, &last_tags, last_row)? {
                    return Ok(Some((resumed_batch_index - 1, last_row)));
                }
            }

            // quick check if all rows are the same by comparing the first and last row in this batch
            if tags.equal_at(0, &tags, num_rows - 1)? {
                resumed_batch_index += 1;
                continue;
            }

            let mut same_until = 0;
            while same_until < num_rows - 1 {
                if !tags.equal_at(same_until, &tags, same_until + 1)? {
                    break;
                }
                same_until += 1;
            }

            if same_until + 1 >= num_rows {
                // all rows are the same, inspect next batch
                resumed_batch_index += 1;
            } else {
                return Ok(Some((resumed_batch_index, same_until)));
            }
        }

        self.inspect_start = resumed_batch_index;
        Ok(None)
    }
}

#[cfg(test)]
mod test {
    use std::collections::HashMap;

    use datafusion::arrow::array::{
        DictionaryArray, Int32Array, Int64Array, LargeStringArray, PrimitiveArray, StringArray,
        StringViewArray, UInt32Array, UInt64Array,
    };
    use datafusion::arrow::datatypes::{
        ArrowDictionaryKeyType, DataType, Field, Int8Type, Int16Type, Int32Type, Int64Type, Schema,
        UInt8Type, UInt16Type, UInt32Type, UInt64Type,
    };
    use datafusion::common::ToDFSchema;
    use datafusion::datasource::memory::MemorySourceConfig;
    use datafusion::datasource::source::DataSourceExec;
    use datafusion::logical_expr::{EmptyRelation, LogicalPlan};
    use datafusion::prelude::SessionContext;

    use super::*;

    fn assert_concat_matches_reference(
        schema: SchemaRef,
        batches: Vec<RecordBatch>,
        tags: &[usize],
    ) {
        let expected = compute::concat_batches(&schema, &batches).unwrap();
        let actual = concat_series_batches(&schema, &batches, tags).unwrap();
        assert_eq!(actual.schema(), schema);
        assert_eq!(actual.num_rows(), expected.num_rows());

        for (index, (actual, expected)) in
            actual.columns().iter().zip(expected.columns()).enumerate()
        {
            assert_eq!(actual.data_type(), expected.data_type(), "column {index}");
            let (actual, expected) = match actual.data_type() {
                DataType::Dictionary(_, value_type) => (
                    compute::cast(actual.as_ref(), value_type).unwrap(),
                    compute::cast(expected.as_ref(), value_type).unwrap(),
                ),
                _ => (actual.clone(), expected.clone()),
            };
            assert_eq!(actual.to_data(), expected.to_data(), "column {index}");
        }
    }

    #[test]
    fn test_concat_series_batches() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("tag", DataType::Utf8, true),
            Field::new("value", DataType::Int64, true),
        ]));
        let batches = vec![
            RecordBatch::try_new(
                schema.clone(),
                vec![
                    Arc::new(StringArray::from(vec![Some("tag"), Some("tag")])),
                    Arc::new(Int64Array::from(vec![Some(1), None])),
                ],
            )
            .unwrap(),
            RecordBatch::try_new(
                schema.clone(),
                vec![
                    Arc::new(StringArray::from(vec![Some("tag"), Some("tag")])),
                    Arc::new(Int64Array::from(vec![Some(3), Some(4)])),
                ],
            )
            .unwrap(),
        ];
        assert_concat_matches_reference(schema, batches, &[0]);
    }

    #[cfg(debug_assertions)]
    #[test]
    #[should_panic(expected = "series batch tag endpoints must match")]
    fn test_concat_series_batches_mismatched_tag_endpoints_panics() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("tag", DataType::Utf8, true),
            Field::new("value", DataType::Int64, true),
        ]));
        let batches = vec![
            RecordBatch::try_new(
                schema.clone(),
                vec![
                    Arc::new(StringArray::from(vec!["first"])),
                    Arc::new(Int64Array::from(vec![1])),
                ],
            )
            .unwrap(),
            RecordBatch::try_new(
                schema.clone(),
                vec![
                    Arc::new(StringArray::from(vec!["last"])),
                    Arc::new(Int64Array::from(vec![2])),
                ],
            )
            .unwrap(),
        ];

        concat_series_batches(&schema, &batches, &[0]).unwrap();
    }

    #[test]
    fn test_concat_series_batches_string_tags() {
        for (data_type, batches) in [
            (
                DataType::LargeUtf8,
                vec![
                    Arc::new(LargeStringArray::from(vec![
                        Some("long tag"),
                        Some("long tag"),
                    ])) as ArrayRef,
                    Arc::new(LargeStringArray::from(vec![Some("long tag")])) as ArrayRef,
                ],
            ),
            (
                DataType::Utf8View,
                vec![
                    Arc::new(StringViewArray::from(vec![
                        Some("view tag longer than twelve bytes"),
                        Some("view tag longer than twelve bytes"),
                    ])) as ArrayRef,
                    Arc::new(StringViewArray::from(vec![Some(
                        "view tag longer than twelve bytes",
                    )])) as ArrayRef,
                ],
            ),
            (
                DataType::Utf8,
                vec![
                    Arc::new(StringArray::from(vec![None::<&str>, None])) as ArrayRef,
                    Arc::new(StringArray::from(vec![None::<&str>])) as ArrayRef,
                ],
            ),
            (
                DataType::Utf8,
                vec![
                    Arc::new(StringArray::from(vec![Some(""), Some("")])) as ArrayRef,
                    Arc::new(StringArray::from(vec![Some("")])) as ArrayRef,
                ],
            ),
        ] {
            let schema = Arc::new(Schema::new(vec![
                Field::new("tag", data_type, true),
                Field::new("value", DataType::Int64, true),
            ]));
            let batches = batches
                .into_iter()
                .enumerate()
                .map(|(index, tag)| {
                    RecordBatch::try_new(
                        schema.clone(),
                        vec![
                            tag.clone(),
                            Arc::new(Int64Array::from_iter_values(
                                (0..tag.len()).map(|row| (index * 10 + row) as i64),
                            )),
                        ],
                    )
                    .unwrap()
                })
                .collect();
            assert_concat_matches_reference(schema, batches, &[0]);
        }
    }

    fn assert_dictionary_concat<K: ArrowDictionaryKeyType>(
        keys: Vec<Option<K::Native>>,
        values: ArrayRef,
        other_keys: Vec<Option<K::Native>>,
        other_values: ArrayRef,
    ) where
        PrimitiveArray<K>: From<Vec<Option<K::Native>>>,
    {
        let data_type = DataType::Dictionary(
            Box::new(PrimitiveArray::<K>::from(keys.clone()).data_type().clone()),
            Box::new(values.data_type().clone()),
        );
        let schema = Arc::new(Schema::new(vec![
            Field::new("tag", data_type, true),
            Field::new("value", DataType::Int64, true),
        ]));
        let first = Arc::new(DictionaryArray::<K>::new(
            PrimitiveArray::<K>::from(keys),
            values,
        )) as ArrayRef;
        let second = Arc::new(DictionaryArray::<K>::new(
            PrimitiveArray::<K>::from(other_keys),
            other_values,
        )) as ArrayRef;
        let batches = [first, second]
            .into_iter()
            .enumerate()
            .map(|(index, tag)| {
                RecordBatch::try_new(
                    schema.clone(),
                    vec![
                        tag.clone(),
                        Arc::new(Int64Array::from_iter_values(
                            (0..tag.len()).map(|row| (index * 10 + row) as i64),
                        )),
                    ],
                )
                .unwrap()
            })
            .collect();
        assert_concat_matches_reference(schema, batches, &[0]);
    }

    #[test]
    fn test_concat_series_batches_dictionary_tags() {
        macro_rules! dictionary_cases {
            ($($key_type:ty, $key:expr),+ $(,)?) => {
                $(for (values, other_values) in [
                    (
                        Arc::new(StringArray::from(vec!["other", "tag"])) as ArrayRef,
                        Arc::new(StringArray::from(vec!["tag", "other"])) as ArrayRef,
                    ),
                    (
                        Arc::new(LargeStringArray::from(vec!["other", "tag"])) as ArrayRef,
                        Arc::new(LargeStringArray::from(vec!["tag", "other"])) as ArrayRef,
                    ),
                    (
                        Arc::new(StringViewArray::from(vec![
                            "other",
                            "view tag longer than twelve bytes",
                        ])) as ArrayRef,
                        Arc::new(StringViewArray::from(vec![
                            "view tag longer than twelve bytes",
                            "other",
                        ])) as ArrayRef,
                    ),
                ] {
                    assert_dictionary_concat::<$key_type>(
                        vec![Some(($key)(1)), Some(($key)(1))],
                        values,
                        vec![Some(($key)(0))],
                        other_values,
                    );
                })+
            };
        }
        dictionary_cases!(
            Int8Type,
            |value| value as i8,
            Int16Type,
            |value| value as i16,
            Int32Type,
            |value: i32| value,
            Int64Type,
            |value| value as i64,
            UInt8Type,
            |value| value as u8,
            UInt16Type,
            |value| value as u16,
            UInt32Type,
            |value| value as u32,
            UInt64Type,
            |value| value as u64,
        );
        assert_dictionary_concat::<UInt16Type>(
            vec![Some(0), Some(0)],
            Arc::new(LargeStringArray::from(vec![None::<&str>])),
            vec![None],
            Arc::new(LargeStringArray::from(vec!["unused"])),
        );
        // Compare logical values rather than dictionary keys for the reverse null orientation.
        assert_dictionary_concat::<UInt16Type>(
            vec![None],
            Arc::new(StringArray::from(vec!["unused"])),
            vec![Some(0)],
            Arc::new(StringArray::from(vec![None::<&str>])),
        );
    }

    #[test]
    fn test_concat_series_batches_dictionary_tag_drops_unused_values() {
        let active = "active view tag longer than twelve bytes";
        let unused = "unused view tag ".repeat(1024);
        let values = Arc::new(StringViewArray::from(vec![unused.as_str(), active])) as ArrayRef;
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "tag",
                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8View)),
                false,
            ),
            Field::new("value", DataType::Int64, false),
        ]));
        let batches = [vec![1, 1], vec![1]]
            .into_iter()
            .enumerate()
            .map(|(batch, keys)| {
                let num_rows = keys.len();
                RecordBatch::try_new(
                    schema.clone(),
                    vec![
                        Arc::new(DictionaryArray::<UInt32Type>::new(
                            UInt32Array::from(keys),
                            values.clone(),
                        )),
                        Arc::new(Int64Array::from_iter_values(
                            (0..num_rows).map(|row| (batch * 10 + row) as i64),
                        )),
                    ],
                )
                .unwrap()
            })
            .collect::<Vec<_>>();

        assert_concat_matches_reference(schema.clone(), batches.clone(), &[0]);

        let actual = concat_series_batches(&schema, &batches, &[0]).unwrap();
        let tag = actual
            .column(0)
            .as_any()
            .downcast_ref::<DictionaryArray<UInt32Type>>()
            .unwrap();
        let values = tag
            .values()
            .as_any()
            .downcast_ref::<StringViewArray>()
            .unwrap();
        assert_eq!(values.len(), 1);
        assert_eq!(values.value(0), active);
        assert_eq!(
            values
                .data_buffers()
                .iter()
                .map(|buffer| buffer.len())
                .sum::<usize>(),
            active.len()
        );
    }

    #[test]
    fn test_concat_series_batches_uint64_nullable_tag() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("tsid", DataType::UInt64, true),
            Field::new("value", DataType::Int64, true),
        ]));
        for tags in [vec![Some(42), Some(42)], vec![None, None]] {
            let batches = [tags.clone(), tags]
                .into_iter()
                .enumerate()
                .map(|(batch, tags)| {
                    RecordBatch::try_new(
                        schema.clone(),
                        vec![
                            Arc::new(UInt64Array::from(tags)),
                            Arc::new(Int64Array::from_iter_values(
                                (0..2).map(|row| (batch * 10 + row) as i64),
                            )),
                        ],
                    )
                    .unwrap()
                })
                .collect();
            assert_concat_matches_reference(schema.clone(), batches, &[0]);
        }
    }

    #[test]
    fn test_concat_series_batches_leading_empty_and_sliced_batches() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("tag", DataType::Utf8, true),
            Field::new("value", DataType::Int64, true),
        ]));
        let make_sliced_batch = |values| {
            RecordBatch::try_new(
                schema.clone(),
                vec![
                    Arc::new(StringArray::from(vec!["discard", "tag", "tag"])),
                    Arc::new(Int64Array::from(values)),
                ],
            )
            .unwrap()
            .slice(1, 2)
        };
        let batches = vec![
            RecordBatch::new_empty(schema.clone()),
            make_sliced_batch(vec![0, 1, 2]),
            make_sliced_batch(vec![3, 4, 5]),
        ];
        assert_concat_matches_reference(schema, batches, &[0]);
    }

    #[test]
    fn test_concat_series_batches_interleaved_tags_and_schema_metadata() {
        let schema = Arc::new(Schema::new_with_metadata(
            vec![
                Field::new("value_before", DataType::Int64, true),
                Field::new("host", DataType::Utf8, true),
                Field::new("value_between", DataType::UInt64, true),
                Field::new("path", DataType::Utf8, true),
                Field::new("value_after", DataType::Int32, true),
            ],
            HashMap::from([("source".to_string(), "concat test".to_string())]),
        ));
        let batches = [
            (
                vec![Some(1), None],
                vec![Some(10), Some(11)],
                vec![100, 101],
            ),
            (
                vec![Some(2), Some(3)],
                vec![Some(12), Some(13)],
                vec![102, 103],
            ),
        ]
        .into_iter()
        .map(|(before, between, after)| {
            RecordBatch::try_new(
                schema.clone(),
                vec![
                    Arc::new(Int64Array::from(before)),
                    Arc::new(StringArray::from(vec!["host-a", "host-a"])),
                    Arc::new(UInt64Array::from(between)),
                    Arc::new(StringArray::from(vec!["/metrics", "/metrics"])),
                    Arc::new(Int32Array::from(after)),
                ],
            )
            .unwrap()
        })
        .collect();
        assert_concat_matches_reference(schema, batches, &[1, 3]);
    }

    #[test]
    fn test_concat_series_batches_fallbacks() {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Int64,
            true,
        )]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![Arc::new(Int64Array::from(vec![Some(1), None]))],
        )
        .unwrap();
        assert_concat_matches_reference(schema.clone(), vec![batch.clone()], &[0]);
        assert_concat_matches_reference(schema.clone(), vec![batch.clone(), batch], &[]);

        let empty_schema = Arc::new(Schema::empty());
        let empty_batch = RecordBatch::try_new_with_options(
            empty_schema.clone(),
            vec![],
            &datafusion::arrow::record_batch::RecordBatchOptions::new().with_row_count(Some(2)),
        )
        .unwrap();
        assert_concat_matches_reference(
            empty_schema.clone(),
            vec![empty_batch.clone(), empty_batch],
            &[],
        );

        let zero_rows = RecordBatch::new_empty(schema.clone());
        assert_concat_matches_reference(schema, vec![zero_rows.clone(), zero_rows], &[0]);
    }

    #[test]
    fn test_dictionary_tag_child_null_comparison() {
        let dictionary: ArrayRef = Arc::new(DictionaryArray::<UInt32Type>::new(
            UInt32Array::from(vec![Some(0), None, Some(1)]),
            Arc::new(StringArray::from(vec![None, Some("")])),
        ));
        let tags = RawTagColumn::try_new(&dictionary).unwrap();

        assert!(tags.equal_at(0, &tags, 1));
        assert!(!tags.equal_at(0, &tags, 2));
        assert!(!tags.equal_at(1, &tags, 2));
    }

    #[test]
    fn test_dictionary_tag_with_non_uint32_keys() {
        let dictionary: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::new(
            Int32Array::from(vec![0, 1]),
            Arc::new(StringArray::from(vec!["host-a", "host-b"])),
        ));

        let tags = RawTagColumn::try_new(&dictionary).unwrap();

        assert!(tags.equal_at(0, &tags, 0));
        assert!(!tags.equal_at(0, &tags, 1));
    }

    fn prepare_test_data() -> DataSourceExec {
        let schema = Arc::new(Schema::new(vec![
            Field::new("host", DataType::Utf8, true),
            Field::new("path", DataType::Utf8, true),
            Field::new(
                "time_index",
                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
                false,
            ),
        ]));

        let path_column_1 = Arc::new(StringArray::from(vec![
            "foo", "foo", "foo", "bar", "bar", "bar", "bar", "bar", "bar", "bla", "bla", "bla",
        ])) as _;
        let host_column_1 = Arc::new(StringArray::from(vec![
            "000", "000", "001", "002", "002", "002", "002", "002", "003", "005", "005", "005",
        ])) as _;
        let time_index_column_1 = Arc::new(
            datafusion::arrow::array::TimestampMillisecondArray::from(vec![
                1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000, 9000, 10000, 11000, 12000,
            ]),
        ) as _;

        let path_column_2 = Arc::new(StringArray::from(vec!["bla", "bla", "bla"])) as _;
        let host_column_2 = Arc::new(StringArray::from(vec!["005", "005", "005"])) as _;
        let time_index_column_2 = Arc::new(
            datafusion::arrow::array::TimestampMillisecondArray::from(vec![13000, 14000, 15000]),
        ) as _;

        let path_column_3 = Arc::new(StringArray::from(vec![
            "bla", "🥺", "🥺", "🥺", "🥺", "🥺", "🫠", "🫠",
        ])) as _;
        let host_column_3 = Arc::new(StringArray::from(vec![
            "005", "001", "001", "001", "001", "001", "001", "001",
        ])) as _;
        let time_index_column_3 =
            Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                vec![16000, 17000, 18000, 19000, 20000, 21000, 22000, 23000],
            )) as _;

        let data_1 = RecordBatch::try_new(
            schema.clone(),
            vec![path_column_1, host_column_1, time_index_column_1],
        )
        .unwrap();
        let data_2 = RecordBatch::try_new(
            schema.clone(),
            vec![path_column_2, host_column_2, time_index_column_2],
        )
        .unwrap();
        let data_3 = RecordBatch::try_new(
            schema.clone(),
            vec![path_column_3, host_column_3, time_index_column_3],
        )
        .unwrap();

        DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![data_1, data_2, data_3]], schema, None).unwrap(),
        ))
    }

    #[test]
    fn pruning_should_keep_tags_and_time_index_columns_for_exec() {
        let df_schema = prepare_test_data().schema().to_dfschema_ref().unwrap();
        let input = LogicalPlan::EmptyRelation(EmptyRelation {
            produce_one_row: false,
            schema: df_schema,
        });
        let plan = SeriesDivide::new(
            vec!["host".to_string(), "path".to_string()],
            "time_index".to_string(),
            input,
        );

        // Simulate a parent projection requesting only the `host` column.
        let output_columns = [0usize];
        let required = plan.necessary_children_exprs(&output_columns).unwrap();
        let required = &required[0];
        assert_eq!(required.as_slice(), &[0, 1, 2]);
    }

    #[tokio::test]
    async fn overall_data() {
        let memory_exec = Arc::new(prepare_test_data());
        let divide_exec = Arc::new(SeriesDivideExec {
            tag_columns: vec!["host".to_string(), "path".to_string()],
            time_index_column: "time_index".to_string(),
            input: memory_exec,
            metric: ExecutionPlanMetricsSet::new(),
        });
        let session_context = SessionContext::default();
        let result = datafusion::physical_plan::collect(divide_exec, session_context.task_ctx())
            .await
            .unwrap();
        let result_literal = datatypes::arrow::util::pretty::pretty_format_batches(&result)
            .unwrap()
            .to_string();

        let expected = String::from(
            "+------+------+---------------------+\
            \n| host | path | time_index          |\
            \n+------+------+---------------------+\
            \n| foo  | 000  | 1970-01-01T00:00:01 |\
            \n| foo  | 000  | 1970-01-01T00:00:02 |\
            \n| foo  | 001  | 1970-01-01T00:00:03 |\
            \n| bar  | 002  | 1970-01-01T00:00:04 |\
            \n| bar  | 002  | 1970-01-01T00:00:05 |\
            \n| bar  | 002  | 1970-01-01T00:00:06 |\
            \n| bar  | 002  | 1970-01-01T00:00:07 |\
            \n| bar  | 002  | 1970-01-01T00:00:08 |\
            \n| bar  | 003  | 1970-01-01T00:00:09 |\
            \n| bla  | 005  | 1970-01-01T00:00:10 |\
            \n| bla  | 005  | 1970-01-01T00:00:11 |\
            \n| bla  | 005  | 1970-01-01T00:00:12 |\
            \n| bla  | 005  | 1970-01-01T00:00:13 |\
            \n| bla  | 005  | 1970-01-01T00:00:14 |\
            \n| bla  | 005  | 1970-01-01T00:00:15 |\
            \n| bla  | 005  | 1970-01-01T00:00:16 |\
            \n| 🥺   | 001  | 1970-01-01T00:00:17 |\
            \n| 🥺   | 001  | 1970-01-01T00:00:18 |\
            \n| 🥺   | 001  | 1970-01-01T00:00:19 |\
            \n| 🥺   | 001  | 1970-01-01T00:00:20 |\
            \n| 🥺   | 001  | 1970-01-01T00:00:21 |\
            \n| 🫠   | 001  | 1970-01-01T00:00:22 |\
            \n| 🫠   | 001  | 1970-01-01T00:00:23 |\
            \n+------+------+---------------------+",
        );
        assert_eq!(result_literal, expected);
    }

    #[tokio::test]
    async fn per_batch_data() {
        let memory_exec = Arc::new(prepare_test_data());
        let divide_exec = Arc::new(SeriesDivideExec {
            tag_columns: vec!["host".to_string(), "path".to_string()],
            time_index_column: "time_index".to_string(),
            input: memory_exec,
            metric: ExecutionPlanMetricsSet::new(),
        });
        let mut divide_stream = divide_exec
            .execute(0, SessionContext::default().task_ctx())
            .unwrap();

        let mut expectations = vec![
            String::from(
                "+------+------+---------------------+\
                \n| host | path | time_index          |\
                \n+------+------+---------------------+\
                \n| foo  | 000  | 1970-01-01T00:00:01 |\
                \n| foo  | 000  | 1970-01-01T00:00:02 |\
                \n+------+------+---------------------+",
            ),
            String::from(
                "+------+------+---------------------+\
                \n| host | path | time_index          |\
                \n+------+------+---------------------+\
                \n| foo  | 001  | 1970-01-01T00:00:03 |\
                \n+------+------+---------------------+",
            ),
            String::from(
                "+------+------+---------------------+\
                \n| host | path | time_index          |\
                \n+------+------+---------------------+\
                \n| bar  | 002  | 1970-01-01T00:00:04 |\
                \n| bar  | 002  | 1970-01-01T00:00:05 |\
                \n| bar  | 002  | 1970-01-01T00:00:06 |\
                \n| bar  | 002  | 1970-01-01T00:00:07 |\
                \n| bar  | 002  | 1970-01-01T00:00:08 |\
                \n+------+------+---------------------+",
            ),
            String::from(
                "+------+------+---------------------+\
                \n| host | path | time_index          |\
                \n+------+------+---------------------+\
                \n| bar  | 003  | 1970-01-01T00:00:09 |\
                \n+------+------+---------------------+",
            ),
            String::from(
                "+------+------+---------------------+\
                \n| host | path | time_index          |\
                \n+------+------+---------------------+\
                \n| bla  | 005  | 1970-01-01T00:00:10 |\
                \n| bla  | 005  | 1970-01-01T00:00:11 |\
                \n| bla  | 005  | 1970-01-01T00:00:12 |\
                \n| bla  | 005  | 1970-01-01T00:00:13 |\
                \n| bla  | 005  | 1970-01-01T00:00:14 |\
                \n| bla  | 005  | 1970-01-01T00:00:15 |\
                \n| bla  | 005  | 1970-01-01T00:00:16 |\
                \n+------+------+---------------------+",
            ),
            String::from(
                "+------+------+---------------------+\
                \n| host | path | time_index          |\
                \n+------+------+---------------------+\
                \n| 🥺   | 001  | 1970-01-01T00:00:17 |\
                \n| 🥺   | 001  | 1970-01-01T00:00:18 |\
                \n| 🥺   | 001  | 1970-01-01T00:00:19 |\
                \n| 🥺   | 001  | 1970-01-01T00:00:20 |\
                \n| 🥺   | 001  | 1970-01-01T00:00:21 |\
                \n+------+------+---------------------+",
            ),
            String::from(
                "+------+------+---------------------+\
                \n| host | path | time_index          |\
                \n+------+------+---------------------+\
                \n| 🫠   | 001  | 1970-01-01T00:00:22 |\
                \n| 🫠   | 001  | 1970-01-01T00:00:23 |\
                \n+------+------+---------------------+",
            ),
        ];
        expectations.reverse();

        while let Some(batch) = divide_stream.next().await {
            let formatted =
                datatypes::arrow::util::pretty::pretty_format_batches(&[batch.unwrap()])
                    .unwrap()
                    .to_string();
            let expected = expectations.pop().unwrap();
            assert_eq!(formatted, expected);
        }
    }

    #[tokio::test]
    async fn test_all_batches_same_combination() {
        // Create a schema with host and path columns, same as prepare_test_data
        let schema = Arc::new(Schema::new(vec![
            Field::new("host", DataType::Utf8, true),
            Field::new("path", DataType::Utf8, true),
            Field::new(
                "time_index",
                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
                false,
            ),
        ]));

        // Create batches with three different combinations
        // Each batch contains only one combination
        // Batches with the same combination are adjacent

        // First combination: "server1", "/var/log"
        let batch1 = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(StringArray::from(vec!["server1", "server1", "server1"])) as _,
                Arc::new(StringArray::from(vec!["/var/log", "/var/log", "/var/log"])) as _,
                Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                    vec![1000, 2000, 3000],
                )) as _,
            ],
        )
        .unwrap();

        let batch2 = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(StringArray::from(vec!["server1", "server1"])) as _,
                Arc::new(StringArray::from(vec!["/var/log", "/var/log"])) as _,
                Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                    vec![4000, 5000],
                )) as _,
            ],
        )
        .unwrap();

        // Second combination: "server2", "/var/data"
        let batch3 = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(StringArray::from(vec!["server2", "server2", "server2"])) as _,
                Arc::new(StringArray::from(vec![
                    "/var/data",
                    "/var/data",
                    "/var/data",
                ])) as _,
                Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                    vec![6000, 7000, 8000],
                )) as _,
            ],
        )
        .unwrap();

        let batch4 = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(StringArray::from(vec!["server2"])) as _,
                Arc::new(StringArray::from(vec!["/var/data"])) as _,
                Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                    vec![9000],
                )) as _,
            ],
        )
        .unwrap();

        // Third combination: "server3", "/opt/logs"
        let batch5 = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(StringArray::from(vec!["server3", "server3"])) as _,
                Arc::new(StringArray::from(vec!["/opt/logs", "/opt/logs"])) as _,
                Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                    vec![10000, 11000],
                )) as _,
            ],
        )
        .unwrap();

        let batch6 = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(StringArray::from(vec!["server3", "server3", "server3"])) as _,
                Arc::new(StringArray::from(vec![
                    "/opt/logs",
                    "/opt/logs",
                    "/opt/logs",
                ])) as _,
                Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                    vec![12000, 13000, 14000],
                )) as _,
            ],
        )
        .unwrap();

        // Create MemoryExec with these batches, keeping same combinations adjacent
        let memory_exec = DataSourceExec::from_data_source(
            MemorySourceConfig::try_new(
                &[vec![batch1, batch2, batch3, batch4, batch5, batch6]],
                schema.clone(),
                None,
            )
            .unwrap(),
        );

        // Create SeriesDivideExec
        let divide_exec = Arc::new(SeriesDivideExec {
            tag_columns: vec!["host".to_string(), "path".to_string()],
            time_index_column: "time_index".to_string(),
            input: memory_exec,
            metric: ExecutionPlanMetricsSet::new(),
        });

        // Execute the division
        let session_context = SessionContext::default();
        let result =
            datafusion::physical_plan::collect(divide_exec.clone(), session_context.task_ctx())
                .await
                .unwrap();

        // Verify that we got 3 batches (one for each combination)
        assert_eq!(result.len(), 3);

        // First batch should have 5 rows (3 + 2 from the "server1" combination)
        assert_eq!(result[0].num_rows(), 5);

        // Second batch should have 4 rows (3 + 1 from the "server2" combination)
        assert_eq!(result[1].num_rows(), 4);

        // Third batch should have 5 rows (2 + 3 from the "server3" combination)
        assert_eq!(result[2].num_rows(), 5);

        // Verify values in first batch (server1, /var/log)
        let host_array1 = result[0]
            .column(0)
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap();
        let path_array1 = result[0]
            .column(1)
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap();
        let time_index_array1 = result[0]
            .column(2)
            .as_any()
            .downcast_ref::<datafusion::arrow::array::TimestampMillisecondArray>()
            .unwrap();

        for i in 0..5 {
            assert_eq!(host_array1.value(i), "server1");
            assert_eq!(path_array1.value(i), "/var/log");
            assert_eq!(time_index_array1.value(i), 1000 + (i as i64) * 1000);
        }

        // Verify values in second batch (server2, /var/data)
        let host_array2 = result[1]
            .column(0)
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap();
        let path_array2 = result[1]
            .column(1)
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap();
        let time_index_array2 = result[1]
            .column(2)
            .as_any()
            .downcast_ref::<datafusion::arrow::array::TimestampMillisecondArray>()
            .unwrap();

        for i in 0..4 {
            assert_eq!(host_array2.value(i), "server2");
            assert_eq!(path_array2.value(i), "/var/data");
            assert_eq!(time_index_array2.value(i), 6000 + (i as i64) * 1000);
        }

        // Verify values in third batch (server3, /opt/logs)
        let host_array3 = result[2]
            .column(0)
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap();
        let path_array3 = result[2]
            .column(1)
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap();
        let time_index_array3 = result[2]
            .column(2)
            .as_any()
            .downcast_ref::<datafusion::arrow::array::TimestampMillisecondArray>()
            .unwrap();

        for i in 0..5 {
            assert_eq!(host_array3.value(i), "server3");
            assert_eq!(path_array3.value(i), "/opt/logs");
            assert_eq!(time_index_array3.value(i), 10000 + (i as i64) * 1000);
        }

        // Also verify streaming behavior
        let mut divide_stream = divide_exec
            .execute(0, SessionContext::default().task_ctx())
            .unwrap();

        // Should produce three batches, one for each combination
        let batch1 = divide_stream.next().await.unwrap().unwrap();
        assert_eq!(batch1.num_rows(), 5); // server1 combination

        let batch2 = divide_stream.next().await.unwrap().unwrap();
        assert_eq!(batch2.num_rows(), 4); // server2 combination

        let batch3 = divide_stream.next().await.unwrap().unwrap();
        assert_eq!(batch3.num_rows(), 5); // server3 combination

        // No more batches should be produced
        assert!(divide_stream.next().await.is_none());
    }

    #[tokio::test]
    async fn test_dictionary_tags_across_batches_and_eof() {
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "tag",
                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
                false,
            ),
            Field::new("value", DataType::Int64, false),
            Field::new(
                "time_index",
                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
                false,
            ),
        ]));
        let make_batch = |values: Vec<&str>, keys: Vec<u32>, payload: Vec<i64>| {
            RecordBatch::try_new(
                schema.clone(),
                vec![
                    Arc::new(DictionaryArray::<UInt32Type>::new(
                        UInt32Array::from(keys),
                        Arc::new(StringArray::from(values)),
                    )),
                    Arc::new(Int64Array::from(payload.clone())),
                    Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                        payload,
                    )),
                ],
            )
            .unwrap()
        };
        let memory_exec: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(
                &[vec![
                    make_batch(vec!["b", "a", "unused-a"], vec![1, 1, 0], vec![1, 2, 3]),
                    make_batch(vec!["c", "b", "unused-b"], vec![1, 1, 0], vec![4, 5, 6]),
                    make_batch(vec!["unused-c", "c"], vec![1, 1], vec![7, 8]),
                ]],
                schema,
                None,
            )
            .unwrap(),
        )));
        let divide_exec = Arc::new(SeriesDivideExec {
            tag_columns: vec!["tag".to_string()],
            time_index_column: "time_index".to_string(),
            input: memory_exec,
            metric: ExecutionPlanMetricsSet::new(),
        });
        let mut stream = divide_exec
            .execute(0, SessionContext::default().task_ctx())
            .unwrap();

        for (expected_tag, expected_payload, concatenated) in [
            ("a", vec![1, 2], false),
            ("b", vec![3, 4, 5], true),
            ("c", vec![6, 7, 8], true),
        ] {
            let batch = stream.next().await.unwrap().unwrap();
            assert_eq!(batch.num_rows(), expected_payload.len());
            assert_eq!(
                (0..batch.num_rows())
                    .map(|row| string_array_value_at_index(batch.column(0), row).unwrap())
                    .collect::<Vec<_>>(),
                vec![expected_tag; expected_payload.len()]
            );
            assert_eq!(
                batch
                    .column(1)
                    .as_any()
                    .downcast_ref::<Int64Array>()
                    .unwrap()
                    .iter()
                    .flatten()
                    .collect::<Vec<_>>(),
                expected_payload
            );

            if concatenated {
                let tag = batch
                    .column(0)
                    .as_any()
                    .downcast_ref::<DictionaryArray<UInt32Type>>()
                    .unwrap();
                let values = tag.values().as_any().downcast_ref::<StringArray>().unwrap();
                assert_eq!(values.len(), 1, "tag {expected_tag}");
                assert_eq!(values.value(0), expected_tag);
                assert!(tag.keys().iter().all(|key| key == Some(0)));
            }
        }
        assert!(stream.next().await.is_none());
    }

    #[tokio::test]
    async fn test_string_tag_column_types() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("tag_large", DataType::LargeUtf8, false),
            Field::new("tag_view", DataType::Utf8View, false),
            Field::new(
                "time_index",
                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
                false,
            ),
        ]));

        let batch1 = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(LargeStringArray::from(vec!["a", "a", "a", "a"])),
                Arc::new(StringViewArray::from(vec!["x", "x", "y", "y"])),
                Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                    vec![1000, 2000, 1000, 2000],
                )),
            ],
        )
        .unwrap();

        let batch2 = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(LargeStringArray::from(vec!["b", "b"])),
                Arc::new(StringViewArray::from(vec!["x", "x"])),
                Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                    vec![1000, 2000],
                )),
            ],
        )
        .unwrap();

        let memory_exec: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![batch1, batch2]], schema.clone(), None).unwrap(),
        )));

        let divide_exec = Arc::new(SeriesDivideExec {
            tag_columns: vec!["tag_large".to_string(), "tag_view".to_string()],
            time_index_column: "time_index".to_string(),
            input: memory_exec,
            metric: ExecutionPlanMetricsSet::new(),
        });

        let session_context = SessionContext::default();
        let result = datafusion::physical_plan::collect(divide_exec, session_context.task_ctx())
            .await
            .unwrap();

        assert_eq!(result.len(), 3);
        for ((expected_large, expected_view), batch) in [("a", "x"), ("a", "y"), ("b", "x")]
            .into_iter()
            .zip(result.iter())
        {
            assert_eq!(batch.num_rows(), 2);

            let tag_large_array = batch
                .column(0)
                .as_any()
                .downcast_ref::<LargeStringArray>()
                .unwrap();
            let tag_view_array = batch
                .column(1)
                .as_any()
                .downcast_ref::<StringViewArray>()
                .unwrap();

            for row in 0..batch.num_rows() {
                assert_eq!(tag_large_array.value(row), expected_large);
                assert_eq!(tag_view_array.value(row), expected_view);
            }
        }
    }

    #[tokio::test]
    async fn test_u64_tag_column() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("tsid", DataType::UInt64, false),
            Field::new(
                "time_index",
                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
                false,
            ),
        ]));

        let batch1 = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(UInt64Array::from(vec![1, 1, 2, 2])),
                Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                    vec![1000, 2000, 1000, 2000],
                )),
            ],
        )
        .unwrap();

        let batch2 = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(UInt64Array::from(vec![3, 3])),
                Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                    vec![1000, 2000],
                )),
            ],
        )
        .unwrap();

        let memory_exec: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(&[vec![batch1, batch2]], schema.clone(), None).unwrap(),
        )));

        let divide_exec = Arc::new(SeriesDivideExec {
            tag_columns: vec!["tsid".to_string()],
            time_index_column: "time_index".to_string(),
            input: memory_exec,
            metric: ExecutionPlanMetricsSet::new(),
        });

        let session_context = SessionContext::default();
        let result = datafusion::physical_plan::collect(divide_exec, session_context.task_ctx())
            .await
            .unwrap();

        assert_eq!(result.len(), 3);
        for (expected_tsid, batch) in [1u64, 2u64, 3u64].into_iter().zip(result.iter()) {
            assert_eq!(batch.num_rows(), 2);
            let tsid_array = batch
                .column(0)
                .as_any()
                .downcast_ref::<UInt64Array>()
                .unwrap();
            assert!(tsid_array.iter().all(|v| v == Some(expected_tsid)));
        }
    }

    /// The input stream may legally start with a batch that has zero rows.
    /// `SeriesDivideStream` must skip it and keep working on the real rows that
    /// follow (here: two batches of the very same tagged series).
    #[tokio::test]
    async fn test_stream_leading_empty_batch_then_same_tags() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("host", DataType::Utf8, false),
            Field::new("path", DataType::Utf8, false),
            Field::new(
                "time_index",
                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
                false,
            ),
        ]));

        let make_batch = |hosts: Vec<&str>, paths: Vec<&str>, times: Vec<i64>| {
            RecordBatch::try_new(
                schema.clone(),
                vec![
                    Arc::new(StringArray::from(hosts)),
                    Arc::new(StringArray::from(paths)),
                    Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                        times,
                    )),
                ],
            )
            .unwrap()
        };

        let memory_exec: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(
                &[vec![
                    RecordBatch::new_empty(schema.clone()),
                    make_batch(
                        vec!["server1", "server1"],
                        vec!["/var/log", "/var/log"],
                        vec![1000, 2000],
                    ),
                    make_batch(vec!["server1"], vec!["/var/log"], vec![3000]),
                ]],
                schema.clone(),
                None,
            )
            .unwrap(),
        )));

        let divide_exec = Arc::new(SeriesDivideExec {
            tag_columns: vec!["host".to_string(), "path".to_string()],
            time_index_column: "time_index".to_string(),
            input: memory_exec,
            metric: ExecutionPlanMetricsSet::new(),
        });

        let mut stream = divide_exec
            .execute(0, SessionContext::default().task_ctx())
            .unwrap();

        let batch = stream.next().await.unwrap().unwrap();
        assert_eq!(batch.num_rows(), 3);
        let expected = [
            ("server1", "/var/log", 1000i64),
            ("server1", "/var/log", 2000),
            ("server1", "/var/log", 3000),
        ];
        for (row, (host, path, time)) in expected.into_iter().enumerate() {
            assert_eq!(
                string_array_value_at_index(batch.column(0), row),
                Some(host)
            );
            assert_eq!(
                string_array_value_at_index(batch.column(1), row),
                Some(path)
            );
            assert_eq!(
                batch
                    .column(2)
                    .as_any()
                    .downcast_ref::<datafusion::arrow::array::TimestampMillisecondArray>()
                    .unwrap()
                    .value(row),
                time
            );
        }
        assert!(stream.next().await.is_none());
    }

    /// Several leading empty batches (and one empty batch in the middle of a
    /// series) must not disturb series division: the two series are still
    /// emitted separately, including the last one flushed at end of input.
    #[tokio::test]
    async fn test_stream_repeated_empty_batches_and_two_series() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("host", DataType::Utf8, false),
            Field::new("value", DataType::Int64, false),
            Field::new(
                "time_index",
                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
                false,
            ),
        ]));

        let make_batch = |hosts: Vec<&str>, values: Vec<i64>| {
            RecordBatch::try_new(
                schema.clone(),
                vec![
                    Arc::new(StringArray::from(hosts)),
                    Arc::new(Int64Array::from(values.clone())),
                    Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
                        values.iter().map(|v| v * 1000).collect::<Vec<_>>(),
                    )),
                ],
            )
            .unwrap()
        };

        let memory_exec: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(
                &[vec![
                    RecordBatch::new_empty(schema.clone()),
                    RecordBatch::new_empty(schema.clone()),
                    make_batch(vec!["server1", "server1"], vec![1, 2]),
                    RecordBatch::new_empty(schema.clone()),
                    make_batch(vec!["server1"], vec![3]),
                    make_batch(vec!["server2", "server2"], vec![4, 5]),
                ]],
                schema.clone(),
                None,
            )
            .unwrap(),
        )));

        let divide_exec = Arc::new(SeriesDivideExec {
            tag_columns: vec!["host".to_string()],
            time_index_column: "time_index".to_string(),
            input: memory_exec,
            metric: ExecutionPlanMetricsSet::new(),
        });

        let mut stream = divide_exec
            .execute(0, SessionContext::default().task_ctx())
            .unwrap();

        for (expected_host, expected_values) in
            [("server1", vec![1i64, 2, 3]), ("server2", vec![4, 5])]
        {
            let batch = stream.next().await.unwrap().unwrap();
            assert_eq!(batch.num_rows(), expected_values.len());
            for (row, value) in expected_values.into_iter().enumerate() {
                assert_eq!(
                    string_array_value_at_index(batch.column(0), row),
                    Some(expected_host)
                );
                assert_eq!(
                    batch
                        .column(1)
                        .as_any()
                        .downcast_ref::<Int64Array>()
                        .unwrap()
                        .value(row),
                    value
                );
            }
        }
        assert!(stream.next().await.is_none());
    }

    /// A stream that only ever yields zero-row batches must terminate with no
    /// batches instead of failing.
    #[tokio::test]
    async fn test_stream_all_empty_batches_yields_nothing() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("host", DataType::Utf8, false),
            Field::new("value", DataType::Int64, false),
            Field::new(
                "time_index",
                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
                false,
            ),
        ]));

        let memory_exec: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
            MemorySourceConfig::try_new(
                &[vec![
                    RecordBatch::new_empty(schema.clone()),
                    RecordBatch::new_empty(schema.clone()),
                    RecordBatch::new_empty(schema.clone()),
                ]],
                schema.clone(),
                None,
            )
            .unwrap(),
        )));

        let divide_exec = Arc::new(SeriesDivideExec {
            tag_columns: vec!["host".to_string()],
            time_index_column: "time_index".to_string(),
            input: memory_exec,
            metric: ExecutionPlanMetricsSet::new(),
        });

        let mut stream = divide_exec
            .execute(0, SessionContext::default().task_ctx())
            .unwrap();
        assert!(stream.next().await.is_none());

        let collected =
            datafusion::physical_plan::collect(divide_exec, SessionContext::default().task_ctx())
                .await
                .unwrap();
        assert!(collected.is_empty());
    }
}