datum-sql 0.10.3

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

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

#[tokio::test]
async fn windowed_join_matches_out_of_order_inputs() {
    let schema = event_schema();
    let left = event_batch(
        Arc::clone(&schema),
        &[1, 2],
        &[10, 20],
        &[10_000_000_000, 2_000_000_000],
    );
    let right = event_batch(
        Arc::clone(&schema),
        &[1, 2],
        &[100, 200],
        &[8_000_000_000, 20_000_000_000],
    );
    let context = context_with_windowed_sources(
        schema,
        Source::from_iter([left]),
        Source::from_iter([right]),
    );

    let output = collect_data_events(
        context
            .streaming_source(
                "SELECT l.id, l.value AS left_value, r.value AS right_value \
                 FROM lefts l JOIN rights r \
                 ON l.id = r.id \
                 AND l.event_ts >= r.event_ts - INTERVAL '5 seconds' \
                 AND l.event_ts <= r.event_ts + INTERVAL '5 seconds'",
            )
            .await
            .expect("join lowers"),
    );

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

#[tokio::test]
async fn watermark_min_reduction_holds_back_on_idle_input() {
    let schema = event_schema();
    let left = event_batch(Arc::clone(&schema), &[1], &[10], &[10_000_000_000]);
    let context = context_with_windowed_sources(schema, Source::from_iter([left]), Source::empty());

    let events = context
        .streaming_source(
            "SELECT l.id FROM lefts l JOIN rights r \
             ON l.id = r.id \
             AND l.event_ts >= r.event_ts - INTERVAL '5 seconds' \
             AND l.event_ts <= r.event_ts + INTERVAL '5 seconds'",
        )
        .await
        .expect("join lowers")
        .run_collect()
        .expect("events collect");

    assert!(
        events
            .iter()
            .all(|event| !matches!(event, SqlEvent::Watermark(_))),
        "an idle input without a watermark must hold back the joined watermark"
    );
}

#[tokio::test]
async fn windowed_join_evicts_state_on_min_watermark() {
    let schema = event_schema();
    let left = event_batch(
        Arc::clone(&schema),
        &[1, 2],
        &[10, 20],
        &[0, 20_000_000_000],
    );
    let right = event_batch(Arc::clone(&schema), &[3], &[30], &[20_000_000_000]);
    let context = context_with_windowed_sources(
        schema,
        Source::from_iter([left]),
        Source::from_iter([right]),
    );

    let (source, metrics) = context
        .streaming_source_with_metrics(
            "SELECT l.id FROM lefts l JOIN rights r \
             ON l.id = r.id \
             AND l.event_ts >= r.event_ts - INTERVAL '5 seconds' \
             AND l.event_ts <= r.event_ts + INTERVAL '5 seconds'",
        )
        .await
        .expect("join lowers");
    let _events = source.run_collect().expect("events collect");

    assert_eq!(metrics.streaming_join_metrics().evicted_rows(), 1);
    assert_eq!(metrics.streaming_join_metrics().state_rows(), 2);
}

#[tokio::test]
async fn bounded_state_join_errors_at_explicit_limit() {
    let schema = plain_schema();
    let left = plain_batch(Arc::clone(&schema), &[1, 2], &[10, 20]);
    let context = context_with_bounded_sources(
        schema,
        Source::from_iter([left]),
        Source::empty(),
        StreamingJoinStateLimits::new(1, 10),
    );

    let source = context
        .streaming_source("SELECT l.id FROM lefts l JOIN rights r ON l.id = r.id")
        .await
        .expect("join lowers");
    let error = source
        .run_collect()
        .expect_err("second left row should exceed side limit");

    assert!(
        error
            .to_string()
            .contains("streaming join state limit exceeded"),
        "{error}"
    );
}

#[tokio::test]
async fn join_filter_is_applied_after_equi_probe() {
    let schema = plain_schema();
    let left = plain_batch(Arc::clone(&schema), &[1, 2], &[10, 30]);
    let right = plain_batch(Arc::clone(&schema), &[1, 2], &[20, 10]);
    let context = context_with_bounded_sources(
        schema,
        Source::from_iter([left]),
        Source::from_iter([right]),
        StreamingJoinStateLimits::new(10, 20),
    );

    let output = collect_data_events(
        context
            .streaming_source(
                "SELECT l.id, l.value AS left_value, r.value AS right_value \
                 FROM lefts l JOIN rights r ON l.id = r.id AND l.value < r.value",
            )
            .await
            .expect("join lowers"),
    );

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

#[tokio::test]
async fn changelog_join_is_rejected_at_planning() {
    let schema = plain_schema();
    let changes = ChangelogBatch::try_new(
        vec![ChangeOp::Insert],
        plain_batch(Arc::clone(&schema), &[1], &[10]),
    )
    .expect("changes build");
    let context = DatumSqlContext::new().with_streaming_join_config(
        StreamingJoinConfig::bounded_state(StreamingJoinStateLimits::new(10, 20)),
    );
    context
        .register_changelog_source("lefts", Arc::clone(&schema), Source::from_iter([changes]))
        .expect("left registers");
    context
        .register_source("rights", schema, Source::empty())
        .expect("right registers");

    let error = match context
        .streaming_source("SELECT l.id FROM lefts l JOIN rights r ON l.id = r.id")
        .await
    {
        Ok(_) => panic!("changelog join should fail planning"),
        Err(error) => error,
    };

    assert!(
        error
            .to_string()
            .contains("streaming joins do not support changelog inputs yet"),
        "{error}"
    );
}

#[tokio::test]
async fn continuous_join_cancels_while_waiting_on_inputs() {
    let schema = event_schema();
    let left = event_batch(Arc::clone(&schema), &[1], &[10], &[1_000_000_000]);
    let right = event_batch(Arc::clone(&schema), &[1], &[20], &[1_000_000_000]);
    let context =
        context_with_windowed_sources(schema, Source::repeat(left), Source::repeat(right));

    let mut handle = context
        .execute_streaming(
            "SELECT l.id FROM lefts l JOIN rights r \
             ON l.id = r.id \
             AND l.event_ts >= r.event_ts - INTERVAL '5 seconds' \
             AND l.event_ts <= r.event_ts + INTERVAL '5 seconds'",
        )
        .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_join_min_reduces_kafka_partitions_before_late_drop() {
    let schema = event_schema();
    let left_high_partition = partitioned_committable_event_batch(
        "lefts",
        Arc::clone(&schema),
        &[99],
        &[990],
        &[20_000_000_000],
        1,
        &[0, 1],
    );
    let left_slow_partition = partitioned_committable_event_batch(
        "lefts",
        Arc::clone(&schema),
        &[1],
        &[10],
        &[0],
        0,
        &[0, 1],
    );
    let right_match = partitioned_committable_event_batch(
        "rights",
        Arc::clone(&schema),
        &[1],
        &[100],
        &[0],
        0,
        &[0, 1],
    );
    let context = DatumSqlContext::new().with_streaming_join_config(StreamingJoinConfig::windowed(
        StreamingJoinWindow::new("event_ts", Duration::from_secs(5)),
    ));
    context
        .register_committable_source_with_event_time(
            "lefts",
            Arc::clone(&schema),
            Source::from_iter([left_high_partition, left_slow_partition]),
            EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::ZERO),
        )
        .expect("left registers");
    context
        .register_committable_source_with_event_time(
            "rights",
            schema,
            Source::from_iter([right_match]),
            EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::ZERO),
        )
        .expect("right registers");

    let (source, metrics) = context
        .streaming_source_with_metrics(
            "SELECT l.id, l.value AS left_value, r.value AS right_value \
             FROM lefts l JOIN rights r \
             ON l.id = r.id \
             AND l.event_ts >= r.event_ts - INTERVAL '5 seconds' \
             AND l.event_ts <= r.event_ts + INTERVAL '5 seconds'",
        )
        .await
        .expect("join lowers");
    let output = collect_data_events(source);

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

fn context_with_windowed_sources(
    schema: Arc<Schema>,
    left: Source<RecordBatch>,
    right: Source<RecordBatch>,
) -> DatumSqlContext {
    let context = DatumSqlContext::new().with_streaming_join_config(StreamingJoinConfig::windowed(
        StreamingJoinWindow::new("event_ts", Duration::from_secs(5)),
    ));
    context
        .register_source_with_event_time(
            "lefts",
            Arc::clone(&schema),
            left,
            EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::ZERO),
        )
        .expect("left registers");
    context
        .register_source_with_event_time(
            "rights",
            schema,
            right,
            EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::ZERO),
        )
        .expect("right registers");
    context
}

fn context_with_bounded_sources(
    schema: Arc<Schema>,
    left: Source<RecordBatch>,
    right: Source<RecordBatch>,
    limits: StreamingJoinStateLimits,
) -> DatumSqlContext {
    let context = DatumSqlContext::new()
        .with_streaming_join_config(StreamingJoinConfig::bounded_state(limits));
    context
        .register_source("lefts", Arc::clone(&schema), left)
        .expect("left registers");
    context
        .register_source("rights", schema, right)
        .expect("right registers");
    context
}

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()
}

fn event_schema() -> Arc<Schema> {
    Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new("value", DataType::Int64, false),
        Field::new(
            "event_ts",
            DataType::Timestamp(TimeUnit::Nanosecond, None),
            false,
        ),
    ]))
}

fn event_batch(
    schema: Arc<Schema>,
    ids: &[i64],
    values: &[i64],
    timestamps_ns: &[i64],
) -> RecordBatch {
    RecordBatch::try_new(
        schema,
        vec![
            Arc::new(Int64Array::from(ids.to_vec())),
            Arc::new(Int64Array::from(values.to_vec())),
            Arc::new(TimestampNanosecondArray::from(timestamps_ns.to_vec())),
        ],
    )
    .expect("event batch builds")
}

fn plain_schema() -> Arc<Schema> {
    Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new("value", DataType::Int64, false),
    ]))
}

fn plain_batch(schema: Arc<Schema>, ids: &[i64], values: &[i64]) -> RecordBatch {
    RecordBatch::try_new(
        schema,
        vec![
            Arc::new(Int64Array::from(ids.to_vec())),
            Arc::new(Int64Array::from(values.to_vec())),
        ],
    )
    .expect("plain batch builds")
}

#[cfg(feature = "mq")]
fn partitioned_committable_event_batch(
    topic: &str,
    schema: Arc<Schema>,
    ids: &[i64],
    values: &[i64],
    timestamps_ns: &[i64],
    partition: i32,
    active_partitions: &[i32],
) -> CommittableRecordBatch {
    let batch = event_batch(schema, ids, values, timestamps_ns);
    let position = KafkaSourcePosition::from_offsets_with_row_partitions(
        [KafkaPartitionOffset {
            topic: topic.to_owned(),
            partition,
            first_offset: 0,
            last_offset: ids.len() as i64 - 1,
        }],
        std::iter::repeat_n(partition, ids.len()),
        active_partitions
            .iter()
            .map(|partition| KafkaTopicPartition::new(topic, *partition)),
    );
    CommittableRecordBatch::new(
        batch,
        Some(SqlSourcePosition::Kafka(position)),
        0,
        SourceCommit::noop(),
    )
}

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

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();
    }
}