use std::cmp::Ordering as CmpOrdering;
use std::collections::{BTreeMap, HashMap, hash_map::Entry};
use std::env;
use std::fmt;
use std::mem;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::Instant;
use arrow::array::{
Array, ArrayRef, Int64Array, RecordBatch, TimestampMicrosecondArray, TimestampMillisecondArray,
TimestampNanosecondArray, TimestampSecondArray,
types::{IntervalDayTimeType, IntervalMonthDayNanoType},
};
use arrow::datatypes::{DataType, Schema, SchemaRef, TimeUnit};
use datafusion::common::{DataFusionError, Result, ScalarValue};
use datafusion::logical_expr::Accumulator;
use datafusion::physical_expr::expressions::Literal;
use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr};
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode};
use datum::{Source, StreamResult};
use crate::{ChangeOp, ChangelogBatch, SqlEvent, Watermark, stream_error};
#[derive(Clone, Default)]
pub struct WindowedAggregationMetrics {
late_dropped_rows: Arc<AtomicU64>,
}
impl WindowedAggregationMetrics {
pub(crate) fn new(late_dropped_rows: Arc<AtomicU64>) -> Self {
Self { late_dropped_rows }
}
#[must_use]
pub fn late_dropped_rows(&self) -> u64 {
self.late_dropped_rows.load(Ordering::Relaxed)
}
fn record_late_row(&self) {
self.late_dropped_rows.fetch_add(1, Ordering::Relaxed);
}
}
impl fmt::Debug for WindowedAggregationMetrics {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("WindowedAggregationMetrics")
.field("late_dropped_rows", &self.late_dropped_rows())
.finish()
}
}
pub(crate) fn windowed_aggregate_source(
source: Source<SqlEvent<RecordBatch>>,
aggregate: &AggregateExec,
metrics: WindowedAggregationMetrics,
) -> Result<Source<SqlEvent<RecordBatch>>> {
let stage = WindowedAggregationStage::try_new(aggregate, metrics, false)?;
Ok(source.try_stateful_map_concat(stage, |stage, event| stage.apply_record_event(event)))
}
pub(crate) fn windowed_changelog_aggregate_source(
source: Source<SqlEvent<ChangelogBatch>>,
aggregate: &AggregateExec,
metrics: WindowedAggregationMetrics,
) -> Result<Source<SqlEvent<RecordBatch>>> {
let stage = WindowedAggregationStage::try_new(aggregate, metrics, true)?;
Ok(source.try_stateful_map_concat(stage, |stage, event| stage.apply_changelog_event(event)))
}
struct WindowedAggregationStage {
plan: Arc<WindowAggregatePlan>,
metrics: WindowedAggregationMetrics,
profile: WindowProfile,
retractable_input: bool,
current_watermark_ns: Option<i64>,
windows: HashMap<WindowGroupKey, WindowEntry>,
windows_by_end: BTreeMap<i64, Vec<WindowGroupKey>>,
}
impl Clone for WindowedAggregationStage {
fn clone(&self) -> Self {
Self {
plan: Arc::clone(&self.plan),
metrics: self.metrics.clone(),
profile: self.profile.clone(),
retractable_input: self.retractable_input,
current_watermark_ns: None,
windows: HashMap::new(),
windows_by_end: BTreeMap::new(),
}
}
}
impl Drop for WindowedAggregationStage {
fn drop(&mut self) {
self.profile.report();
}
}
impl WindowedAggregationStage {
fn try_new(
aggregate: &AggregateExec,
metrics: WindowedAggregationMetrics,
retractable_input: bool,
) -> Result<Self> {
Ok(Self {
plan: Arc::new(WindowAggregatePlan::try_new(aggregate, retractable_input)?),
metrics,
profile: WindowProfile::new(),
retractable_input,
current_watermark_ns: None,
windows: HashMap::new(),
windows_by_end: BTreeMap::new(),
})
}
fn apply_record_event(
&mut self,
event: SqlEvent<RecordBatch>,
) -> StreamResult<Vec<SqlEvent<RecordBatch>>> {
match event {
SqlEvent::Data(batch) => self.apply_batch(&batch, None).map_err(stream_error),
SqlEvent::Watermark(watermark) => self.apply_watermark(watermark).map_err(stream_error),
SqlEvent::Barrier(barrier) => Ok(vec![SqlEvent::Barrier(barrier)]),
}
}
fn apply_changelog_event(
&mut self,
event: SqlEvent<ChangelogBatch>,
) -> StreamResult<Vec<SqlEvent<RecordBatch>>> {
match event {
SqlEvent::Data(changes) => self
.apply_batch(changes.batch(), Some(changes.ops()))
.map_err(stream_error),
SqlEvent::Watermark(watermark) => self.apply_watermark(watermark).map_err(stream_error),
SqlEvent::Barrier(barrier) => Ok(vec![SqlEvent::Barrier(barrier)]),
}
}
fn apply_batch(
&mut self,
batch: &RecordBatch,
ops: Option<&[ChangeOp]>,
) -> Result<Vec<SqlEvent<RecordBatch>>> {
if batch.num_rows() == 0 {
return Ok(Vec::new());
}
if let Some(ops) = ops
&& ops.len() != batch.num_rows()
{
return Err(DataFusionError::Plan(format!(
"windowed aggregate received {} changelog ops for {} rows",
ops.len(),
batch.num_rows()
)));
}
self.profile.record_batch(batch.num_rows());
let prepare_start = self.profile.start_timer();
let prepared = self.plan.prepare_batch(batch)?;
self.profile.record_prepare(prepare_start);
if self.plan.use_batch_grouping() {
return self.apply_grouped_batch(batch.num_rows(), &prepared, ops);
}
let plan = Arc::clone(&self.plan);
let profile = self.profile.clone();
let mut assignments = Vec::new();
for row in 0..batch.num_rows() {
let event_time_ns = timestamp_ns_from_array(&prepared.event_times, row)?;
if self.row_is_late(event_time_ns) {
self.metrics.record_late_row();
self.profile.record_late_row();
continue;
}
let assignment_start = self.profile.start_timer();
self.plan
.window
.assignments_into(event_time_ns, &mut assignments)?;
self.profile
.record_assignments(assignment_start, assignments.len());
for assignment in assignments.drain(..) {
let key_start = self.profile.start_timer();
let key = self.plan.group_key(&prepared, row, assignment.start_ns)?;
let change = ops.map(|ops| ops[row]).unwrap_or(ChangeOp::Insert);
let change = [(row, change)];
let entry = self.entry_for_key(key, assignment.end_ns, change[0].1)?;
profile.record_key_lookup(key_start);
let update_start = profile.start_timer();
plan.apply_rows(entry, &prepared, &change)?;
profile.record_accumulator_update(update_start);
}
}
Ok(Vec::new())
}
fn apply_grouped_batch(
&mut self,
rows: usize,
prepared: &PreparedBatch,
ops: Option<&[ChangeOp]>,
) -> Result<Vec<SqlEvent<RecordBatch>>> {
let mut pending = HashMap::<WindowGroupKey, PendingWindowGroup>::new();
let mut assignments = Vec::new();
for row in 0..rows {
let event_time_ns = timestamp_ns_from_array(&prepared.event_times, row)?;
if self.row_is_late(event_time_ns) {
self.metrics.record_late_row();
self.profile.record_late_row();
continue;
}
let assignment_start = self.profile.start_timer();
self.plan
.window
.assignments_into(event_time_ns, &mut assignments)?;
self.profile
.record_assignments(assignment_start, assignments.len());
for assignment in assignments.drain(..) {
let key = self.plan.group_key(prepared, row, assignment.start_ns)?;
let change = ops.map(|ops| ops[row]).unwrap_or(ChangeOp::Insert);
pending
.entry(key)
.or_insert_with(|| PendingWindowGroup::new(assignment.end_ns))
.changes
.push((row, change));
}
}
let plan = Arc::clone(&self.plan);
let profile = self.profile.clone();
for (key, group) in pending {
let key_start = profile.start_timer();
let first_change = group
.changes
.first()
.map(|(_row, change)| *change)
.unwrap_or(ChangeOp::Insert);
let entry = self.entry_for_key(key, group.window_end_ns, first_change)?;
profile.record_key_lookup(key_start);
let update_start = profile.start_timer();
plan.apply_rows(entry, prepared, &group.changes)?;
profile.record_accumulator_update(update_start);
}
Ok(Vec::new())
}
fn entry_for_key(
&mut self,
key: WindowGroupKey,
window_end_ns: i64,
first_change: ChangeOp,
) -> Result<&mut WindowEntry> {
let open_entries = self.windows.len() + 1;
let plan = Arc::clone(&self.plan);
let profile = self.profile.clone();
match self.windows.entry(key) {
Entry::Occupied(entry) => Ok(entry.into_mut()),
Entry::Vacant(entry) => {
if first_change.is_retraction() {
return Err(DataFusionError::Plan(
"windowed aggregate received a retraction for an absent open window/key"
.into(),
));
}
let key = entry.key().clone();
let window_entry = WindowEntry::new(window_end_ns, plan.create_aggregate_states()?);
profile.record_entry_insert(open_entries, &key, &window_entry);
self.windows_by_end
.entry(window_end_ns)
.or_default()
.push(key);
Ok(entry.insert(window_entry))
}
}
}
fn apply_watermark(&mut self, watermark: Watermark) -> Result<Vec<SqlEvent<RecordBatch>>> {
let watermark_ns = self
.current_watermark_ns
.map_or(watermark.timestamp_ns(), |current| {
current.max(watermark.timestamp_ns())
});
self.current_watermark_ns = Some(watermark_ns);
let mut out = self.emit_ready_windows(watermark_ns)?;
out.push(SqlEvent::Watermark(Watermark::new(watermark_ns)));
Ok(out)
}
fn row_is_late(&self, event_time_ns: i64) -> bool {
self.current_watermark_ns
.is_some_and(|watermark_ns| event_time_ns <= watermark_ns)
}
fn emit_ready_windows(&mut self, watermark_ns: i64) -> Result<Vec<SqlEvent<RecordBatch>>> {
let ready_ends = self
.windows_by_end
.keys()
.take_while(|end_ns| **end_ns <= watermark_ns)
.copied()
.collect::<Vec<_>>();
let ready_len = ready_ends
.iter()
.filter_map(|end_ns| self.windows_by_end.get(end_ns).map(Vec::len))
.sum();
let mut rows = Vec::with_capacity(ready_len);
for end_ns in ready_ends {
if let Some(mut keys) = self.windows_by_end.remove(&end_ns) {
keys.sort_by(WindowGroupKey::sort_cmp);
for key in keys {
if let Some(mut entry) = self.windows.remove(&key) {
rows.push(self.plan.evaluate_entry(&key, &mut entry)?);
}
}
}
}
if rows.is_empty() {
return Ok(Vec::new());
}
Ok(vec![SqlEvent::Data(self.plan.build_output_batch(rows)?)])
}
}
#[derive(Clone)]
struct WindowProfile {
enabled: bool,
counters: Arc<WindowProfileCounters>,
}
impl WindowProfile {
fn new() -> Self {
Self {
enabled: window_profile_enabled(),
counters: Arc::new(WindowProfileCounters::default()),
}
}
fn record_batch(&self, rows: usize) {
if self.enabled {
self.counters.batches.fetch_add(1, Ordering::Relaxed);
self.counters
.input_rows
.fetch_add(rows as u64, Ordering::Relaxed);
}
}
fn record_late_row(&self) {
if self.enabled {
self.counters.late_rows.fetch_add(1, Ordering::Relaxed);
}
}
fn start_timer(&self) -> Option<Instant> {
self.enabled.then(Instant::now)
}
fn record_prepare(&self, start: Option<Instant>) {
self.record_elapsed(&self.counters.prepare_ns, start);
}
fn record_assignments(&self, start: Option<Instant>, assignments: usize) {
if self.enabled {
self.counters
.assignments
.fetch_add(assignments as u64, Ordering::Relaxed);
self.record_elapsed(&self.counters.assignment_ns, start);
}
}
fn record_key_lookup(&self, start: Option<Instant>) {
self.record_elapsed(&self.counters.key_lookup_ns, start);
}
fn record_accumulator_update(&self, start: Option<Instant>) {
if self.enabled {
self.counters.update_calls.fetch_add(1, Ordering::Relaxed);
self.record_elapsed(&self.counters.update_ns, start);
}
}
fn record_entry_insert(&self, open_entries: usize, key: &WindowGroupKey, entry: &WindowEntry) {
if self.enabled {
self.counters.entry_inserts.fetch_add(1, Ordering::Relaxed);
self.counters
.entry_estimated_bytes
.fetch_add(estimated_entry_bytes(key, entry) as u64, Ordering::Relaxed);
update_max(&self.counters.max_open_entries, open_entries as u64);
}
}
fn record_elapsed(&self, counter: &AtomicU64, start: Option<Instant>) {
if let Some(start) = start {
counter.fetch_add(start.elapsed().as_nanos() as u64, Ordering::Relaxed);
}
}
fn report(&self) {
if !self.enabled || self.counters.input_rows.load(Ordering::Relaxed) == 0 {
return;
}
if self.counters.reported.swap(true, Ordering::Relaxed) {
return;
}
let batches = self.counters.batches.load(Ordering::Relaxed);
let input_rows = self.counters.input_rows.load(Ordering::Relaxed);
let assignments = self.counters.assignments.load(Ordering::Relaxed);
let update_calls = self.counters.update_calls.load(Ordering::Relaxed);
let entry_inserts = self.counters.entry_inserts.load(Ordering::Relaxed);
let estimated_bytes = self.counters.entry_estimated_bytes.load(Ordering::Relaxed);
let avg_entry_bytes = estimated_bytes.checked_div(entry_inserts).unwrap_or(0);
eprintln!(
"DATUM_SQL_WINDOW_PROFILE batches={batches} input_rows={input_rows} late_rows={} assignments={assignments} fanout_per_row={:.3} update_calls={update_calls} entry_inserts={entry_inserts} max_open_entries={} avg_entry_estimated_bytes={avg_entry_bytes} prepare_ms={:.3} assignment_ms={:.3} key_lookup_ms={:.3} accumulator_update_ms={:.3}",
self.counters.late_rows.load(Ordering::Relaxed),
assignments as f64 / input_rows.max(1) as f64,
self.counters.max_open_entries.load(Ordering::Relaxed),
ns_to_ms(self.counters.prepare_ns.load(Ordering::Relaxed)),
ns_to_ms(self.counters.assignment_ns.load(Ordering::Relaxed)),
ns_to_ms(self.counters.key_lookup_ns.load(Ordering::Relaxed)),
ns_to_ms(self.counters.update_ns.load(Ordering::Relaxed)),
);
}
}
#[derive(Default)]
struct WindowProfileCounters {
reported: AtomicBool,
batches: AtomicU64,
input_rows: AtomicU64,
late_rows: AtomicU64,
assignments: AtomicU64,
update_calls: AtomicU64,
entry_inserts: AtomicU64,
max_open_entries: AtomicU64,
entry_estimated_bytes: AtomicU64,
prepare_ns: AtomicU64,
assignment_ns: AtomicU64,
key_lookup_ns: AtomicU64,
update_ns: AtomicU64,
}
fn window_profile_enabled() -> bool {
static ENABLED: OnceLock<bool> = OnceLock::new();
*ENABLED.get_or_init(|| {
env::var("DATUM_SQL_WINDOW_PROFILE")
.is_ok_and(|value| value == "1" || value.eq_ignore_ascii_case("true"))
})
}
fn update_max(counter: &AtomicU64, value: u64) {
let mut current = counter.load(Ordering::Relaxed);
while value > current {
match counter.compare_exchange_weak(current, value, Ordering::Relaxed, Ordering::Relaxed) {
Ok(_) => break,
Err(next) => current = next,
}
}
}
fn estimated_entry_bytes(key: &WindowGroupKey, entry: &WindowEntry) -> usize {
let key_heap = match key {
WindowGroupKey::FastI64 { values, .. } => match values {
FastI64GroupValues::None | FastI64GroupValues::One(_) => 0,
FastI64GroupValues::Many(values) => values.capacity() * mem::size_of::<i64>(),
},
WindowGroupKey::Scalar(values) => values
.capacity()
.saturating_mul(mem::size_of::<ScalarValue>())
.saturating_add(
values
.iter()
.map(|value| value.size().saturating_sub(mem::size_of::<ScalarValue>()))
.sum::<usize>(),
),
};
let state_heap = entry
.states
.capacity()
.saturating_mul(mem::size_of::<AggregateState>())
.saturating_add(
entry
.states
.iter()
.map(AggregateState::estimated_size)
.sum::<usize>(),
);
mem::size_of::<WindowGroupKey>()
.saturating_add(key_heap)
.saturating_add(mem::size_of::<WindowEntry>())
.saturating_add(state_heap)
}
fn ns_to_ms(ns: u64) -> f64 {
ns as f64 / 1_000_000.0
}
struct WindowAggregatePlan {
output_schema: SchemaRef,
window: WindowSpec,
group_exprs: Vec<Arc<dyn PhysicalExpr>>,
group_types: Vec<DataType>,
aggregate_exprs: Vec<Arc<datafusion::physical_expr::aggregate::AggregateFunctionExpr>>,
aggregate_kinds: Vec<AggregateKind>,
key_kind: WindowGroupKeyKind,
use_batch_grouping: bool,
group_count: usize,
aggregate_count: usize,
retractable_input: bool,
}
impl WindowAggregatePlan {
fn try_new(aggregate: &AggregateExec, retractable_input: bool) -> Result<Self> {
match aggregate.mode() {
AggregateMode::Single | AggregateMode::SinglePartitioned => {}
other => {
return Err(DataFusionError::NotImplemented(format!(
"datum-sql windowed aggregation lowers single-phase aggregates only, found {other:?}"
)));
}
}
if !aggregate.group_expr().is_single() || aggregate.group_expr().groups().len() != 1 {
return Err(DataFusionError::NotImplemented(
"datum-sql windowed aggregation does not support grouping sets, cube, or rollup"
.into(),
));
}
if aggregate.filter_expr().iter().any(Option::is_some) {
return Err(DataFusionError::NotImplemented(
"datum-sql windowed aggregation does not support aggregate FILTER clauses yet"
.into(),
));
}
let group_exprs = aggregate
.group_expr()
.expr()
.iter()
.map(|(expr, _name)| Arc::clone(expr))
.collect::<Vec<_>>();
let input_schema = aggregate.input_schema();
let group_types = group_exprs
.iter()
.map(|expr| expr.data_type(input_schema.as_ref()))
.collect::<Result<Vec<_>>>()?;
let (window_index, window) = find_window_spec(&group_exprs, input_schema.as_ref())?;
if window_index == usize::MAX {
return Err(DataFusionError::Internal(
"window index sentinel escaped validation".into(),
));
}
let key_kind =
if group_types.iter().enumerate().all(|(index, data_type)| {
index == window_index || matches!(data_type, DataType::Int64)
}) {
WindowGroupKeyKind::FastI64
} else {
WindowGroupKeyKind::Scalar
};
let use_batch_grouping = retractable_input
&& matches!(&window, WindowSpec::Tumble { .. })
&& group_exprs.len().saturating_sub(1) > 0;
let aggregate_exprs = aggregate.aggr_expr().to_vec();
let mut aggregate_kinds = Vec::with_capacity(aggregate_exprs.len());
for aggregate_expr in &aggregate_exprs {
let name = aggregate_expr.fun().name().to_ascii_lowercase();
if !matches!(name.as_str(), "count" | "sum" | "avg" | "min" | "max") {
return Err(DataFusionError::NotImplemented(format!(
"datum-sql windowed aggregation supports COUNT, SUM, AVG, MIN, and MAX for now, found {}",
aggregate_expr.fun().name()
)));
}
if aggregate_expr.is_distinct() {
return Err(DataFusionError::NotImplemented(
"datum-sql windowed aggregation does not support DISTINCT aggregates yet"
.into(),
));
}
if !aggregate_expr.order_bys().is_empty() {
return Err(DataFusionError::NotImplemented(
"datum-sql windowed aggregation does not support ORDER BY aggregates yet"
.into(),
));
}
if retractable_input {
aggregate_expr.create_sliding_accumulator()?;
}
aggregate_kinds.push(aggregate_kind(
aggregate_expr,
input_schema.as_ref(),
retractable_input,
)?);
}
Ok(Self {
output_schema: aggregate.schema(),
window,
group_types,
group_count: group_exprs.len(),
aggregate_count: aggregate_exprs.len(),
group_exprs,
aggregate_exprs,
aggregate_kinds,
key_kind,
use_batch_grouping,
retractable_input,
})
}
fn use_batch_grouping(&self) -> bool {
self.use_batch_grouping
}
fn prepare_batch(&self, batch: &RecordBatch) -> Result<PreparedBatch> {
let event_times = self
.window
.event_time_expr()
.evaluate(batch)?
.into_array(batch.num_rows())?;
let group_values = self
.group_exprs
.iter()
.map(|expr| expr.evaluate(batch)?.into_array(batch.num_rows()))
.collect::<Result<Vec<_>>>()?;
let aggregate_values = self
.aggregate_exprs
.iter()
.map(|aggregate_expr| {
aggregate_expr
.expressions()
.iter()
.map(|expr| expr.evaluate(batch)?.into_array(batch.num_rows()))
.collect::<Result<Vec<_>>>()
})
.collect::<Result<Vec<_>>>()?;
Ok(PreparedBatch {
event_times,
group_values,
aggregate_values,
})
}
fn group_key(
&self,
prepared: &PreparedBatch,
row: usize,
window_start_ns: i64,
) -> Result<WindowGroupKey> {
match self.key_kind {
WindowGroupKeyKind::FastI64 => {
let mut values = FastI64GroupValues::None;
for (index, array) in prepared.group_values.iter().enumerate() {
if index == self.window.group_index() {
continue;
}
let array = array.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
DataFusionError::Internal(
"fast window group key expected Int64Array".into(),
)
})?;
if array.is_null(row) {
return Err(DataFusionError::Plan(format!(
"window group expression contains null at row {row}"
)));
}
values.push(array.value(row));
}
Ok(WindowGroupKey::FastI64 {
window_start_ns,
values,
})
}
WindowGroupKeyKind::Scalar => Ok(WindowGroupKey::Scalar(
prepared
.group_values
.iter()
.enumerate()
.map(|(index, array)| {
if index == self.window.group_index() {
timestamp_scalar_from_ns(window_start_ns, self.window.output_type())
} else {
ScalarValue::try_from_array(array.as_ref(), row)
}
})
.collect::<Result<Vec<_>>>()?,
)),
}
}
fn group_values_from_key(&self, key: &WindowGroupKey) -> Result<Vec<ScalarValue>> {
match key {
WindowGroupKey::FastI64 {
window_start_ns,
values,
} => {
let mut out = Vec::with_capacity(self.group_count);
let mut fast_index = 0;
for (index, data_type) in self.group_types.iter().enumerate() {
if index == self.window.group_index() {
out.push(timestamp_scalar_from_ns(
*window_start_ns,
self.window.output_type(),
)?);
} else {
let value = values.get(fast_index).ok_or_else(|| {
DataFusionError::Internal("fast group key value missing".into())
})?;
if !matches!(data_type, DataType::Int64) {
return Err(DataFusionError::Internal(
"fast group key carried a non-Int64 field".into(),
));
}
out.push(ScalarValue::Int64(Some(value)));
fast_index += 1;
}
}
Ok(out)
}
WindowGroupKey::Scalar(values) => Ok(values.clone()),
}
}
fn create_aggregate_states(&self) -> Result<Vec<AggregateState>> {
self.aggregate_exprs
.iter()
.zip(&self.aggregate_kinds)
.map(|(expr, kind)| {
Ok(match kind {
AggregateKind::Count => AggregateState::Count { count: 0 },
AggregateKind::SumInt64 => AggregateState::SumInt64 { sum: 0, count: 0 },
AggregateKind::MinInt64 => AggregateState::MinInt64 { value: None },
AggregateKind::MaxInt64 => AggregateState::MaxInt64 { value: None },
AggregateKind::Generic => {
let accumulator = if self.retractable_input {
expr.create_sliding_accumulator()
} else {
expr.create_accumulator()
}?;
AggregateState::Generic(accumulator)
}
})
})
.collect()
}
fn apply_rows(
&self,
entry: &mut WindowEntry,
prepared: &PreparedBatch,
changes: &[(usize, ChangeOp)],
) -> Result<()> {
for (aggregate_index, state) in entry.states.iter_mut().enumerate() {
state.apply_rows(&prepared.aggregate_values[aggregate_index], changes)?;
}
Ok(())
}
fn evaluate_entry(
&self,
key: &WindowGroupKey,
entry: &mut WindowEntry,
) -> Result<Vec<ScalarValue>> {
let mut row = Vec::with_capacity(self.group_count + self.aggregate_count);
row.extend(self.group_values_from_key(key)?);
for state in &mut entry.states {
row.push(state.evaluate()?);
}
Ok(row)
}
fn build_output_batch(&self, rows: Vec<Vec<ScalarValue>>) -> Result<RecordBatch> {
let column_count = self.output_schema.fields().len();
let mut columns = vec![Vec::with_capacity(rows.len()); column_count];
for row in rows {
if row.len() != column_count {
return Err(DataFusionError::Internal(format!(
"windowed aggregate produced {} values for {column_count} output columns",
row.len()
)));
}
for (index, value) in row.into_iter().enumerate() {
columns[index].push(value);
}
}
let arrays = columns
.into_iter()
.map(ScalarValue::iter_to_array)
.collect::<Result<Vec<_>>>()?;
RecordBatch::try_new(Arc::clone(&self.output_schema), arrays).map_err(DataFusionError::from)
}
}
struct PreparedBatch {
event_times: ArrayRef,
group_values: Vec<ArrayRef>,
aggregate_values: Vec<Vec<ArrayRef>>,
}
struct PendingWindowGroup {
window_end_ns: i64,
changes: Vec<(usize, ChangeOp)>,
}
impl PendingWindowGroup {
fn new(window_end_ns: i64) -> Self {
Self {
window_end_ns,
changes: Vec::new(),
}
}
}
struct WindowEntry {
states: Vec<AggregateState>,
}
impl WindowEntry {
fn new(_window_end_ns: i64, states: Vec<AggregateState>) -> Self {
Self { states }
}
}
enum AggregateState {
Count { count: i64 },
SumInt64 { sum: i64, count: u64 },
MinInt64 { value: Option<i64> },
MaxInt64 { value: Option<i64> },
Generic(Box<dyn Accumulator>),
}
impl AggregateState {
fn apply_rows(&mut self, values: &[ArrayRef], changes: &[(usize, ChangeOp)]) -> Result<()> {
match self {
Self::Count { count } => {
for (row, change) in changes {
if row_is_counted(values, *row) {
if change.is_positive() {
*count += 1;
} else {
*count -= 1;
}
}
}
}
Self::SumInt64 { sum, count } => {
let array = single_int64_input(values, "SUM")?;
for (row, change) in changes {
if array.is_null(*row) {
continue;
}
let value = array.value(*row);
if change.is_positive() {
*sum = sum.wrapping_add(value);
*count += 1;
} else {
*sum = sum.wrapping_sub(value);
*count = count.saturating_sub(1);
}
}
}
Self::MinInt64 { value } => {
let array = single_int64_input(values, "MIN")?;
for (row, change) in changes {
if change.is_retraction() {
return Err(DataFusionError::Internal(
"fast MIN state cannot retract; generic sliding state should have been selected"
.into(),
));
}
if !array.is_null(*row) {
let next = array.value(*row);
*value = Some(value.map_or(next, |current| current.min(next)));
}
}
}
Self::MaxInt64 { value } => {
let array = single_int64_input(values, "MAX")?;
for (row, change) in changes {
if change.is_retraction() {
return Err(DataFusionError::Internal(
"fast MAX state cannot retract; generic sliding state should have been selected"
.into(),
));
}
if !array.is_null(*row) {
let next = array.value(*row);
*value = Some(value.map_or(next, |current| current.max(next)));
}
}
}
Self::Generic(accumulator) => {
for (row, change) in changes {
let values = values
.iter()
.map(|array| single_value_array(array.as_ref(), *row))
.collect::<Result<Vec<_>>>()?;
if change.is_positive() {
accumulator.update_batch(&values)?;
} else {
accumulator.retract_batch(&values)?;
}
}
}
}
Ok(())
}
fn evaluate(&mut self) -> Result<ScalarValue> {
match self {
Self::Count { count } => Ok(ScalarValue::Int64(Some(*count))),
Self::SumInt64 { sum, count } => {
if *count == 0 {
Ok(ScalarValue::Int64(None))
} else {
Ok(ScalarValue::Int64(Some(*sum)))
}
}
Self::MinInt64 { value } | Self::MaxInt64 { value } => Ok(ScalarValue::Int64(*value)),
Self::Generic(accumulator) => accumulator.evaluate(),
}
}
fn estimated_size(&self) -> usize {
match self {
Self::Count { .. } => mem::size_of::<Self>(),
Self::SumInt64 { .. } => mem::size_of::<Self>(),
Self::MinInt64 { .. } | Self::MaxInt64 { .. } => mem::size_of::<Self>(),
Self::Generic(accumulator) => mem::size_of::<Self>() + accumulator.size(),
}
}
}
#[derive(Clone, Copy)]
enum AggregateKind {
Count,
SumInt64,
MinInt64,
MaxInt64,
Generic,
}
#[derive(Clone, Copy)]
enum WindowGroupKeyKind {
FastI64,
Scalar,
}
#[derive(Clone, Eq, Hash, PartialEq)]
enum WindowGroupKey {
FastI64 {
window_start_ns: i64,
values: FastI64GroupValues,
},
Scalar(Vec<ScalarValue>),
}
impl WindowGroupKey {
fn sort_cmp(left: &Self, right: &Self) -> CmpOrdering {
match (left, right) {
(
Self::FastI64 {
window_start_ns: left_window,
values: left_values,
},
Self::FastI64 {
window_start_ns: right_window,
values: right_values,
},
) => left_window
.cmp(right_window)
.then_with(|| left_values.cmp(right_values)),
_ => left.sort_key().cmp(&right.sort_key()),
}
}
fn sort_key(&self) -> String {
match self {
Self::FastI64 {
window_start_ns,
values,
} => format!("{window_start_ns}|{}", values.sort_key()),
Self::Scalar(values) => values
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join("|"),
}
}
}
#[derive(Clone, Eq, Hash, PartialEq)]
enum FastI64GroupValues {
None,
One(i64),
Many(Vec<i64>),
}
impl FastI64GroupValues {
fn push(&mut self, value: i64) {
*self = match self {
Self::None => Self::One(value),
Self::One(current) => Self::Many(vec![*current, value]),
Self::Many(values) => {
values.push(value);
return;
}
};
}
fn get(&self, index: usize) -> Option<i64> {
match (self, index) {
(Self::None, _) => None,
(Self::One(value), 0) => Some(*value),
(Self::One(_), _) => None,
(Self::Many(values), index) => values.get(index).copied(),
}
}
fn sort_key(&self) -> String {
match self {
Self::None => String::new(),
Self::One(value) => value.to_string(),
Self::Many(values) => values
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join("|"),
}
}
}
impl PartialOrd for FastI64GroupValues {
fn partial_cmp(&self, other: &Self) -> Option<CmpOrdering> {
Some(self.cmp(other))
}
}
impl Ord for FastI64GroupValues {
fn cmp(&self, other: &Self) -> CmpOrdering {
match (self, other) {
(Self::None, Self::None) => CmpOrdering::Equal,
(Self::None, _) => CmpOrdering::Less,
(_, Self::None) => CmpOrdering::Greater,
(Self::One(left), Self::One(right)) => left.cmp(right),
(Self::One(left), Self::Many(right)) => {
std::slice::from_ref(left).cmp(right.as_slice())
}
(Self::Many(left), Self::One(right)) => {
left.as_slice().cmp(std::slice::from_ref(right))
}
(Self::Many(left), Self::Many(right)) => left.cmp(right),
}
}
}
struct WindowAssignment {
start_ns: i64,
end_ns: i64,
}
#[derive(Clone)]
enum WindowSpec {
Tumble {
group_index: usize,
width_ns: i64,
origin_ns: i64,
event_time_expr: Arc<dyn PhysicalExpr>,
output_type: DataType,
},
Hop {
group_index: usize,
slide_ns: i64,
width_ns: i64,
origin_ns: i64,
event_time_expr: Arc<dyn PhysicalExpr>,
output_type: DataType,
},
}
impl WindowSpec {
fn event_time_expr(&self) -> &Arc<dyn PhysicalExpr> {
match self {
Self::Tumble {
event_time_expr, ..
}
| Self::Hop {
event_time_expr, ..
} => event_time_expr,
}
}
fn output_type(&self) -> &DataType {
match self {
Self::Tumble { output_type, .. } | Self::Hop { output_type, .. } => output_type,
}
}
fn group_index(&self) -> usize {
match self {
Self::Tumble { group_index, .. } | Self::Hop { group_index, .. } => *group_index,
}
}
fn assignments_into(&self, event_time_ns: i64, out: &mut Vec<WindowAssignment>) -> Result<()> {
out.clear();
match self {
Self::Tumble {
width_ns,
origin_ns,
..
} => {
let start_ns = floor_to_origin(event_time_ns, *width_ns, *origin_ns)?;
out.push(WindowAssignment {
start_ns,
end_ns: checked_add_ns(start_ns, *width_ns)?,
});
Ok(())
}
Self::Hop {
slide_ns,
width_ns,
origin_ns,
..
} => {
let latest_start = floor_to_origin(event_time_ns, *slide_ns, *origin_ns)?;
let hop_count = width_ns.checked_div(*slide_ns).ok_or_else(|| {
DataFusionError::Plan("hop slide interval cannot be zero".into())
})?;
out.reserve(hop_count as usize);
for hop in 0..hop_count {
let start_ns = latest_start
.checked_sub(hop.checked_mul(*slide_ns).ok_or_else(|| {
DataFusionError::Plan(
"hop assignment underflowed i64 nanoseconds".into(),
)
})?)
.ok_or_else(|| {
DataFusionError::Plan(
"hop assignment underflowed i64 nanoseconds".into(),
)
})?;
let end_ns = checked_add_ns(start_ns, *width_ns)?;
if event_time_ns >= start_ns && event_time_ns < end_ns {
out.push(WindowAssignment { start_ns, end_ns });
}
}
Ok(())
}
}
}
}
fn aggregate_kind(
aggregate_expr: &datafusion::physical_expr::aggregate::AggregateFunctionExpr,
input_schema: &Schema,
retractable_input: bool,
) -> Result<AggregateKind> {
let name = aggregate_expr.fun().name().to_ascii_lowercase();
match name.as_str() {
"count" => Ok(AggregateKind::Count),
"sum" => {
if aggregate_expr.expressions().len() == 1
&& aggregate_expr.expressions()[0].data_type(input_schema)? == DataType::Int64
&& aggregate_expr.field().data_type() == &DataType::Int64
{
Ok(AggregateKind::SumInt64)
} else {
Ok(AggregateKind::Generic)
}
}
"min" => {
if !retractable_input
&& aggregate_expr.expressions().len() == 1
&& aggregate_expr.expressions()[0].data_type(input_schema)? == DataType::Int64
&& aggregate_expr.field().data_type() == &DataType::Int64
{
Ok(AggregateKind::MinInt64)
} else {
Ok(AggregateKind::Generic)
}
}
"max" => {
if !retractable_input
&& aggregate_expr.expressions().len() == 1
&& aggregate_expr.expressions()[0].data_type(input_schema)? == DataType::Int64
&& aggregate_expr.field().data_type() == &DataType::Int64
{
Ok(AggregateKind::MaxInt64)
} else {
Ok(AggregateKind::Generic)
}
}
_ => Ok(AggregateKind::Generic),
}
}
fn find_window_spec(
group_exprs: &[Arc<dyn PhysicalExpr>],
input_schema: &Schema,
) -> Result<(usize, WindowSpec)> {
let mut found = None;
for (index, expr) in group_exprs.iter().enumerate() {
if let Some(spec) = window_spec_from_expr(index, expr, input_schema)? {
if found.is_some() {
return Err(DataFusionError::NotImplemented(
"datum-sql windowed aggregation supports exactly one window group expression"
.into(),
));
}
found = Some((index, spec));
}
}
found.ok_or_else(|| {
DataFusionError::NotImplemented(
"datum-sql windowed aggregation requires GROUP BY date_bin(...), tumble(...), or hop(...)"
.into(),
)
})
}
fn window_spec_from_expr(
group_index: usize,
expr: &Arc<dyn PhysicalExpr>,
input_schema: &Schema,
) -> Result<Option<WindowSpec>> {
let Some(function) = expr.downcast_ref::<ScalarFunctionExpr>() else {
return Ok(None);
};
match function.name().to_ascii_lowercase().as_str() {
"date_bin" => {
if function.args().len() != 2 {
return Err(DataFusionError::Plan(format!(
"date_bin window grouping expects two arguments in WP-SQL-3' lowering, found {}",
function.args().len()
)));
}
let width_ns = fixed_interval_ns(&function.args()[0])?;
let output_type = expr.data_type(input_schema)?;
validate_timestamp_output_type(&output_type, "date_bin")?;
Ok(Some(WindowSpec::Tumble {
group_index,
width_ns,
origin_ns: 0,
event_time_expr: Arc::clone(&function.args()[1]),
output_type,
}))
}
"tumble" => {
if function.args().len() != 2 {
return Err(DataFusionError::Plan(format!(
"tumble window grouping expects two arguments, found {}",
function.args().len()
)));
}
let width_ns = fixed_interval_ns(&function.args()[0])?;
let output_type = expr.data_type(input_schema)?;
validate_timestamp_output_type(&output_type, "tumble")?;
Ok(Some(WindowSpec::Tumble {
group_index,
width_ns,
origin_ns: 0,
event_time_expr: Arc::clone(&function.args()[1]),
output_type,
}))
}
"hop" => {
if function.args().len() != 3 {
return Err(DataFusionError::Plan(format!(
"hop window grouping expects three arguments, found {}",
function.args().len()
)));
}
let slide_ns = fixed_interval_ns(&function.args()[0])?;
let width_ns = fixed_interval_ns(&function.args()[1])?;
if width_ns % slide_ns != 0 {
return Err(DataFusionError::Plan(format!(
"hop window width {width_ns}ns must be an integer multiple of slide {slide_ns}ns"
)));
}
let output_type = expr.data_type(input_schema)?;
validate_timestamp_output_type(&output_type, "hop")?;
Ok(Some(WindowSpec::Hop {
group_index,
slide_ns,
width_ns,
origin_ns: 0,
event_time_expr: Arc::clone(&function.args()[2]),
output_type,
}))
}
_ => Ok(None),
}
}
fn fixed_interval_ns(expr: &Arc<dyn PhysicalExpr>) -> Result<i64> {
let literal = expr.downcast_ref::<Literal>().ok_or_else(|| {
DataFusionError::Plan(
"window interval must be a literal INTERVAL for WP-SQL-3' lowering".into(),
)
})?;
let nanos = match literal.value() {
ScalarValue::IntervalDayTime(Some(value)) => {
let (days, millis) = IntervalDayTimeType::to_parts(*value);
i64::from(days)
.checked_mul(86_400_000_000_000)
.and_then(|base| base.checked_add(i64::from(millis) * 1_000_000))
.ok_or_else(|| {
DataFusionError::Plan("window interval overflowed i64 nanoseconds".into())
})?
}
ScalarValue::IntervalMonthDayNano(Some(value)) => {
let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(*value);
if months != 0 {
return Err(DataFusionError::NotImplemented(
"datum-sql windowed aggregation supports fixed day/time intervals only, not month intervals"
.into(),
));
}
i64::from(days)
.checked_mul(86_400_000_000_000)
.and_then(|base| base.checked_add(nanos))
.ok_or_else(|| {
DataFusionError::Plan("window interval overflowed i64 nanoseconds".into())
})?
}
ScalarValue::DurationSecond(Some(value)) => {
value.checked_mul(1_000_000_000).ok_or_else(|| {
DataFusionError::Plan("window interval overflowed i64 nanoseconds".into())
})?
}
ScalarValue::DurationMillisecond(Some(value)) => {
value.checked_mul(1_000_000).ok_or_else(|| {
DataFusionError::Plan("window interval overflowed i64 nanoseconds".into())
})?
}
ScalarValue::DurationMicrosecond(Some(value)) => {
value.checked_mul(1_000).ok_or_else(|| {
DataFusionError::Plan("window interval overflowed i64 nanoseconds".into())
})?
}
ScalarValue::DurationNanosecond(Some(value)) => *value,
other => {
return Err(DataFusionError::Plan(format!(
"window interval must be a non-null fixed interval, found {other:?}"
)));
}
};
if nanos <= 0 {
return Err(DataFusionError::Plan(format!(
"window interval must be positive, found {nanos}ns"
)));
}
Ok(nanos)
}
fn validate_timestamp_output_type(data_type: &DataType, function: &str) -> Result<()> {
if matches!(data_type, DataType::Timestamp(_, _)) {
Ok(())
} else {
Err(DataFusionError::Plan(format!(
"{function} window grouping must return a timestamp, found {data_type:?}"
)))
}
}
fn floor_to_origin(value: i64, stride: i64, origin: i64) -> Result<i64> {
let diff = value.checked_sub(origin).ok_or_else(|| {
DataFusionError::Plan("window assignment overflowed i64 nanoseconds".into())
})?;
let mut delta = diff - (diff % stride);
if diff < 0 && stride > 1 && delta != diff {
delta -= stride;
}
origin
.checked_add(delta)
.ok_or_else(|| DataFusionError::Plan("window assignment overflowed i64 nanoseconds".into()))
}
fn checked_add_ns(left: i64, right: i64) -> Result<i64> {
left.checked_add(right)
.ok_or_else(|| DataFusionError::Plan("window end overflowed i64 nanoseconds".into()))
}
fn single_value_array(array: &dyn Array, row: usize) -> Result<ArrayRef> {
ScalarValue::iter_to_array([ScalarValue::try_from_array(array, row)?])
}
fn row_is_counted(values: &[ArrayRef], row: usize) -> bool {
values.iter().all(|array| !array.is_null(row))
}
fn single_int64_input<'a>(values: &'a [ArrayRef], function: &str) -> Result<&'a Int64Array> {
if values.len() != 1 {
return Err(DataFusionError::Internal(format!(
"fast {function} state expected one input array, found {}",
values.len()
)));
}
values[0]
.as_any()
.downcast_ref::<Int64Array>()
.ok_or_else(|| {
DataFusionError::Internal(format!("fast {function} state expected Int64Array"))
})
}
fn timestamp_ns_from_array(array: &ArrayRef, row: usize) -> Result<i64> {
match array.data_type() {
DataType::Timestamp(TimeUnit::Second, _) => timestamp_value_ns::<TimestampSecondArray>(
array.as_any().downcast_ref::<TimestampSecondArray>(),
row,
1_000_000_000,
),
DataType::Timestamp(TimeUnit::Millisecond, _) => {
timestamp_value_ns::<TimestampMillisecondArray>(
array.as_any().downcast_ref::<TimestampMillisecondArray>(),
row,
1_000_000,
)
}
DataType::Timestamp(TimeUnit::Microsecond, _) => {
timestamp_value_ns::<TimestampMicrosecondArray>(
array.as_any().downcast_ref::<TimestampMicrosecondArray>(),
row,
1_000,
)
}
DataType::Timestamp(TimeUnit::Nanosecond, _) => {
timestamp_value_ns::<TimestampNanosecondArray>(
array.as_any().downcast_ref::<TimestampNanosecondArray>(),
row,
1,
)
}
other => Err(DataFusionError::Plan(format!(
"window event-time expression must evaluate to Timestamp, found {other:?}"
))),
}
}
fn timestamp_value_ns<T>(array: Option<&T>, row: usize, multiplier: i64) -> Result<i64>
where
T: Array + TimestampArrayValue,
{
let array =
array.ok_or_else(|| DataFusionError::Internal("timestamp array type mismatch".into()))?;
if array.is_null(row) {
return Err(DataFusionError::Plan(format!(
"window event-time expression contains null at row {row}"
)));
}
array
.timestamp_value(row)
.checked_mul(multiplier)
.ok_or_else(|| DataFusionError::Plan("timestamp overflowed i64 nanoseconds".into()))
}
trait TimestampArrayValue {
fn timestamp_value(&self, row: usize) -> i64;
}
impl TimestampArrayValue for TimestampSecondArray {
fn timestamp_value(&self, row: usize) -> i64 {
self.value(row)
}
}
impl TimestampArrayValue for TimestampMillisecondArray {
fn timestamp_value(&self, row: usize) -> i64 {
self.value(row)
}
}
impl TimestampArrayValue for TimestampMicrosecondArray {
fn timestamp_value(&self, row: usize) -> i64 {
self.value(row)
}
}
impl TimestampArrayValue for TimestampNanosecondArray {
fn timestamp_value(&self, row: usize) -> i64 {
self.value(row)
}
}
fn timestamp_scalar_from_ns(timestamp_ns: i64, data_type: &DataType) -> Result<ScalarValue> {
match data_type {
DataType::Timestamp(TimeUnit::Second, timezone) => Ok(ScalarValue::TimestampSecond(
checked_unit_value(timestamp_ns, 1_000_000_000)?,
timezone.clone(),
)),
DataType::Timestamp(TimeUnit::Millisecond, timezone) => {
Ok(ScalarValue::TimestampMillisecond(
checked_unit_value(timestamp_ns, 1_000_000)?,
timezone.clone(),
))
}
DataType::Timestamp(TimeUnit::Microsecond, timezone) => {
Ok(ScalarValue::TimestampMicrosecond(
checked_unit_value(timestamp_ns, 1_000)?,
timezone.clone(),
))
}
DataType::Timestamp(TimeUnit::Nanosecond, timezone) => Ok(
ScalarValue::TimestampNanosecond(Some(timestamp_ns), timezone.clone()),
),
other => Err(DataFusionError::Plan(format!(
"window start output type must be Timestamp, found {other:?}"
))),
}
}
fn checked_unit_value(timestamp_ns: i64, divisor: i64) -> Result<Option<i64>> {
if timestamp_ns % divisor != 0 {
return Err(DataFusionError::Plan(format!(
"window start {timestamp_ns}ns cannot be represented in timestamp unit divisor {divisor}"
)));
}
Ok(Some(timestamp_ns / divisor))
}