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,
SqlEventPayload, assign_event_time_watermarks,
};
#[tokio::test]
async fn event_time_watermarks_propagate_through_projection_and_filter() {
let schema = event_schema();
let batches = vec![
event_batch(Arc::clone(&schema), &[1], &[10_000_000_000]),
event_batch(Arc::clone(&schema), &[2], &[8_000_000_000]),
event_batch(Arc::clone(&schema), &[3], &[16_000_000_000]),
];
let context = DatumSqlContext::new();
context
.register_source_with_event_time(
"events",
Arc::clone(&schema),
Source::from_iter(batches),
EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::from_secs(5)),
)
.expect("source registers");
let events = context
.streaming_source("SELECT id FROM events WHERE id >= 1")
.await
.expect("query lowers")
.run_collect()
.expect("events collect");
let mut ids = Vec::new();
let mut watermarks = Vec::new();
for event in events {
match event {
SqlEvent::Data(batch) => {
assert_eq!(batch.num_columns(), 1);
ids.extend(int_values(&batch, 0));
}
SqlEvent::Watermark(watermark) => watermarks.push(watermark.timestamp_ns()),
SqlEvent::Barrier(_) => panic!("test source should not emit barriers"),
}
}
assert_eq!(ids, vec![1, 2, 3]);
assert_eq!(watermarks, vec![5_000_000_000, 11_000_000_000]);
}
#[test]
fn partitioned_event_time_watermarks_min_reduce_active_partitions() {
let schema = event_schema();
let batches = vec![
partitioned_event_batch(Arc::clone(&schema), &[1], &[20_000_000_000], &[1], &[0, 1]),
partitioned_event_batch(Arc::clone(&schema), &[2], &[0], &[0], &[0, 1]),
partitioned_event_batch(Arc::clone(&schema), &[3], &[40_000_000_000], &[0], &[0, 1]),
];
let events = assign_event_time_watermarks(
Source::from_iter(batches),
schema.as_ref(),
EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::ZERO),
)
.expect("watermark assigner builds")
.run_collect()
.expect("events collect");
let watermarks = events
.iter()
.filter_map(|event| {
event
.as_watermark()
.map(|watermark| watermark.timestamp_ns())
})
.collect::<Vec<_>>();
assert_eq!(watermarks, vec![0, 20_000_000_000]);
}
#[test]
fn changelog_payloads_compose_with_sql_event_envelope() {
let schema = event_schema();
let changes = ChangelogBatch::try_new(
vec![
ChangeOp::Insert,
ChangeOp::UpdateDelete,
ChangeOp::UpdateInsert,
],
event_batch(
Arc::clone(&schema),
&[1, 2, 2],
&[10_000_000_000, 12_000_000_000, 16_000_000_000],
),
)
.expect("changelog batch builds");
let events = assign_event_time_watermarks(
Source::from_iter([changes]),
schema.as_ref(),
EventTimeConfig::bounded_out_of_orderness("event_ts", Duration::from_secs(5)),
)
.expect("watermark assigner builds")
.run_collect()
.expect("events collect");
assert_eq!(events.len(), 2);
match &events[0] {
SqlEvent::Data(changes) => {
assert_eq!(
changes.ops(),
&[
ChangeOp::Insert,
ChangeOp::UpdateDelete,
ChangeOp::UpdateInsert,
]
);
assert_eq!(int_values(changes.batch(), 0), vec![1, 2, 2]);
}
SqlEvent::Watermark(_) | SqlEvent::Barrier(_) => {
panic!("first event should carry changelog data")
}
}
assert_eq!(
events[1]
.as_watermark()
.map(|watermark| watermark.timestamp_ns()),
Some(11_000_000_000)
);
}
#[tokio::test]
async fn continuous_query_handle_cancels_unbounded_query() {
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![Arc::new(Int64Array::from(vec![1]))],
)
.expect("batch builds");
let context = DatumSqlContext::new();
context
.register_source("ticks", schema, Source::repeat(batch))
.expect("source registers");
let mut handle = context
.execute_streaming("SELECT id FROM ticks")
.await
.expect("query materializes");
handle.cancel();
let result = wait_for_completion(&mut handle, Duration::from_secs(2));
assert!(matches!(result, Ok(NotUsed)));
}
fn event_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new(
"event_ts",
DataType::Timestamp(TimeUnit::Nanosecond, None),
false,
),
]))
}
fn event_batch(schema: Arc<Schema>, ids: &[i64], timestamps_ns: &[i64]) -> RecordBatch {
RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(ids.to_vec())),
Arc::new(TimestampNanosecondArray::from(timestamps_ns.to_vec())),
],
)
.expect("event batch builds")
}
#[derive(Clone)]
struct PartitionedBatch {
batch: RecordBatch,
row_partitions: Vec<i64>,
active_partitions: Vec<i64>,
}
impl SqlEventPayload for PartitionedBatch {
fn event_time_batch(&self) -> &RecordBatch {
&self.batch
}
fn event_time_partition(&self, row: usize) -> Option<i64> {
self.row_partitions.get(row).copied()
}
fn event_time_active_partitions(&self) -> Option<Vec<i64>> {
Some(self.active_partitions.clone())
}
}
fn partitioned_event_batch(
schema: Arc<Schema>,
ids: &[i64],
timestamps_ns: &[i64],
row_partitions: &[i64],
active_partitions: &[i64],
) -> PartitionedBatch {
assert_eq!(ids.len(), row_partitions.len());
PartitionedBatch {
batch: event_batch(schema, ids, timestamps_ns),
row_partitions: row_partitions.to_vec(),
active_partitions: active_partitions.to_vec(),
}
}
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();
}
}