datum-sql 0.10.3

use std::sync::Arc;
use std::thread;
use std::time::{Duration, Instant};

use arrow::array::{Array, Int64Array, TimestampNanosecondArray, UInt64Array};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
use arrow::record_batch::RecordBatch;
use datum::{NotUsed, Source, StreamResult};
use datum_sql::{
    ChangeOp, ChangelogBatch, ContinuousQueryHandle, DatumSqlContext, EventTimeConfig, SqlEvent,
};
#[cfg(feature = "mq")]
use datum_sql::{
    CommittableRecordBatch, SourceCommit, SqlSourcePosition,
    connect::mq::{KafkaPartitionOffset, KafkaSourcePosition, KafkaTopicPartition},
};

#[tokio::test]
async fn tumble_windows_emit_once_on_watermark_with_boundary_rows() {
    let schema = bids_schema();
    let batch = bid_batch(
        Arc::clone(&schema),
        &[1, 1, 1],
        &[7, 11, 13],
        &[0, 9_999_999_999, 10_000_000_000],
    );
    let context = context_with_bids(schema, Source::from_iter([batch]));

    let output = collect_data_events(
        context
            .streaming_source(
                "SELECT date_bin(INTERVAL '10 seconds', event_ts) AS window_start, \
                        category, SUM(price) AS price_sum, COUNT(*) AS bid_count \
                 FROM bids \
                 GROUP BY date_bin(INTERVAL '10 seconds', event_ts), category",
            )
            .await
            .expect("query lowers"),
    );

    assert_eq!(output.len(), 1);
    assert_eq!(timestamp_values(&output[0], 0), vec![0]);
    assert_eq!(int_values(&output[0], 1), vec![1]);
    assert_eq!(int_values(&output[0], 2), vec![18]);
    assert_eq!(int_values(&output[0], 3), vec![2]);
}

#[tokio::test]
async fn hop_window_assigns_each_row_to_multiple_bins() {
    let schema = bids_schema();
    let batch = bid_batch(
        Arc::clone(&schema),
        &[42, 42],
        &[1, 1],
        &[8_000_000_000, 20_000_000_000],
    );
    let context = context_with_bids(schema, Source::from_iter([batch]));

    let output = collect_data_events(
        context
            .streaming_source(
                "SELECT hop(INTERVAL '2 seconds', INTERVAL '10 seconds', event_ts) AS window_start, \
                        COUNT(*) AS bid_count \
                 FROM bids \
                 GROUP BY hop(INTERVAL '2 seconds', INTERVAL '10 seconds', event_ts)",
            )
            .await
            .expect("query lowers"),
    );

    assert_eq!(output.len(), 1);
    assert_eq!(
        timestamp_values(&output[0], 0),
        vec![
            0,
            2_000_000_000,
            4_000_000_000,
            6_000_000_000,
            8_000_000_000
        ]
    );
    assert_eq!(int_values(&output[0], 1), vec![1, 1, 1, 1, 1]);
}

#[tokio::test]
async fn tumble_projection_can_emit_window_end_epoch() {
    let schema = bids_schema();
    let batch = bid_batch(
        Arc::clone(&schema),
        &[1, 1],
        &[5, 11],
        &[1_000_000_000, 20_000_000_000],
    );
    let context = context_with_bids(schema, Source::from_iter([batch]));

    let output = collect_data_events(
        context
            .streaming_source(
                "SELECT CAST(date_part('epoch', date_bin(INTERVAL '10 seconds', event_ts) \
                            + INTERVAL '10 seconds') AS BIGINT) AS window_end_epoch, \
                        SUM(price) AS price_sum \
                 FROM bids \
                 GROUP BY date_bin(INTERVAL '10 seconds', event_ts)",
            )
            .await
            .expect("query lowers"),
    );

    assert_eq!(output.len(), 1);
    assert_eq!(int_values(&output[0], 0), vec![10]);
    assert_eq!(int_values(&output[0], 1), vec![5]);
}

#[tokio::test]
async fn bounded_out_of_order_input_accumulates_until_watermark_passes_window() {
    let schema = bids_schema();
    let batches = vec![
        bid_batch(Arc::clone(&schema), &[1], &[5], &[10_000_000_000]),
        bid_batch(Arc::clone(&schema), &[1], &[7], &[2_000_000_000]),
        bid_batch(Arc::clone(&schema), &[1], &[11], &[20_000_000_000]),
    ];
    let context = DatumSqlContext::new();
    context
        .register_source_with_event_time(
            "bids",
            Arc::clone(&schema),
            Source::from_iter(batches),
            EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::from_secs(10)),
        )
        .expect("source registers");

    let output = collect_data_events(
        context
            .streaming_source(
                "SELECT date_bin(INTERVAL '10 seconds', event_ts) AS window_start, \
                        SUM(price) AS price_sum \
                 FROM bids \
                 GROUP BY date_bin(INTERVAL '10 seconds', event_ts)",
            )
            .await
            .expect("query lowers"),
    );

    assert_eq!(output.len(), 1);
    assert_eq!(timestamp_values(&output[0], 0), vec![0]);
    assert_eq!(int_values(&output[0], 1), vec![7]);
}

#[tokio::test]
async fn late_rows_behind_watermark_are_dropped_and_counted() {
    let schema = bids_schema();
    let batches = vec![
        bid_batch(Arc::clone(&schema), &[1], &[5], &[10_000_000_000]),
        bid_batch(Arc::clone(&schema), &[1], &[7], &[2_000_000_000]),
        bid_batch(Arc::clone(&schema), &[1], &[11], &[20_000_000_000]),
    ];
    let context = context_with_bids(Arc::clone(&schema), Source::from_iter(batches));
    let (source, metrics) = context
        .streaming_source_with_metrics(
            "SELECT date_bin(INTERVAL '10 seconds', event_ts) AS window_start, \
                    SUM(price) AS price_sum \
             FROM bids \
             GROUP BY date_bin(INTERVAL '10 seconds', event_ts)",
        )
        .await
        .expect("query lowers");

    let output = collect_data_events(source);

    assert_eq!(metrics.late_dropped_rows(), 1);
    assert_eq!(output.len(), 1);
    assert_eq!(timestamp_values(&output[0], 0), vec![10_000_000_000]);
    assert_eq!(int_values(&output[0], 1), vec![5]);
}

#[tokio::test]
async fn changelog_window_uses_retractable_accumulators() {
    let schema = bids_schema();
    let changes = ChangelogBatch::try_new(
        vec![ChangeOp::Insert, ChangeOp::Delete, ChangeOp::Insert],
        bid_batch(
            Arc::clone(&schema),
            &[1, 1, 2],
            &[5, 5, 1],
            &[1_000_000_000, 2_000_000_000, 10_000_000_000],
        ),
    )
    .expect("changelog builds");
    let context = DatumSqlContext::new();
    context
        .register_changelog_source_with_event_time(
            "bids",
            Arc::clone(&schema),
            Source::from_iter([changes]),
            EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::ZERO),
        )
        .expect("source registers");

    let output = collect_data_events(
        context
            .streaming_source(
                "SELECT date_bin(INTERVAL '10 seconds', event_ts) AS window_start, \
                        COUNT(*) AS bid_count \
                 FROM bids \
                 GROUP BY date_bin(INTERVAL '10 seconds', event_ts)",
            )
            .await
            .expect("query lowers"),
    );

    assert_eq!(output.len(), 1);
    assert_eq!(timestamp_values(&output[0], 0), vec![0]);
    assert_eq!(int_values(&output[0], 1), vec![0]);
}

#[tokio::test]
async fn grouped_window_updates_compose_with_retractions() {
    let schema = bids_schema();
    let changes = ChangelogBatch::try_new(
        vec![
            ChangeOp::Insert,
            ChangeOp::Insert,
            ChangeOp::Delete,
            ChangeOp::Insert,
            ChangeOp::Insert,
        ],
        bid_batch(
            Arc::clone(&schema),
            &[7, 7, 7, 7, 8],
            &[5, 7, 5, 2, 1],
            &[
                1_000_000_000,
                2_000_000_000,
                3_000_000_000,
                4_000_000_000,
                10_000_000_000,
            ],
        ),
    )
    .expect("changelog builds");
    let context = DatumSqlContext::new();
    context
        .register_changelog_source_with_event_time(
            "bids",
            Arc::clone(&schema),
            Source::from_iter([changes]),
            EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::ZERO),
        )
        .expect("source registers");

    let output = collect_data_events(
        context
            .streaming_source(
                "SELECT date_bin(INTERVAL '10 seconds', event_ts) AS window_start, \
                        category, SUM(price) AS price_sum, COUNT(*) AS bid_count \
                 FROM bids \
                 GROUP BY date_bin(INTERVAL '10 seconds', event_ts), category",
            )
            .await
            .expect("query lowers"),
    );

    assert_eq!(output.len(), 1);
    assert_eq!(timestamp_values(&output[0], 0), vec![0]);
    assert_eq!(int_values(&output[0], 1), vec![7]);
    assert_eq!(int_values(&output[0], 2), vec![9]);
    assert_eq!(int_values(&output[0], 3), vec![2]);
}

#[tokio::test]
async fn updating_input_rejects_unsupported_aggregate_shape() {
    let schema = bids_schema();
    let changes =
        ChangelogBatch::insert_only(bid_batch(Arc::clone(&schema), &[1], &[5], &[1_000_000_000]));
    let context = DatumSqlContext::new();
    context
        .register_changelog_source_with_event_time(
            "bids",
            Arc::clone(&schema),
            Source::from_iter([changes]),
            EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::ZERO),
        )
        .expect("source registers");

    let result = context
        .streaming_source(
            "SELECT date_bin(INTERVAL '10 seconds', event_ts) AS window_start, \
                    COUNT(DISTINCT price) AS distinct_prices \
             FROM bids \
             GROUP BY date_bin(INTERVAL '10 seconds', event_ts)",
        )
        .await;
    let error = match result {
        Ok(_) => panic!("unsupported updating aggregate should fail planning"),
        Err(error) => error,
    };

    assert!(
        error
            .to_string()
            .contains("does not lower updating DataFusion physical node AggregateExec"),
        "{error}"
    );
}

#[tokio::test]
async fn windowed_continuous_query_cancels_while_window_is_open() {
    let schema = bids_schema();
    let batch = bid_batch(Arc::clone(&schema), &[1], &[5], &[1_000_000_000]);
    let context = context_with_bids(schema, Source::repeat(batch));

    let mut handle = context
        .execute_streaming(
            "SELECT date_bin(INTERVAL '10 seconds', event_ts) AS window_start, \
                    SUM(price) AS price_sum \
             FROM bids \
             GROUP BY date_bin(INTERVAL '10 seconds', event_ts)",
        )
        .await
        .expect("query materializes");
    handle.cancel();

    let result = wait_for_completion(&mut handle, Duration::from_secs(2));
    assert!(matches!(result, Ok(NotUsed)));
}

#[cfg(feature = "mq")]
#[tokio::test]
async fn windowed_aggregations_min_reduce_kafka_partitions_before_late_drop() {
    let (q4_output, q4_late_drops) = collect_partitioned_bids(
        "SELECT date_bin(INTERVAL '10 seconds', event_ts) AS window_start, \
                category, SUM(price) AS price_sum, COUNT(*) AS bid_count \
         FROM bids \
         GROUP BY date_bin(INTERVAL '10 seconds', event_ts), category",
    )
    .await;
    assert_eq!(q4_late_drops, 0);
    assert_eq!(q4_output.len(), 1);
    assert_eq!(timestamp_values(&q4_output[0], 0), vec![0]);
    assert_eq!(int_values(&q4_output[0], 1), vec![1]);
    assert_eq!(int_values(&q4_output[0], 2), vec![7]);
    assert_eq!(int_values(&q4_output[0], 3), vec![1]);

    let (q5_output, q5_late_drops) = collect_partitioned_bids(
        "SELECT hop(INTERVAL '2 seconds', INTERVAL '10 seconds', event_ts) AS window_start, \
                COUNT(*) AS bid_count \
         FROM bids \
         GROUP BY hop(INTERVAL '2 seconds', INTERVAL '10 seconds', event_ts)",
    )
    .await;
    assert_eq!(q5_late_drops, 0);
    assert_eq!(q5_output.len(), 1);
    assert_eq!(
        timestamp_values(&q5_output[0], 0),
        vec![
            -8_000_000_000,
            -6_000_000_000,
            -4_000_000_000,
            -2_000_000_000,
            0
        ]
    );
    assert_eq!(int_values(&q5_output[0], 1), vec![1, 1, 1, 1, 1]);

    let (q7_output, q7_late_drops) = collect_partitioned_bids(
        "SELECT date_bin(INTERVAL '10 seconds', event_ts) AS window_start, \
                MAX(price) AS max_price \
         FROM bids \
         GROUP BY date_bin(INTERVAL '10 seconds', event_ts)",
    )
    .await;
    assert_eq!(q7_late_drops, 0);
    assert_eq!(q7_output.len(), 1);
    assert_eq!(timestamp_values(&q7_output[0], 0), vec![0]);
    assert_eq!(int_values(&q7_output[0], 1), vec![7]);
}

fn context_with_bids(source_schema: SchemaRef, source: Source<RecordBatch>) -> DatumSqlContext {
    let context = DatumSqlContext::new();
    context
        .register_source_with_event_time(
            "bids",
            source_schema,
            source,
            EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::ZERO),
        )
        .expect("source registers");
    context
}

#[cfg(feature = "mq")]
async fn collect_partitioned_bids(sql: &str) -> (Vec<RecordBatch>, u64) {
    let schema = bids_schema();
    let batches = vec![
        partitioned_committable_bid_batch(
            Arc::clone(&schema),
            &[9],
            &[90],
            &[20_000_000_000],
            1,
            &[0, 1],
        ),
        partitioned_committable_bid_batch(Arc::clone(&schema), &[1], &[7], &[0], 0, &[0, 1]),
        partitioned_committable_bid_batch(
            Arc::clone(&schema),
            &[2],
            &[11],
            &[20_000_000_000],
            0,
            &[0, 1],
        ),
    ];
    let context = DatumSqlContext::new();
    context
        .register_committable_source_with_event_time(
            "bids",
            schema,
            Source::from_iter(batches),
            EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::ZERO),
        )
        .expect("source registers");
    let (source, metrics) = context
        .streaming_source_with_metrics(sql)
        .await
        .expect("query lowers");
    (collect_data_events(source), metrics.late_dropped_rows())
}

fn collect_data_events(source: Source<SqlEvent<RecordBatch>>) -> Vec<RecordBatch> {
    source
        .run_collect()
        .expect("events collect")
        .into_iter()
        .filter_map(|event| match event {
            SqlEvent::Data(batch) => Some(batch),
            SqlEvent::Watermark(_) | SqlEvent::Barrier(_) => None,
        })
        .collect()
}

#[cfg(feature = "mq")]
fn partitioned_committable_bid_batch(
    schema: SchemaRef,
    categories: &[i64],
    prices: &[i64],
    timestamps_ns: &[i64],
    partition: i32,
    active_partitions: &[i32],
) -> CommittableRecordBatch {
    let batch = bid_batch(schema, categories, prices, timestamps_ns);
    let position = KafkaSourcePosition::from_offsets_with_row_partitions(
        [KafkaPartitionOffset {
            topic: "bids".to_owned(),
            partition,
            first_offset: 0,
            last_offset: categories.len() as i64 - 1,
        }],
        std::iter::repeat_n(partition, categories.len()),
        active_partitions
            .iter()
            .map(|partition| KafkaTopicPartition::new("bids", *partition)),
    );
    CommittableRecordBatch::new(
        batch,
        Some(SqlSourcePosition::Kafka(position)),
        0,
        SourceCommit::noop(),
    )
}

fn bids_schema() -> SchemaRef {
    Arc::new(Schema::new(vec![
        Field::new("category", DataType::Int64, false),
        Field::new("price", DataType::Int64, false),
        Field::new(
            "event_ts",
            DataType::Timestamp(TimeUnit::Nanosecond, None),
            false,
        ),
    ]))
}

fn bid_batch(
    schema: SchemaRef,
    categories: &[i64],
    prices: &[i64],
    timestamps_ns: &[i64],
) -> RecordBatch {
    RecordBatch::try_new(
        schema,
        vec![
            Arc::new(Int64Array::from(categories.to_vec())),
            Arc::new(Int64Array::from(prices.to_vec())),
            Arc::new(TimestampNanosecondArray::from(timestamps_ns.to_vec())),
        ],
    )
    .expect("batch builds")
}

fn timestamp_values(batch: &RecordBatch, column: usize) -> Vec<i64> {
    let array = batch
        .column(column)
        .as_any()
        .downcast_ref::<TimestampNanosecondArray>()
        .expect("column is TimestampNanosecond");
    (0..array.len()).map(|row| array.value(row)).collect()
}

fn int_values(batch: &RecordBatch, column: usize) -> Vec<i64> {
    if let Some(array) = batch.column(column).as_any().downcast_ref::<Int64Array>() {
        return (0..array.len()).map(|row| array.value(row)).collect();
    }
    if let Some(array) = batch.column(column).as_any().downcast_ref::<UInt64Array>() {
        return (0..array.len())
            .map(|row| array.value(row) as i64)
            .collect();
    }
    panic!("column {column} is not an integer array");
}

fn wait_for_completion(
    handle: &mut ContinuousQueryHandle,
    timeout: Duration,
) -> StreamResult<NotUsed> {
    let deadline = Instant::now() + timeout;
    loop {
        if let Some(result) = handle.try_wait() {
            return result;
        }
        assert!(
            Instant::now() < deadline,
            "continuous query did not complete within {timeout:?}"
        );
        thread::yield_now();
    }
}