From fe80d7fd47cd75a6e4568c80f6f0ae26d1fdbcf6 Mon Sep 17 00:00:00 2001 From: zzylol Date: Tue, 29 Sep 2026 20:00:34 +0000 Subject: [PATCH] feat: add native physical operators and shared DAG runtime Add the execution layer of `asap-physical-operators`: - `plan`, `runtime` and `sources`: physical DAGs, run context, bounded backpressure, cancellation, memory reservations and raw-source scans. - `expressions` and `operators`: relational/scalar operators, windows, temporal panes, current-series snapshots, and summary build, merge and readout. A readout is typed: a `SketchQuery`, or an `ExactReadout` whose counter lookback resolves to the run's evaluation range. Deserialized operators are validated before use. - `readout`: readouts over merged exact summary states. There is no physical planner yet; operators are built directly. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 2 + crates/asap-physical-operators/Cargo.toml | 2 + crates/asap-physical-operators/src/error.rs | 3 +- .../src/expressions/arithmetic.rs | 63 ++ .../src/expressions/mod.rs | 348 +++++++++++ .../src/expressions/planner.rs | 549 +++++++++++++++++ crates/asap-physical-operators/src/lib.rs | 11 +- .../src/operators/aggregate/mod.rs | 293 +++++++++ .../src/operators/aggregate/temporal.rs | 300 +++++++++ .../src/operators/aligned_binary.rs | 144 +++++ .../src/operators/common.rs | 75 +++ .../src/operators/current_series.rs | 135 +++++ .../src/operators/filter.rs | 39 ++ .../src/operators/joins/mod.rs | 219 +++++++ .../src/operators/limit.rs | 69 +++ .../src/operators/mod.rs | 316 ++++++++++ .../src/operators/panes.rs | 228 +++++++ .../src/operators/projection.rs | 56 ++ .../src/operators/sort.rs | 172 ++++++ .../src/operators/source.rs | 96 +++ .../src/operators/summary/mod.rs | 567 ++++++++++++++++++ .../src/operators/unchecked.rs | 140 +++++ .../src/operators/vector_binary.rs | 237 ++++++++ .../src/operators/vector_window.rs | 147 +++++ .../asap-physical-operators/src/plan/mod.rs | 160 +++++ .../src/plan/properties.rs | 32 + crates/asap-physical-operators/src/readout.rs | 111 ++++ .../src/runtime/batch_execution.rs | 210 +++++++ .../src/runtime/context.rs | 133 ++++ .../src/runtime/cooperative.rs | 40 ++ .../src/runtime/mod.rs | 272 +++++++++ .../src/runtime/tests.rs | 262 ++++++++ .../src/sources/memory.rs | 44 ++ .../src/sources/mod.rs | 187 ++++++ .../src/summary_kernels/exact.rs | 22 + crates/asap-physical-operators/src/values.rs | 26 +- .../tests/blocking_resources.rs | 242 ++++++++ .../tests/plan_properties.rs | 151 +++++ 38 files changed, 6100 insertions(+), 3 deletions(-) create mode 100644 crates/asap-physical-operators/src/expressions/arithmetic.rs create mode 100644 crates/asap-physical-operators/src/expressions/mod.rs create mode 100644 crates/asap-physical-operators/src/expressions/planner.rs create mode 100644 crates/asap-physical-operators/src/operators/aggregate/mod.rs create mode 100644 crates/asap-physical-operators/src/operators/aggregate/temporal.rs create mode 100644 crates/asap-physical-operators/src/operators/aligned_binary.rs create mode 100644 crates/asap-physical-operators/src/operators/common.rs create mode 100644 crates/asap-physical-operators/src/operators/current_series.rs create mode 100644 crates/asap-physical-operators/src/operators/filter.rs create mode 100644 crates/asap-physical-operators/src/operators/joins/mod.rs create mode 100644 crates/asap-physical-operators/src/operators/limit.rs create mode 100644 crates/asap-physical-operators/src/operators/mod.rs create mode 100644 crates/asap-physical-operators/src/operators/panes.rs create mode 100644 crates/asap-physical-operators/src/operators/projection.rs create mode 100644 crates/asap-physical-operators/src/operators/sort.rs create mode 100644 crates/asap-physical-operators/src/operators/source.rs create mode 100644 crates/asap-physical-operators/src/operators/summary/mod.rs create mode 100644 crates/asap-physical-operators/src/operators/unchecked.rs create mode 100644 crates/asap-physical-operators/src/operators/vector_binary.rs create mode 100644 crates/asap-physical-operators/src/operators/vector_window.rs create mode 100644 crates/asap-physical-operators/src/plan/mod.rs create mode 100644 crates/asap-physical-operators/src/plan/properties.rs create mode 100644 crates/asap-physical-operators/src/readout.rs create mode 100644 crates/asap-physical-operators/src/runtime/batch_execution.rs create mode 100644 crates/asap-physical-operators/src/runtime/context.rs create mode 100644 crates/asap-physical-operators/src/runtime/cooperative.rs create mode 100644 crates/asap-physical-operators/src/runtime/mod.rs create mode 100644 crates/asap-physical-operators/src/runtime/tests.rs create mode 100644 crates/asap-physical-operators/src/sources/memory.rs create mode 100644 crates/asap-physical-operators/src/sources/mod.rs create mode 100644 crates/asap-physical-operators/tests/blocking_resources.rs create mode 100644 crates/asap-physical-operators/tests/plan_properties.rs diff --git a/Cargo.lock b/Cargo.lock index 96a1f63a..a063dd78 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -382,7 +382,9 @@ dependencies = [ "asap-frontend-promql", "asap-types", "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369)", + "futures", "serde", + "serde_json", "thiserror 2.0.18", "tracing", ] diff --git a/crates/asap-physical-operators/Cargo.toml b/crates/asap-physical-operators/Cargo.toml index 51dd6a04..74c172e3 100644 --- a/crates/asap-physical-operators/Cargo.toml +++ b/crates/asap-physical-operators/Cargo.toml @@ -4,9 +4,11 @@ version = "0.1.0" edition = "2021" [dependencies] +futures = "0.3" planner-types = { package = "asap-types", path = "../types" } asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "5f03ccbd798ed5fec62bdd839bcb331123cab369" } serde = { version = "1", features = ["derive", "rc"] } +serde_json = "1" tracing = "0.1" thiserror = "2" diff --git a/crates/asap-physical-operators/src/error.rs b/crates/asap-physical-operators/src/error.rs index 06ce16e9..afee33d8 100644 --- a/crates/asap-physical-operators/src/error.rs +++ b/crates/asap-physical-operators/src/error.rs @@ -1,3 +1,4 @@ +use crate::plan::NodeId; #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum Error { #[error("invalid DAG: {0}")] @@ -6,7 +7,7 @@ pub enum Error { Operator(String), #[error("node {node} ({operation}) failed: {source}")] AtNode { - node: u64, + node: NodeId, operation: String, source: Box, }, diff --git a/crates/asap-physical-operators/src/expressions/arithmetic.rs b/crates/asap-physical-operators/src/expressions/arithmetic.rs new file mode 100644 index 00000000..30e277d4 --- /dev/null +++ b/crates/asap-physical-operators/src/expressions/arithmetic.rs @@ -0,0 +1,63 @@ +//! Float64 arithmetic shared by ASAP execution engines. +//! Preserve IEEE non-finite results; callers own their output policies. + +pub fn evaluate_float64_arithmetic( + operator: &planner_types::pre_asap::ArithmeticOpKind, + left: f64, + right: f64, +) -> f64 { + use planner_types::pre_asap::ArithmeticOpKind::*; + match operator { + Add => left + right, + Sub => left - right, + Mul => left * right, + Div => left / right, + Mod => left % right, + Pow => left.powf(right), + Atan2 => left.atan2(right), + } +} + +/// Execute the Planner binary contract after a deployment has resolved matching rows. +pub fn evaluate_binary( + operator: &planner_types::post_asap::BinaryOperator, + left: f64, + right: f64, +) -> Result { + use crate::{values::Value, Error}; + use planner_types::pre_asap::{ArithmeticOpKind, BinaryOpKind, CompareOpKind}; + let invalid = + || Error::Invalid("unsupported binary operation or invalid checked-division domain".into()); + if operator.vector_match.is_some() { + return Err(invalid()); + } + if operator.checked_relative_division || operator.checked_finite_division { + if operator.kind != BinaryOpKind::Arithmetic(ArithmeticOpKind::Div) + || !left.is_finite() + || !right.is_finite() + || right == 0. + { + return Err(invalid()); + } + let value = left / right; + if !value.is_finite() || (operator.checked_relative_division && !value.is_normal()) { + return Err(invalid()); + } + return Ok(Value::Float64(value)); + } + Ok(match operator.kind { + BinaryOpKind::Arithmetic(ref op) => { + Value::Float64(evaluate_float64_arithmetic(op, left, right)) + } + BinaryOpKind::Compare(ref op) => Value::Bool(match op { + CompareOpKind::Eq => left == right, + CompareOpKind::Ne => left != right, + CompareOpKind::Lt => left < right, + CompareOpKind::Le => left <= right, + CompareOpKind::Gt => left > right, + CompareOpKind::Ge => left >= right, + _ => return Err(invalid()), + }), + _ => return Err(invalid()), + }) +} diff --git a/crates/asap-physical-operators/src/expressions/mod.rs b/crates/asap-physical-operators/src/expressions/mod.rs new file mode 100644 index 00000000..2916cca7 --- /dev/null +++ b/crates/asap-physical-operators/src/expressions/mod.rs @@ -0,0 +1,348 @@ +//! Scalar semantics and typed expression binding. Planner expressions enter through CompiledExpression. +use crate::{ + values::{plain, Schema, Value}, + Error, +}; +use planner_types::pre_asap::{ArithmeticOpKind, DataType}; +pub mod arithmetic; +mod planner; +pub use planner::CompiledExpression; +#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)] +pub enum Expression { + Binary { + operator: planner_types::post_asap::BinaryOperator, + left: Box, + right: Box, + }, + Planner(Box), + Column(usize), + ExactFloat64(usize), + FiniteFloat64(Box), + LabelSet { + column: usize, + labels: Vec, + without: bool, + }, + Literal { + value: Value, + dtype: DataType, + }, + Negate(Box), + Arithmetic { + op: ArithmeticOpKind, + left: Box, + right: Box, + }, + Equal(Box, Box), + Less(Box, Box), + And(Box, Box), + Or(Box, Box), + Not(Box), + IsNull(Box), +} +impl Expression { + pub fn planner(expression: crate::expressions::CompiledExpression) -> Self { + Self::Planner(Box::new(expression)) + } + pub(crate) fn dtype(&self, input: &Schema) -> Result<(DataType, bool), Error> { + use Expression::*; + match self { + Binary { + operator, + left, + right, + } => { + use planner_types::pre_asap::{BinaryOpKind, CompareOpKind}; + let (a, n) = left.dtype(input)?; + let (b, m) = right.dtype(input)?; + if a != DataType::Float64 || b != a || operator.vector_match.is_some() { + return Err(invalid( + "binary expression requires resolved Float64 operands", + )); + } + if (operator.checked_relative_division || operator.checked_finite_division) + && operator.kind != BinaryOpKind::Arithmetic(ArithmeticOpKind::Div) + { + return Err(invalid("checked division contract on non-division")); + } + let dtype = match operator.kind { + BinaryOpKind::Arithmetic(_) => DataType::Float64, + BinaryOpKind::Compare( + CompareOpKind::Eq + | CompareOpKind::Ne + | CompareOpKind::Lt + | CompareOpKind::Le + | CompareOpKind::Gt + | CompareOpKind::Ge, + ) => DataType::Bool, + _ => return Err(invalid("unsupported binary operation")), + }; + Ok((dtype, n || m)) + } + Planner(expression) => { + expression.validate_input(input)?; + Ok(expression.dtype()) + } + FiniteFloat64(expression) => { + if expression.dtype(input)? != (DataType::Float64, false) { + return Err(invalid("finite update requires non-null Float64")); + } + Ok((DataType::Float64, false)) + } + ExactFloat64(column) => { + let (dtype, nullable) = plain(input, *column)?; + if nullable || !matches!(dtype, DataType::Int64 | DataType::Float64) { + return Err(invalid( + "exact Float64 conversion requires non-null numeric input", + )); + } + Ok((DataType::Float64, false)) + } + LabelSet { column, labels, .. } => { + let (dtype, nullable) = plain(input, *column)?; + let expected = DataType::Map { + key: Box::new(DataType::Utf8), + value: Box::new(DataType::Utf8), + value_nullable: false, + }; + if dtype != &expected + || nullable + || labels + .iter() + .collect::>() + .len() + != labels.len() + { + return Err(invalid( + "label projection requires a non-null Utf8 map and unique label names", + )); + } + Ok((expected, false)) + } + Column(i) => { + let (t, n) = plain(input, *i)?; + Ok((t.clone(), n)) + } + Literal { value, dtype } => { + if value.matches(dtype, true) { + Ok((dtype.clone(), matches!(value, Value::Null))) + } else { + Err(invalid("literal type mismatch")) + } + } + Negate(v) => { + let (t, n) = v.dtype(input)?; + if matches!(t, DataType::Int64 | DataType::Float64) { + Ok((t, n)) + } else { + Err(invalid("numeric negation required")) + } + } + Arithmetic { op, left, right } => { + let (a, n) = left.dtype(input)?; + let (b, m) = right.dtype(input)?; + if a == b + && matches!(a, DataType::Int64 | DataType::Float64) + && !(a == DataType::Int64 && *op == ArithmeticOpKind::Atan2) + { + Ok((a, n || m)) + } else { + Err(invalid("arithmetic requires matching numeric types")) + } + } + Equal(a, b) | Less(a, b) => { + let (a, n) = a.dtype(input)?; + let (b, m) = b.dtype(input)?; + if a == b && ordered(&a) { + Ok((DataType::Bool, n || m)) + } else { + Err(invalid("comparison requires matching ordered types")) + } + } + And(a, b) | Or(a, b) => { + let (a, n) = a.dtype(input)?; + let (b, m) = b.dtype(input)?; + if a == DataType::Bool && b == DataType::Bool { + Ok((DataType::Bool, n || m)) + } else { + Err(invalid("boolean operands required")) + } + } + Not(v) => { + let (t, n) = v.dtype(input)?; + if t == DataType::Bool { + Ok((t, n)) + } else { + Err(invalid("boolean operand required")) + } + } + IsNull(v) => { + v.dtype(input)?; + Ok((DataType::Bool, false)) + } + } + } + pub(crate) fn evaluate(&self, row: &[Value]) -> Result { + use Expression::*; + Ok(match self { + FiniteFloat64(expression) => match expression.evaluate(row)? { + Value::Float64(value) if value.is_finite() => Value::Float64(value), + _ => return Err(invalid("summary update must be finite")), + }, + ExactFloat64(column) => match row[*column] { + Value::Float64(value) => Value::Float64(value), + Value::Int64(value) if value.unsigned_abs() <= (1u64 << 53) => { + Value::Float64(value as f64) + } + _ => { + return Err(invalid( + "numeric result cannot be represented exactly as Float64", + )) + } + }, + LabelSet { + column, + labels, + without, + } => { + let Value::Map(entries) = &row[*column] else { + return Err(invalid("label projection requires a map")); + }; + let mut selected = std::collections::BTreeMap::new(); + let mut seen = std::collections::BTreeSet::new(); + for (key, value) in entries.iter() { + let (Value::Utf8(key), Value::Utf8(value)) = (key, value) else { + return Err(invalid("label projection requires Utf8 entries")); + }; + if !seen.insert(key.clone()) { + return Err(invalid("duplicate label name")); + } + let keep = if *without { + key.as_ref() != "__name__" + && !labels.iter().any(|label| label.as_str() == key.as_ref()) + } else { + labels.iter().any(|label| label.as_str() == key.as_ref()) + }; + if keep && !value.is_empty() { + selected.insert(key.clone(), value.clone()); + } + } + Value::Map( + selected + .into_iter() + .map(|(k, v)| (Value::Utf8(k), Value::Utf8(v))) + .collect::>() + .into(), + ) + } + Binary { + operator, + left, + right, + } => { + let (a, b) = (left.evaluate(row)?, right.evaluate(row)?); + if matches!(a, Value::Null) || matches!(b, Value::Null) { + Value::Null + } else { + let (Value::Float64(a), Value::Float64(b)) = (a, b) else { + return Err(invalid("binary value schema mismatch")); + }; + arithmetic::evaluate_binary(operator, a, b)? + } + } + Planner(expression) => expression.evaluate(row)?, + Column(i) => row[*i].clone(), + Literal { value, .. } => value.clone(), + Negate(v) => match v.evaluate(row)? { + Value::Int64(v) => Value::Int64( + v.checked_neg() + .ok_or_else(|| invalid("integer negation overflow"))?, + ), + Value::Float64(v) => Value::Float64(-v), + Value::Null => Value::Null, + _ => return Err(invalid("numeric negation required")), + }, + Arithmetic { op, left, right } => { + numeric(op, left.evaluate(row)?, right.evaluate(row)?)? + } + Equal(a, b) | Less(a, b) => { + let (a, b) = (a.evaluate(row)?, b.evaluate(row)?); + if matches!(a, Value::Null) || matches!(b, Value::Null) { + Value::Null + } else if matches!((&a,&b),(Value::Float64(a),Value::Float64(b)) if a.is_nan() || b.is_nan()) + { + Value::Bool(false) + } else { + let c = a.compare(&b)?; + Value::Bool(if matches!(self, Equal(..)) { + c.is_eq() + } else { + c.is_lt() + }) + } + } + And(a, b) | Or(a, b) => { + let (a, b) = (a.evaluate(row)?, b.evaluate(row)?); + match (a, b, matches!(self, And(..))) { + (Value::Bool(false), _, true) | (_, Value::Bool(false), true) => { + Value::Bool(false) + } + (Value::Bool(true), _, false) | (_, Value::Bool(true), false) => { + Value::Bool(true) + } + (Value::Null, _, _) | (_, Value::Null, _) => Value::Null, + (Value::Bool(a), Value::Bool(b), true) => Value::Bool(a && b), + (Value::Bool(a), Value::Bool(b), false) => Value::Bool(a || b), + _ => return Err(invalid("boolean operands required")), + } + } + Not(v) => match v.evaluate(row)? { + Value::Bool(v) => Value::Bool(!v), + Value::Null => Value::Null, + _ => return Err(invalid("boolean operand required")), + }, + IsNull(v) => Value::Bool(matches!(v.evaluate(row)?, Value::Null)), + }) + } +} +pub(crate) fn ordered(dtype: &DataType) -> bool { + if let DataType::Map { key, value, .. } = dtype { + return ordered(key) && ordered(value); + } + matches!( + dtype, + DataType::Null + | DataType::Int64 + | DataType::Float64 + | DataType::Utf8 + | DataType::Bool + | DataType::Timestamp + | DataType::Date + ) +} +pub(crate) fn numeric(op: &ArithmeticOpKind, a: Value, b: Value) -> Result { + use ArithmeticOpKind::*; + Ok(match (a, b) { + (Value::Null, _) | (_, Value::Null) => Value::Null, + (Value::Float64(a), Value::Float64(b)) => { + Value::Float64(arithmetic::evaluate_float64_arithmetic(op, a, b)) + } + (Value::Int64(a), Value::Int64(b)) => Value::Int64( + match op { + Add => a.checked_add(b), + Sub => a.checked_sub(b), + Mul => a.checked_mul(b), + Div => a.checked_div(b), + Mod => a.checked_rem(b), + Pow => u32::try_from(b).ok().and_then(|b| a.checked_pow(b)), + Atan2 => None, + } + .ok_or_else(|| invalid("invalid integer arithmetic or overflow"))?, + ), + _ => return Err(invalid("arithmetic type mismatch")), + }) +} + +fn invalid(message: &str) -> Error { + Error::Invalid(message.into()) +} diff --git a/crates/asap-physical-operators/src/expressions/planner.rs b/crates/asap-physical-operators/src/expressions/planner.rs new file mode 100644 index 00000000..2130a2f7 --- /dev/null +++ b/crates/asap-physical-operators/src/expressions/planner.rs @@ -0,0 +1,549 @@ +//! Planner scalar expressions evaluated over native typed rows. +use crate::{ + values::{Schema, Value}, + Error, +}; +use planner_types::pre_asap::{ArithmeticOpKind, CompareOpKind, DataType, QueryExpr, ScalarValue}; +use std::{cmp::Ordering, sync::Arc}; + +pub(super) fn evaluate( + expr: &QueryExpr, + row: &[Value], + schema: &planner_types::pre_asap::Schema, +) -> Result { + match expr { + QueryExpr::Column(index) => row.get(*index).cloned().ok_or(Error::Invalid(format!( + "column {index} outside row width {}", + row.len() + ))), + QueryExpr::Literal(value) => Ok(match value { + ScalarValue::Interval { + months, + days, + nanos, + } => Value::Interval { + months: *months, + days: *days, + nanos: *nanos, + }, + ScalarValue::Int64(value) => Value::Int64(*value), + ScalarValue::Float64(value) => Value::Float64(*value), + ScalarValue::Utf8(value) => Value::Utf8(value.clone().into()), + ScalarValue::Boolean(value) => Value::Bool(*value), + ScalarValue::Null => Value::Null, + }), + QueryExpr::Compare { left, op, right } => { + let left = evaluate(left, row, schema)?; + let right = evaluate(right, row, schema)?; + compare(op, left, right) + } + QueryExpr::Arithmetic { op, left, right } => arithmetic( + op, + evaluate(left, row, schema)?, + evaluate(right, row, schema)?, + ), + QueryExpr::BoolAnd(parts) | QueryExpr::BoolOr(parts) => { + let and = matches!(expr, QueryExpr::BoolAnd(_)); + let mut null = false; + for part in parts { + match evaluate(part, row, schema)? { + Value::Bool(value) if value != and => return Ok(Value::Bool(value)), + Value::Bool(_) => {} + Value::Null => null = true, + _ => return Err(Error::Invalid("boolean predicate required".into())), + } + } + Ok(if null { Value::Null } else { Value::Bool(and) }) + } + QueryExpr::Not(value) => match evaluate(value, row, schema)? { + Value::Bool(value) => Ok(Value::Bool(!value)), + Value::Null => Ok(Value::Null), + _ => Err(Error::Invalid("boolean predicate required".into())), + }, + QueryExpr::IsNull(value) => Ok(Value::Bool(matches!( + evaluate(value, row, schema)?, + Value::Null + ))), + QueryExpr::IsNotNull(value) => Ok(Value::Bool(!matches!( + evaluate(value, row, schema)?, + Value::Null + ))), + QueryExpr::FunctionCall { name, args } => { + use planner_types::pre_asap::scalar_signature::MapScalarFunction; + if name.eq_ignore_ascii_case("asap_struct_field") { + expr.scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))?; + let DataType::Struct { fields } = args[0] + .scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))? + .0 + else { + unreachable!() + }; + let offset = match &args[1] { + QueryExpr::Literal(ScalarValue::Int64(index)) => { + usize::try_from(index - 1).ok() + } + QueryExpr::Literal(ScalarValue::Utf8(name)) => { + fields.iter().position(|field| &field.name == name) + } + _ => None, + } + .ok_or_else(|| Error::Invalid("struct field selector".into()))?; + let Value::Struct(values) = evaluate(&args[0], row, schema)? else { + return Err(Error::Invalid("struct field input".into())); + }; + return values + .get(offset) + .cloned() + .ok_or_else(|| Error::Invalid("struct field value".into())); + } + if name.eq_ignore_ascii_case("asap_element_access") { + let (output_type, _) = expr + .scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))?; + if let DataType::List { element } = args[0] + .scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))? + .0 + { + let Value::List(values) = evaluate(&args[0], row, schema)? else { + return Err(Error::Invalid("array access input".into())); + }; + let index = match evaluate(&args[1], row, schema)? { + Value::Null => return Ok(Value::Null), + Value::Int64(index) => index, + _ => return Err(Error::Invalid("array access index".into())), + }; + let offset = if index > 0 { + usize::try_from(index - 1).ok() + } else if index < 0 { + usize::try_from(index.unsigned_abs()) + .ok() + .and_then(|distance| values.len().checked_sub(distance)) + } else { + None + }; + return match offset.and_then(|offset| values.get(offset)) { + Some(value) => Ok(value.clone()), + None => default_collection_element(&output_type, element.nullable), + }; + } + } + let function = (if name.eq_ignore_ascii_case("asap_element_access") { + Some(MapScalarFunction::Access) + } else { + MapScalarFunction::from_name(name) + }) + .ok_or_else(|| Error::Invalid(format!("scalar function {name}")))?; + expr.scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))?; + let values = args + .iter() + .map(|arg| evaluate(arg, row, schema)) + .collect::, _>>()?; + match function { + MapScalarFunction::Construct => { + let mut values = values.into_iter(); + let mut entries = Vec::new(); + while let Some(key) = values.next() { + if !matches!(key, Value::Int64(_) | Value::Utf8(_) | Value::Bool(_)) { + return Err(Error::Invalid("map key value type".into())); + } + entries.push(( + key, + values + .next() + .ok_or_else(|| Error::Invalid("odd map argument count".into()))?, + )); + } + Ok(Value::Map(entries.into())) + } + MapScalarFunction::Concat => { + let mut entries = Vec::new(); + for value in values { + let Value::Map(next) = value else { + return Err(Error::Invalid("map concat argument".into())); + }; + entries.extend(next.iter().cloned()); + } + Ok(Value::Map(entries.into())) + } + MapScalarFunction::Access => { + let [Value::Map(entries), key] = values.as_slice() else { + return Err(Error::Invalid("map access arguments".into())); + }; + if matches!(key, Value::Null) { + return Ok(Value::Null); + } + if !matches!(key, Value::Int64(_) | Value::Utf8(_) | Value::Bool(_)) { + return Err(Error::Invalid("map lookup key type".into())); + } + if let Some((_, value)) = entries + .iter() + .find(|(candidate, _)| cell_cmp(candidate, key) == Some(Ordering::Equal)) + { + return Ok(value.clone()); + } + let ( + DataType::Map { + value, + value_nullable, + .. + }, + _, + ) = args[0] + .scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))? + else { + unreachable!() + }; + default_collection_element(&value, value_nullable) + } + } + } + other => Err(Error::Invalid(format!("scalar expression {other:?}"))), + } +} + +fn default_collection_element(dtype: &DataType, nullable: bool) -> Result { + if nullable { + return Ok(Value::Null); + } + Ok(match dtype { + DataType::Interval | DataType::Date => { + return Err(Error::Invalid("temporal value transport".into())) + } + DataType::Null => Value::Null, + DataType::Int64 => Value::Int64(0), + DataType::Float64 => Value::Float64(0.0), + DataType::Utf8 => Value::Utf8("".into()), + DataType::Bool => Value::Bool(false), + DataType::Map { .. } => Value::Map(Arc::from([])), + DataType::List { .. } => Value::List(Arc::from([])), + DataType::Struct { fields } => Value::Struct( + fields + .iter() + .map(|field| default_collection_element(&field.dtype, field.nullable)) + .collect::, _>>()? + .into(), + ), + _ => { + return Err(Error::Invalid( + "collection missing-element default type".into(), + )) + } + }) +} + +fn compare(op: &CompareOpKind, left: Value, right: Value) -> Result { + if matches!(left, Value::Null) || matches!(right, Value::Null) { + return Ok(Value::Null); + } + // NaN is unordered, not a type mismatch. Match the native scalar path. + if matches!(&left, Value::Float64(v) if v.is_nan()) + || matches!(&right, Value::Float64(v) if v.is_nan()) + { + return match op { + CompareOpKind::Ne => Ok(Value::Bool(true)), + CompareOpKind::Eq + | CompareOpKind::Lt + | CompareOpKind::Le + | CompareOpKind::Gt + | CompareOpKind::Ge => Ok(Value::Bool(false)), + _ => Err(Error::Invalid(format!("comparison {op:?}"))), + }; + } + let ordering = cell_cmp(&left, &right) + .ok_or_else(|| Error::Invalid("comparison of incompatible values".into()))?; + let value = match op { + CompareOpKind::Eq => ordering == Ordering::Equal, + CompareOpKind::Ne => ordering != Ordering::Equal, + CompareOpKind::Lt => ordering == Ordering::Less, + CompareOpKind::Le => ordering != Ordering::Greater, + CompareOpKind::Gt => ordering == Ordering::Greater, + CompareOpKind::Ge => ordering != Ordering::Less, + _ => return Err(Error::Invalid(format!("comparison {op:?}"))), + }; + Ok(Value::Bool(value)) +} + +fn arithmetic(op: &ArithmeticOpKind, left: Value, right: Value) -> Result { + let (left, right) = match (left, right) { + (Value::Int64(a), Value::Float64(b)) => (Value::Float64(a as f64), Value::Float64(b)), + (Value::Float64(a), Value::Int64(b)) => (Value::Float64(a), Value::Float64(b as f64)), + pair => pair, + }; + super::numeric(op, left, right) +} + +fn integer_float_cmp(integer: i64, float: f64) -> Option { + if float.is_nan() { + return None; + } + // These bounds are powers of two, exactly representable as Float64. + if float >= 9_223_372_036_854_775_808.0 { + return Some(Ordering::Less); + } + if float < -9_223_372_036_854_775_808.0 { + return Some(Ordering::Greater); + } + let integral = float as i64; + match integer.cmp(&integral) { + Ordering::Equal => 0.0_f64.partial_cmp(&float.fract()), + other => Some(other), + } +} + +fn cell_cmp(left: &Value, right: &Value) -> Option { + match (left, right) { + (Value::Int64(left), Value::Int64(right)) => Some(left.cmp(right)), + (Value::Float64(left), Value::Float64(right)) => left.partial_cmp(right), + (Value::Int64(left), Value::Float64(right)) => integer_float_cmp(*left, *right), + (Value::Float64(left), Value::Int64(right)) => { + integer_float_cmp(*right, *left).map(Ordering::reverse) + } + (Value::Utf8(left), Value::Utf8(right)) => Some(left.cmp(right)), + (Value::Bool(left), Value::Bool(right)) => Some(left.cmp(right)), + (Value::Timestamp(left), Value::Timestamp(right)) => Some(left.cmp(right)), + (Value::Map(left), Value::Map(right)) => { + for ((left_key, left_value), (right_key, right_value)) in left.iter().zip(right.iter()) + { + let order = cell_cmp(left_key, right_key)?; + if order != Ordering::Equal { + return Some(order); + } + let order = match (left_value, right_value) { + (Value::Null, Value::Null) => Ordering::Equal, + (Value::Null, _) => Ordering::Greater, + (_, Value::Null) => Ordering::Less, + _ => cell_cmp(left_value, right_value)?, + }; + if order != Ordering::Equal { + return Some(order); + } + } + Some(left.len().cmp(&right.len())) + } + _ => None, + } +} + +#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)] +pub struct CompiledExpression { + expression: QueryExpr, + schema: planner_types::pre_asap::Schema, + output: (DataType, bool), +} +impl CompiledExpression { + pub(crate) fn expression(&self) -> &QueryExpr { + &self.expression + } + + pub fn compile(expression: &QueryExpr, input: &Schema) -> Result { + let schema = input + .fields + .iter() + .map(|field| { + let planner_types::post_asap::SummaryFamilyType::Plain(dtype) = &field.dtype else { + return Err(Error::Invalid( + "scalar expression cannot consume opaque summary state".into(), + )); + }; + Ok(planner_types::pre_asap::Column::new( + field.name.clone(), + dtype.clone(), + field.nullable, + )) + }) + .collect::, Error>>()?; + let schema = planner_types::pre_asap::Schema::new(schema); + validate(expression, &schema)?; + let output = expression + .scalar_type(&schema) + .map_err(|e| Error::Invalid(e.to_string()))?; + Ok(Self { + expression: expression.clone(), + schema, + output, + }) + } + pub(crate) fn dtype(&self) -> (DataType, bool) { + self.output.clone() + } + pub(crate) fn validate_input(&self, input: &Schema) -> Result<(), Error> { + let checked = Self::compile(&self.expression, input)?; + if checked.output != self.output { + return Err(Error::Invalid( + "persisted expression type differs from its semantics".into(), + )); + } + if input.fields.len() != self.schema.columns.len() + || input + .fields + .iter() + .zip(&self.schema.columns) + .any(|(field, column)| { + field.dtype + != planner_types::post_asap::SummaryFamilyType::Plain(column.dtype.clone()) + || field.nullable != column.nullable + }) + { + return Err(Error::Invalid( + "expression input differs from its bound schema".into(), + )); + } + Ok(()) + } + /// Evaluate a row under the same typed schema used when binding the expression. + pub fn evaluate(&self, row: &[Value]) -> Result { + if row.len() != self.schema.columns.len() + || row + .iter() + .zip(&self.schema.columns) + .any(|(value, column)| !value.matches(&column.dtype, column.nullable)) + { + return Err(Error::Invalid( + "expression input differs from its bound schema".into(), + )); + } + evaluate(&self.expression, row, &self.schema) + } +} +fn validate(expr: &QueryExpr, schema: &planner_types::pre_asap::Schema) -> Result<(), Error> { + let invalid = || Error::Invalid(format!("unsupported scalar expression: {expr:?}")); + expr.scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))?; + match expr { + QueryExpr::Column(_) | QueryExpr::Literal(_) => Ok(()), + QueryExpr::Arithmetic { left, right, .. } => { + for value in [left, right] { + validate(value, schema)?; + if !matches!( + value + .scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))? + .0, + DataType::Int64 | DataType::Float64 | DataType::Null + ) { + return Err(invalid()); + } + } + Ok(()) + } + QueryExpr::Compare { left, right, op } => { + if !matches!( + op, + CompareOpKind::Eq + | CompareOpKind::Ne + | CompareOpKind::Lt + | CompareOpKind::Le + | CompareOpKind::Gt + | CompareOpKind::Ge + ) { + return Err(invalid()); + } + validate(left, schema)?; + validate(right, schema)?; + let (a, _) = left + .scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))?; + let (b, _) = right + .scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))?; + fn comparable(dtype: &DataType) -> bool { + match dtype { + DataType::Null + | DataType::Int64 + | DataType::Float64 + | DataType::Utf8 + | DataType::Bool + | DataType::Timestamp => true, + DataType::Map { key, value, .. } => comparable(key) && comparable(value), + _ => false, + } + } + let numeric = |dtype: &DataType| matches!(dtype, DataType::Int64 | DataType::Float64); + if !comparable(&a) + || !comparable(&b) + || (a != b + && !matches!(a, DataType::Null) + && !matches!(b, DataType::Null) + && !(numeric(&a) && numeric(&b))) + { + return Err(invalid()); + } + Ok(()) + } + QueryExpr::FunctionCall { name, args } => { + if name != "asap_struct_field" + && name != "asap_element_access" + && planner_types::pre_asap::scalar_signature::MapScalarFunction::from_name(name) + .is_none() + { + return Err(invalid()); + } + for arg in args { + validate(arg, schema)?; + } + Ok(()) + } + QueryExpr::BoolAnd(parts) | QueryExpr::BoolOr(parts) => { + for part in parts { + validate(part, schema)?; + if !matches!( + part.scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))? + .0, + DataType::Bool | DataType::Null + ) { + return Err(invalid()); + } + } + Ok(()) + } + QueryExpr::Not(value) => { + validate(value, schema)?; + if !matches!( + value + .scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))? + .0, + DataType::Bool | DataType::Null + ) { + return Err(invalid()); + } + Ok(()) + } + QueryExpr::IsNull(value) | QueryExpr::IsNotNull(value) => validate(value, schema), + _ => Err(invalid()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn mixed_comparison_preserves_integer_precision_and_boundaries() { + assert_eq!( + integer_float_cmp(9_007_199_254_740_993, 9_007_199_254_740_992.0), + Some(Ordering::Greater) + ); + assert_eq!( + integer_float_cmp(i64::MAX, 9_223_372_036_854_775_808.0), + Some(Ordering::Less) + ); + assert_eq!( + integer_float_cmp(i64::MIN, -9_223_372_036_854_775_808.0), + Some(Ordering::Equal) + ); + assert_eq!(integer_float_cmp(-1, -1.5), Some(Ordering::Greater)); + assert_eq!(integer_float_cmp(1, 1.5), Some(Ordering::Less)); + assert_eq!(integer_float_cmp(0, f64::INFINITY), Some(Ordering::Less)); + assert_eq!( + integer_float_cmp(0, f64::NEG_INFINITY), + Some(Ordering::Greater) + ); + assert_eq!(integer_float_cmp(0, f64::NAN), None); + } +} diff --git a/crates/asap-physical-operators/src/lib.rs b/crates/asap-physical-operators/src/lib.rs index 5026cb6d..0835159c 100644 --- a/crates/asap-physical-operators/src/lib.rs +++ b/crates/asap-physical-operators/src/lib.rs @@ -1,4 +1,5 @@ -//! Shared physical operators. This layer holds summary kernels and typed values. +//! Shared physical operators: summary kernels, typed values, native operators +//! and the DAG runtime. pub mod key_by_label_values; pub mod measurement; @@ -20,3 +21,11 @@ pub use planner_types as planner; mod error; pub use error::Error; pub mod values; + +pub use expressions::arithmetic; +pub mod expressions; +pub mod operators; +pub mod plan; +pub mod readout; +pub mod runtime; +pub mod sources; diff --git a/crates/asap-physical-operators/src/operators/aggregate/mod.rs b/crates/asap-physical-operators/src/operators/aggregate/mod.rs new file mode 100644 index 00000000..71e7134c --- /dev/null +++ b/crates/asap-physical-operators/src/operators/aggregate/mod.rs @@ -0,0 +1,293 @@ +use super::*; +impl Operator { + pub fn aggregate( + input: Schema, + groups: Vec, + measures: Vec<(String, Reduction)>, + ) -> Result { + validate_groups(&input, &groups)?; + let mut fields = groups + .iter() + .map(|&i| input.fields[i].clone()) + .collect::>(); + for (name, reduction) in &measures { + let (t, n) = match reduction { + Reduction::Count => (DataType::Int64, false), + Reduction::Sum(i) | Reduction::Avg(i) => { + let (t, _) = plain(&input, *i)?; + if !matches!(t, DataType::Int64 | DataType::Float64) { + return Err(invalid("numeric aggregate input required")); + } + ( + if matches!(reduction, Reduction::Avg(_)) { + DataType::Float64 + } else { + t.clone() + }, + false, + ) + } + Reduction::Min(i) | Reduction::Max(i) => { + let (t, nullable) = plain(&input, *i)?; + if !ordered(t) { + return Err(invalid("ordered aggregate input required")); + } + (t.clone(), nullable || groups.is_empty()) + } + }; + fields.push(result_field(name, t, n)); + } + Ok(Self { + kind: Kind::Aggregate { + groups, + measures: measures.into_iter().map(|(_, r)| r).collect(), + }, + inputs: vec![input], + output: schema(fields), + }) + } + pub fn window( + input: Schema, + intent: planner_types::pre_asap::AggIntent, + coordinate: usize, + value: usize, + groups: Vec, + window: Option<(i64, i64)>, + ) -> Result { + use planner_types::pre_asap::AggIntent; + validate_groups(&input, &groups)?; + let histogram = matches!(intent, AggIntent::HistogramQuantile { .. }); + if !matches!( + intent, + AggIntent::Rate + | AggIntent::Increase + | AggIntent::Count { .. } + | AggIntent::Sum { col: None } + | AggIntent::Avg { col: None } + | AggIntent::Min { col: None } + | AggIntent::Max { col: None } + | AggIntent::HistogramQuantile { .. } + ) { + return Err(invalid( + "unsupported temporal intent or unresolved value column", + )); + } + let coordinate_type = if histogram { + DataType::Float64 + } else { + DataType::Timestamp + }; + if plain(&input, coordinate)? != (&coordinate_type, false) + || plain(&input, value)? != (&DataType::Float64, false) + { + return Err(invalid("window coordinate/value schema mismatch")); + } + if (!histogram && !matches!(window, Some((start, end)) if start < end)) + || (histogram && window.is_some()) + { + return Err(invalid("invalid temporal window")); + } + let mut fields = groups + .iter() + .map(|i| input.fields[*i].clone()) + .collect::>(); + fields.push(result_field( + "value", + if matches!(intent, AggIntent::Count { .. }) { + DataType::Int64 + } else { + DataType::Float64 + }, + false, + )); + Ok(Self { + kind: Kind::Window { + intent: Box::new(intent), + coordinate, + value, + groups, + window, + }, + inputs: vec![input], + output: schema(fields), + }) + } +} +#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)] +pub enum Reduction { + Count, + Sum(usize), + Avg(usize), + Min(usize), + Max(usize), +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + Ok(futures::stream::once(async move { + let (rows, _memory) = collect_rows(input, &context).await?; + let result = match &operator.kind { + Kind::Window { + intent, + coordinate, + value, + groups, + window, + } => { + crate::operators::aggregate::temporal::reduce( + rows, + intent, + groups, + *coordinate, + *value, + *window, + &context, + ) + .await? + } + Kind::Aggregate { groups, measures } => { + reduce(rows, groups, measures, &operator.inputs[0], &context).await? + } + _ => unreachable!(), + }; + Batch::try_new(output, result) + }) + .boxed_local()) +} + +pub(super) mod temporal; +async fn reduce( + rows: Vec>, + groups: &[usize], + measures: &[Reduction], + input: &Schema, + context: &RunContext, +) -> Result>, Error> { + let mut work = Cooperative::new(context); + let mut workspace = Workspace::new(context)?; + let mut grouped = BTreeMap::>, Vec>>::new(); + if rows.is_empty() && groups.is_empty() { + grouped.insert(vec![], vec![]); + } + for row in rows { + work.checkpoint().await?; + let key = group_key(&row, groups)?; + workspace.grow(std::mem::size_of::>())?; + if !grouped.contains_key(&key) { + workspace.grow(key_bytes(&key))?; + } + grouped.entry(key).or_default().push(row); + } + let mut output = Vec::new(); + for rows in grouped.into_values() { + work.checkpoint().await?; + let mut result = groups + .iter() + .map(|&i| rows[0][i].clone()) + .collect::>(); + for measure in measures { + result.push(reduce_one(&rows, measure, input, &mut work).await?); + } + workspace.grow(row_bytes(&result))?; + output.push(result); + } + Ok(output) +} + +async fn reduce_one( + rows: &[Vec], + measure: &Reduction, + input: &Schema, + work: &mut Cooperative, +) -> Result { + let column = match measure { + Reduction::Count => { + return Ok(Value::Int64( + i64::try_from(rows.len()).map_err(|_| invalid("count overflow"))?, + )) + } + Reduction::Sum(i) | Reduction::Avg(i) | Reduction::Min(i) | Reduction::Max(i) => *i, + }; + let values = rows + .iter() + .map(|r| &r[column]) + .filter(|v| !matches!(v, Value::Null)); + if matches!(measure, Reduction::Min(_) | Reduction::Max(_)) { + if plain(input, column)?.0 == &DataType::Float64 { + // Match exact-state kernels: ignore NaN when a numeric value exists. + let mut best: Option = None; + for value in values { + work.checkpoint().await?; + let Value::Float64(value) = value else { + return Err(invalid("floating aggregate value required")); + }; + best = Some(best.map_or(*value, |old| { + if matches!(measure, Reduction::Min(_)) { + old.min(*value) + } else { + old.max(*value) + } + })); + } + return Ok(best.map(Value::Float64).unwrap_or(Value::Null)); + } + let mut best: Option<&Value> = None; + for value in values { + work.checkpoint().await?; + if best + .map(|b| value.compare(b)) + .transpose()? + .is_none_or(|order| { + if matches!(measure, Reduction::Min(_)) { + order.is_lt() + } else { + order.is_gt() + } + }) + { + best = Some(value); + } + } + return Ok(best.cloned().unwrap_or(Value::Null)); + } + let mut count = 0usize; + let dtype = plain(input, column)?.0; + if dtype == &DataType::Int64 { + let mut sum = 0i128; + for v in values { + work.checkpoint().await?; + let Value::Int64(v) = v else { + return Err(invalid("integer aggregate value required")); + }; + sum = sum + .checked_add(i128::from(*v)) + .ok_or_else(|| invalid("integer aggregate overflow"))?; + count += 1; + } + return if matches!(measure, Reduction::Avg(_)) { + Ok(Value::Float64(sum as f64 / count as f64)) + } else { + Ok(Value::Int64( + i64::try_from(sum).map_err(|_| invalid("integer sum overflow"))?, + )) + }; + } + let mut sum = -0.0; + for v in values { + work.checkpoint().await?; + let Value::Float64(v) = v else { + return Err(invalid("floating aggregate value required")); + }; + sum += v; + count += 1; + } + Ok(Value::Float64(if matches!(measure, Reduction::Avg(_)) { + sum / count as f64 + } else { + sum + })) +} diff --git a/crates/asap-physical-operators/src/operators/aggregate/temporal.rs b/crates/asap-physical-operators/src/operators/aggregate/temporal.rs new file mode 100644 index 00000000..dbf66108 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/aggregate/temporal.rs @@ -0,0 +1,300 @@ +//! Windowed computations use Planner intents; deployments supply the input window. +use crate::{ + operators::{ + common::{key_bytes, row_bytes, Workspace}, + sort::cooperative_sort, + }, + runtime::{Cooperative, RunContext}, +}; +use crate::{ + values::{group_key, Value}, + Error, +}; +use planner_types::pre_asap::{AggIntent, ColumnRef}; +use std::collections::BTreeMap; + +pub(in crate::operators) async fn reduce( + rows: Vec>, + intent: &AggIntent, + groups: &[usize], + coordinate: usize, + value: usize, + window: Option<(i64, i64)>, + context: &RunContext, +) -> Result>, Error> { + let mut work = Cooperative::new(context); + let mut workspace = Workspace::new(context)?; + let mut grouped = + BTreeMap::>, (Vec, Vec<(f64, f64)>, Vec<(i64, f64)>)>::new(); + for row in rows { + work.checkpoint().await?; + let key = group_key(&row, groups)?; + workspace.grow(32)?; + if !grouped.contains_key(&key) { + workspace.grow(key_bytes(&key) + row_bytes(&row))?; + } + let entry = grouped.entry(key).or_insert_with(|| { + ( + groups.iter().map(|i| row[*i].clone()).collect(), + vec![], + vec![], + ) + }); + let Value::Float64(v) = row[value] else { + return Err(Error::Invalid("window value must be Float64".into())); + }; + match row[coordinate] { + Value::Timestamp(t) => entry.2.push((t, v)), + Value::Float64(bound) => entry.1.push((bound, v)), + _ => return Err(Error::Invalid("invalid window coordinate".into())), + } + } + let mut output = Vec::new(); + for (_, (mut keys, buckets, points)) in grouped { + work.checkpoint().await?; + let result = if let AggIntent::HistogramQuantile { q } = intent { + Some(Value::Float64(bucket_quantile(*q, buckets, context).await?)) + } else { + let points = cooperative_sort(points, |a, b| a.0.cmp(&b.0), context).await?; + let (start, end) = + window.ok_or_else(|| Error::Invalid("missing temporal window".into()))?; + if points.iter().any(|p| p.0 < start || p.0 > end) + || points.windows(2).any(|p| p[0].0 == p[1].0) + { + return Err(Error::Invalid( + "duplicate or out-of-window timestamp".into(), + )); + } + match intent { + AggIntent::Rate => rate(&points, start, end).map(Value::Float64), + AggIntent::Increase => rate(&points, start, end) + .map(|v| Value::Float64(v * (end as f64 - start as f64) / 1000.)), + AggIntent::Count { .. } => Some(Value::Int64( + i64::try_from(points.len()) + .map_err(|_| Error::Invalid("count overflow".into()))?, + )), + AggIntent::Sum { .. } => Some(Value::Float64(points.iter().map(|p| p.1).sum())), + AggIntent::Avg { .. } => Some(Value::Float64( + points.iter().map(|p| p.1).sum::() / points.len() as f64, + )), + AggIntent::Min { .. } => { + Some(Value::Float64(points.iter().fold(f64::NAN, |a, p| { + if a.is_nan() || p.1 < a { + p.1 + } else { + a + } + }))) + } + AggIntent::Max { .. } => { + Some(Value::Float64(points.iter().fold(f64::NAN, |a, p| { + if a.is_nan() || p.1 > a { + p.1 + } else { + a + } + }))) + } + _ => return Err(Error::Invalid("unsupported temporal intent".into())), + } + }; + if let Some(result) = result { + keys.push(result); + output.push(keys); + } + } + Ok(output) +} + +fn rate(points: &[(i64, f64)], start: i64, end: i64) -> Option { + if points.len() < 2 { + return None; + } + let (first_t, first) = points[0]; + let (last_t, last) = *points.last()?; + let span = (last_t as f64 - first_t as f64) / 1000.; + if span <= 0. { + return None; + } + let mut delta = last - first; + for pair in points.windows(2) { + if pair[1].1 < pair[0].1 { + delta += pair[0].1; + } + } + let average = span / (points.len() - 1) as f64; + let mut to_start = (first_t as f64 - start as f64) / 1000.; + let mut to_end = (end as f64 - last_t as f64) / 1000.; + if to_start >= average * 1.1 { + to_start = average / 2.; + } + // Apply the zero bound after the sparse-window half-interval cap. + if delta > 0. && first >= 0. { + to_start = to_start.min(span * first / delta); + } + if to_end >= average * 1.1 { + to_end = average / 2.; + } + Some(delta * (span + to_start + to_end) / span / ((end as f64 - start as f64) / 1000.)) +} + +async fn bucket_quantile( + q: f64, + mut b: Vec<(f64, f64)>, + context: &RunContext, +) -> Result { + let mut work = Cooperative::new(context); + let _scratch = context.reserve(b.len().checked_mul(16).ok_or(Error::MemoryLimit)?)?; + if q.is_nan() { + return Ok(f64::NAN); + } + if q < 0. { + return Ok(f64::NEG_INFINITY); + } + if q > 1. { + return Ok(f64::INFINITY); + } + b.retain(|p| !p.0.is_nan()); + b = cooperative_sort(b, |a, b| a.0.total_cmp(&b.0), context).await?; + let mut buckets: Vec<(f64, f64)> = Vec::new(); + for p in b { + work.checkpoint().await?; + if let Some(last) = buckets.last_mut() { + if last.0 == p.0 { + last.1 += p.1; + continue; + } + } + buckets.push(p); + } + if buckets.len() < 2 || buckets.last().unwrap().0 != f64::INFINITY { + return Ok(f64::NAN); + } + let mut prev = buckets[0].1; + for p in buckets.iter_mut().skip(1) { + work.checkpoint().await?; + if p.1 < prev || (p.1 - prev).abs() <= 1e-12 * (p.1.abs() + prev.abs()) { + p.1 = prev; + } + prev = p.1; + } + let count = buckets.last().unwrap().1; + if count == 0. { + return Ok(f64::NAN); + } + let rank = q * count; + let idx = buckets[..buckets.len() - 1].partition_point(|p| p.1 < rank); + if idx == buckets.len() - 1 { + return Ok(buckets[idx - 1].0); + } + if idx == 0 && buckets[0].0 <= 0. { + return Ok(buckets[0].0); + } + let (start, base) = if idx == 0 { (0., 0.) } else { buckets[idx - 1] }; + let (end, upper) = buckets[idx]; + Ok(start + (end - start) * (rank - base) / (upper - base)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + operators::Operator, + runtime::{batch_execution::evaluate_batch, Limits, RunContext, Scope}, + values::Batch, + }; + use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::DataType, + types::AccuracyTarget, + }; + use std::sync::Arc; + + // The same window operator must give the same answer in either engine phase. + #[test] + fn temporal_windows_execute_in_both_phases_and_count_is_integer() { + let schema = Arc::new(SummarySchema { + fields: vec![ + SummaryField { + name: "time".into(), + dtype: SummaryFamilyType::Plain(DataType::Timestamp), + nullable: false, + }, + SummaryField { + name: "value".into(), + dtype: SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }, + ], + time_index: Some(0), + }); + for scope in [ + Scope::Query { + evaluation_time_ms: 2000, + revision: 1, + }, + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 2000, + revision: 1, + }, + ] { + for (intent, expected) in [ + (AggIntent::Rate, Value::Float64(2.)), + (AggIntent::Increase, Value::Float64(4.)), + ( + AggIntent::Count { + accuracy: AccuracyTarget::Exact, + }, + Value::Int64(3), + ), + ] { + let batch = Batch::try_new( + schema.clone(), + vec![ + vec![Value::Timestamp(0), Value::Float64(2.)], + vec![Value::Timestamp(1000), Value::Float64(4.)], + vec![Value::Timestamp(2000), Value::Float64(2.)], + ], + ) + .unwrap(); + let operator = + Operator::window(schema.clone(), intent, 0, 1, vec![], Some((0, 2000))) + .unwrap(); + let result = evaluate_batch( + batch, + vec![operator], + RunContext::new(scope.clone(), Limits::default()).unwrap(), + ) + .unwrap(); + assert_eq!( + format!("{:?}", result[0].rows()[0][0]), + format!("{expected:?}") + ); + } + } + assert!(Operator::window(schema, AggIntent::Rate, 0, 1, vec![], Some((1, 1))).is_err()); + } + + // Histogram interpolation requires an infinite terminal bucket and coalesces duplicates. + #[test] + fn histogram_boundaries_and_duplicate_buckets() { + let context = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits::default(), + ) + .unwrap(); + let bucket_quantile = |q, buckets| { + futures::executor::block_on(super::bucket_quantile(q, buckets, &context)).unwrap() + }; + assert_eq!( + bucket_quantile(0.5, vec![(1., 1.), (1., 1.), (2., 4.), (f64::INFINITY, 4.)]), + 1. + ); + assert!(bucket_quantile(0.5, vec![(1., 2.), (2., 4.)]).is_nan()); + assert_eq!(bucket_quantile(-0.1, vec![]), f64::NEG_INFINITY); + } +} diff --git a/crates/asap-physical-operators/src/operators/aligned_binary.rs b/crates/asap-physical-operators/src/operators/aligned_binary.rs new file mode 100644 index 00000000..941a9f61 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/aligned_binary.rs @@ -0,0 +1,144 @@ +//! Arithmetic on complete, aligned population/window rows used by precomputation. +use super::*; +use planner_types::{post_asap::BinaryOperator, pre_asap::BinaryOpKind}; +use std::collections::BTreeSet; + +impl Operator { + /// Match every row by the declared identity columns. Unlike an inner join, + /// incomplete or duplicate keys are errors: dropping an update changes state. + pub fn aligned_binary( + left: Schema, + right: Schema, + keys: Vec<(usize, usize)>, + values: (usize, usize), + operator: BinaryOperator, + ) -> Result { + if keys.is_empty() + || !matches!(operator.kind, BinaryOpKind::Arithmetic(_)) + || operator.vector_match.is_some() + { + return Err(invalid( + "aligned arithmetic requires explicit keys and arithmetic semantics", + )); + } + for (input, value) in [(&left, values.0), (&right, values.1)] { + if input.fields.get(value).is_none_or(|f| { + f.nullable || f.dtype != SummaryFamilyType::Plain(DataType::Float64) + }) { + return Err(invalid( + "aligned arithmetic requires non-null Float64 values", + )); + } + } + let mut left_keys = BTreeSet::new(); + let mut right_keys = BTreeSet::new(); + for &(l, r) in &keys { + if l == values.0 + || r == values.1 + || !left_keys.insert(l) + || !right_keys.insert(r) + || left + .fields + .get(l) + .zip(right.fields.get(r)) + .is_none_or(|(l, r)| l.nullable || r.nullable || l.dtype != r.dtype) + { + return Err(invalid("invalid aligned arithmetic keys")); + } + } + if left_keys.len() + 1 != left.fields.len() || right_keys.len() + 1 != right.fields.len() { + return Err(invalid( + "aligned arithmetic must account for every input column", + )); + } + Ok(Self { + output: left.clone(), + inputs: vec![left, right], + kind: Kind::AlignedBinary { + keys, + values, + operator, + }, + }) + } +} + +pub(super) fn execute<'a>( + op: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let Kind::AlignedBinary { + keys, + values, + operator, + } = &op.kind + else { + unreachable!() + }; + let right = inputs + .pop() + .ok_or_else(|| invalid("missing aligned right input"))?; + let left = inputs + .pop() + .ok_or_else(|| invalid("missing aligned left input"))?; + Ok(futures::stream::once(async move { + let ((left, _left_memory), (right, _right_memory)) = + futures::try_join!(collect_rows(left, &context), collect_rows(right, &context))?; + if left.is_empty() || left.len() != right.len() { + return Err(invalid( + "aligned arithmetic requires matching nonempty key sets", + )); + } + let mut work = Cooperative::new(&context); + let mut workspace = Workspace::new(&context)?; + let columns = |side: bool| { + keys.iter() + .map(|&(l, r)| if side { r } else { l }) + .collect::>() + }; + let left_columns = columns(false); + let right_columns = columns(true); + let mut indexed = BTreeMap::new(); + for row in right { + work.checkpoint().await?; + let key = group_key(&row, &right_columns)?; + let Value::Float64(value) = row[values.1] else { + return Err(invalid("invalid aligned value")); + }; + if !value.is_finite() { + return Err(invalid("aligned arithmetic input is non-finite")); + } + workspace.grow(key_bytes(&key) + 64)?; + if indexed.insert(key, value).is_some() { + return Err(invalid("aligned arithmetic input has duplicate keys")); + } + } + let mut rows = Vec::new(); + for mut row in left { + work.checkpoint().await?; + let key = group_key(&row, &left_columns)?; + let right = indexed + .remove(&key) + .ok_or_else(|| invalid("aligned arithmetic input has missing or duplicate keys"))?; + let Value::Float64(left) = row[values.0] else { + return Err(invalid("invalid aligned value")); + }; + if !left.is_finite() { + return Err(invalid("aligned arithmetic input is non-finite")); + } + let result = crate::expressions::arithmetic::evaluate_binary(operator, left, right)?; + if !matches!(result, Value::Float64(value) if value.is_finite()) { + return Err(invalid("aligned arithmetic produced a non-finite update")); + } + row[values.0] = result; + workspace.grow(std::mem::size_of::>())?; + rows.push(row); + } + if !indexed.is_empty() { + return Err(invalid("aligned arithmetic has unmatched input keys")); + } + Batch::try_new(op.output.clone(), rows) + }) + .boxed_local()) +} diff --git a/crates/asap-physical-operators/src/operators/common.rs b/crates/asap-physical-operators/src/operators/common.rs new file mode 100644 index 00000000..9f1709da --- /dev/null +++ b/crates/asap-physical-operators/src/operators/common.rs @@ -0,0 +1,75 @@ +use super::*; +pub(super) fn invalid(message: &str) -> Error { + Error::Invalid(message.into()) +} +pub(super) fn schema(fields: Vec) -> Schema { + Arc::new(SummarySchema { + fields, + time_index: None, + }) +} +pub(super) fn result_field(name: &str, dtype: DataType, nullable: bool) -> SummaryField { + SummaryField { + name: name.into(), + dtype: SummaryFamilyType::Plain(dtype), + nullable, + } +} + +pub(super) fn validate_groups(input: &Schema, groups: &[usize]) -> Result<(), Error> { + for &i in groups { + plain(input, i)?; + } + if groups + .iter() + .collect::>() + .len() + != groups.len() + { + return Err(invalid("duplicate group columns")); + } + Ok(()) +} +pub(super) async fn collect_rows( + mut input: Input<'_, Batch>, + context: &RunContext, +) -> Result<(Vec>, Vec), Error> { + let mut rows = Vec::new(); + let mut work = Cooperative::new(context); + let mut reservations = Vec::new(); + while let Some(batch) = input.next().await { + let batch = batch?; + reservations.push(context.reserve(batch.bytes())?); + for row in batch.rows() { + work.checkpoint().await?; + rows.push(row.clone()); + } + } + Ok((rows, reservations)) +} +/// Estimates retained workspace before growing collections. It is not an RSS limit. +pub(super) struct Workspace { + reservation: Reservation, + bytes: usize, +} +impl Workspace { + pub(super) fn new(context: &RunContext) -> Result { + Ok(Self { + reservation: context.reserve(0)?, + bytes: 0, + }) + } + pub(super) fn grow(&mut self, bytes: usize) -> Result<(), Error> { + self.bytes = self.bytes.checked_add(bytes).ok_or(Error::MemoryLimit)?; + self.reservation.resize(self.bytes) + } +} +pub(super) fn row_bytes(row: &[Value]) -> usize { + std::mem::size_of::>() + row.iter().map(Value::bytes).sum::() +} +pub(super) fn key_bytes(key: &[Vec]) -> usize { + 64 + key + .iter() + .map(|part| std::mem::size_of::>() + part.len()) + .sum::() +} diff --git a/crates/asap-physical-operators/src/operators/current_series.rs b/crates/asap-physical-operators/src/operators/current_series.rs new file mode 100644 index 00000000..937036c5 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/current_series.rs @@ -0,0 +1,135 @@ +//! A bounded instant-vector snapshot: select latest before removing stale markers. +use super::*; + +impl Operator { + pub fn current_series( + input: Schema, + identity: usize, + coordinate: usize, + value: usize, + lookback_ms: i64, + ) -> Result { + if lookback_ms <= 0 + || plain(&input, identity)? != (&DataType::Utf8, false) + || plain(&input, coordinate)? != (&DataType::Timestamp, false) + || plain(&input, value)? != (&DataType::Float64, false) + { + return Err(invalid("invalid current-series input contract")); + } + Ok(Self { + kind: Kind::CurrentSeries { + identity, + coordinate, + value, + lookback_ms, + }, + inputs: vec![input.clone()], + output: input, + }) + } +} + +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let Kind::CurrentSeries { + identity, + coordinate, + value, + lookback_ms, + } = operator.kind + else { + unreachable!() + }; + let input = inputs + .pop() + .ok_or_else(|| invalid("current-series input missing"))?; + let output = operator.output.clone(); + let (start, end) = window(lookback_ms, &context)?; + Ok(futures::stream::once(async move { + let (rows, _memory) = collect_rows(input, &context).await?; + let mut latest = BTreeMap::, usize>::new(); + let mut work = Cooperative::new(&context); + let mut workspace = Workspace::new(&context)?; + for (index, row) in rows.iter().enumerate() { + work.checkpoint().await?; + let Value::Timestamp(timestamp) = row[coordinate] else { + unreachable!() + }; + if timestamp <= start || timestamp > end { + continue; + } + let key = row[identity].key()?; + if let Some(&previous) = latest.get(&key) { + let Value::Timestamp(previous_time) = rows[previous][coordinate] else { + unreachable!() + }; + if timestamp < previous_time { + continue; + } + if timestamp == previous_time { + let (Value::Float64(a), Value::Float64(b)) = + (&row[value], &rows[previous][value]) + else { + unreachable!() + }; + if a.to_bits() != b.to_bits() { + return Err(invalid("conflicting samples for one series timestamp")); + } + continue; + } + } else { + workspace.grow(64 + key.len())?; + } + latest.insert(key, index); + } + let mut result = Vec::new(); + for index in latest.into_values() { + work.checkpoint().await?; + let Value::Float64(sample) = rows[index][value] else { + unreachable!() + }; + if sample.to_bits() == 0x7ff0_0000_0000_0002 { + continue; + } + workspace.grow(row_bytes(&rows[index]))?; + let mut row = rows[index].clone(); + row[coordinate] = Value::Timestamp(end); + result.push(row); + } + Batch::try_new(output, result) + }) + .boxed_local()) +} + +fn window(lookback_ms: i64, context: &RunContext) -> Result<(i64, i64), Error> { + let end = match context.scope { + crate::runtime::Scope::Query { + evaluation_time_ms, .. + } => evaluation_time_ms, + crate::runtime::Scope::Ingestion { window_end_ms, .. } => window_end_ms, + }; + let start = end + .checked_sub(lookback_ms) + .ok_or_else(|| invalid("current-series window overflows"))?; + if let crate::runtime::Scope::Ingestion { + window_start_ms, .. + } = context.scope + { + if window_start_ms != start { + return Err(invalid( + "current-series maintenance window differs from lookback", + )); + } + } + Ok((start, end)) +} + +pub(super) fn validate_context(operator: &Operator, context: &RunContext) -> Result<(), Error> { + if let Kind::CurrentSeries { lookback_ms, .. } = operator.kind { + window(lookback_ms, context)?; + } + Ok(()) +} diff --git a/crates/asap-physical-operators/src/operators/filter.rs b/crates/asap-physical-operators/src/operators/filter.rs new file mode 100644 index 00000000..8862aacc --- /dev/null +++ b/crates/asap-physical-operators/src/operators/filter.rs @@ -0,0 +1,39 @@ +use super::*; +impl Operator { + pub fn filter(input: Schema, predicate: Expression) -> Result { + if predicate.dtype(&input)?.0 != DataType::Bool { + return Err(invalid("filter predicate must be boolean")); + } + Ok(Self { + kind: Kind::Filter(predicate), + inputs: vec![input.clone()], + output: input, + }) + } +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + match &operator.kind { + Kind::Filter(predicate) => Ok(input + .map(move |batch| { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + let batch = batch?; + let mut rows = Vec::new(); + for row in batch.rows() { + if matches!(predicate.evaluate(row)?, Value::Bool(true)) { + rows.push(row.clone()); + } + } + Batch::try_new(output.clone(), rows) + }) + .boxed_local()), + _ => unreachable!(), + } +} diff --git a/crates/asap-physical-operators/src/operators/joins/mod.rs b/crates/asap-physical-operators/src/operators/joins/mod.rs new file mode 100644 index 00000000..608fd8d7 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/joins/mod.rs @@ -0,0 +1,219 @@ +use super::*; +impl Operator { + pub fn semi_join( + left: Schema, + right: Schema, + keys: Vec<(usize, usize)>, + ) -> Result { + if keys.is_empty() { + return Err(invalid("semi-join needs matching keys")); + } + for &(l, r) in &keys { + if plain(&left, l)?.0 != plain(&right, r)?.0 { + return Err(invalid("join key types differ")); + } + } + Ok(Self { + kind: Kind::SemiJoin { + keys, + require_complete_right: false, + }, + inputs: vec![left.clone(), right], + output: left, + }) + } + pub(crate) fn require_complete_right(mut self) -> Self { + if let Kind::SemiJoin { + require_complete_right, + .. + } = &mut self.kind + { + *require_complete_right = true; + } + self + } + pub fn relational_join( + left: Schema, + right: Schema, + kind: planner_types::pre_asap::JoinKind, + predicate: &planner_types::pre_asap::Predicate, + output: Schema, + ) -> Result { + use planner_types::pre_asap::JoinKind; + let mut joined = left.fields.clone(); + joined.extend(right.fields.clone()); + let predicate = + crate::expressions::CompiledExpression::compile(&predicate.0, &schema(joined.clone()))?; + if predicate.dtype().0 != DataType::Bool { + return Err(invalid("join predicate must be boolean")); + } + let fields = if matches!(kind, JoinKind::Semi | JoinKind::Anti) { + left.fields.clone() + } else { + for field in &mut joined[..left.fields.len()] { + if matches!(kind, JoinKind::Right | JoinKind::Full) { + field.nullable = true; + } + } + for field in &mut joined[left.fields.len()..] { + if matches!(kind, JoinKind::Left | JoinKind::Full) { + field.nullable = true; + } + } + joined + }; + Self { + kind: Kind::Join { + kind, + predicate: Box::new(predicate), + }, + inputs: vec![left, right], + output: schema(fields), + } + .with_output_schema(output) + } +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + if let Kind::Join { kind, predicate } = &operator.kind { + let right = inputs.pop().ok_or_else(|| invalid("right input missing"))?; + let left = inputs.pop().ok_or_else(|| invalid("left input missing"))?; + return Ok(futures::stream::once(async move { + use planner_types::pre_asap::JoinKind; + let ((left, _left_memory), (right, _right_memory)) = + futures::try_join!(collect_rows(left, &context), collect_rows(right, &context))?; + let mut workspace = Workspace::new(&context)?; + let mut work = Cooperative::new(&context); + workspace.grow(right.len())?; + let mut result = Vec::new(); + let mut right_matched = vec![false; right.len()]; + for left_row in &left { + work.checkpoint().await?; + let mut matched = false; + for (i, right_row) in right.iter().enumerate() { + work.checkpoint().await?; + let mut joined = left_row.clone(); + joined.extend(right_row.iter().cloned()); + if *kind == JoinKind::Cross + || matches!(predicate.evaluate(&joined)?, Value::Bool(true)) + { + matched = true; + right_matched[i] = true; + match kind { + JoinKind::Semi => { + workspace.grow(row_bytes(left_row))?; + result.push(left_row.clone()); + break; + } + JoinKind::Anti => break, + _ => { + workspace.grow(row_bytes(&joined))?; + result.push(joined); + } + } + } + } + if !matched { + match kind { + JoinKind::Left | JoinKind::Full => { + let mut joined = left_row.clone(); + joined.resize( + joined.len() + operator.inputs[1].fields.len(), + Value::Null, + ); + workspace.grow(row_bytes(&joined))?; + result.push(joined); + } + JoinKind::Anti => { + workspace.grow(row_bytes(left_row))?; + result.push(left_row.clone()); + } + _ => {} + } + } + } + if matches!(kind, JoinKind::Right | JoinKind::Full) { + for (matched, row) in right_matched.into_iter().zip(right) { + work.checkpoint().await?; + if !matched { + let mut joined = vec![Value::Null; operator.inputs[0].fields.len()]; + joined.extend(row); + workspace.grow(row_bytes(&joined))?; + result.push(joined); + } + } + } + Batch::try_new(output, result) + }) + .boxed_local()); + } + if let Kind::SemiJoin { + keys, + require_complete_right, + } = &operator.kind + { + let right = inputs.pop().ok_or_else(|| invalid("right input missing"))?; + let left = inputs.pop().ok_or_else(|| invalid("left input missing"))?; + return Ok(futures::stream::once(async move { + // Poll both branches together: either may depend on a common producer. + let ((left, _left_memory), (right, _right_memory)) = + futures::try_join!(collect_rows(left, &context), collect_rows(right, &context))?; + let right_cols = keys.iter().map(|(_, r)| *r).collect::>(); + let left_cols = keys.iter().map(|(l, _)| *l).collect::>(); + let mut members = std::collections::BTreeSet::new(); + let mut workspace = Workspace::new(&context)?; + let mut work = Cooperative::new(&context); + for row in &right { + work.checkpoint().await?; + if right_cols.iter().all(|&i| matchable_key(&row[i])) { + let key = group_key(row, &right_cols)?; + if !members.contains(&key) { + workspace.grow(key_bytes(&key))?; + members.insert(key); + } + } else if *require_complete_right { + return Err(invalid( + "certified pruning candidate has an unmatchable key", + )); + } + } + let mut rows = Vec::new(); + let mut covered = std::collections::BTreeSet::new(); + for row in left { + work.checkpoint().await?; + if left_cols.iter().all(|&i| matchable_key(&row[i])) + && members.contains(&group_key(&row, &left_cols)?) + { + if *require_complete_right { + let key = group_key(&row, &left_cols)?; + if !covered.contains(&key) { + workspace.grow(key_bytes(&key))?; + covered.insert(key); + } + } + workspace.grow(std::mem::size_of::>())?; + rows.push(row); + } + } + if *require_complete_right && members != covered { + return Err(invalid("certified pruning key has no authoritative value")); + } + Batch::try_new(output, rows) + }) + .boxed_local()); + } + unreachable!() +} + +// Group keys canonicalize NaNs, but equality joins must not match them. +fn matchable_key(value: &Value) -> bool { + match value { + Value::Null => false, + Value::Float64(v) => !v.is_nan(), + _ => true, + } +} diff --git a/crates/asap-physical-operators/src/operators/limit.rs b/crates/asap-physical-operators/src/operators/limit.rs new file mode 100644 index 00000000..5f5e601f --- /dev/null +++ b/crates/asap-physical-operators/src/operators/limit.rs @@ -0,0 +1,69 @@ +use super::*; +impl Operator { + pub fn limit(input: Schema, n: u64, offset: u64, groups: Vec) -> Result { + validate_groups(&input, &groups)?; + Ok(Self { + kind: Kind::Limit { n, offset, groups }, + inputs: vec![input.clone()], + output: input, + }) + } +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + match &operator.kind { + Kind::Limit { n, offset, groups } => { + let counts = BTreeMap::>, u64>::new(); + Ok(futures::stream::try_unfold( + (input, counts, Vec::::new(), false), + move |(mut input, mut counts, mut memory, done)| { + let output = output.clone(); + let context = context.clone(); + async move { + if done || *n == 0 { + return Ok(None); + } + let Some(batch) = input.next().await else { + return Ok(None); + }; + let batch = batch?; + let mut rows = Vec::new(); + for row in batch.rows() { + let key = group_key(row, groups)?; + if !counts.contains_key(&key) { + memory.push( + context.reserve( + key.iter() + .map(|part| part.len() + std::mem::size_of::>()) + .sum::() + + 64, + )?, + ); + } + let count = counts.entry(key).or_default(); + if *count >= *offset && count.saturating_sub(*offset) < *n { + rows.push(row.clone()); + } + *count = count.saturating_add(1); + } + let done = groups.is_empty() + && counts + .get(&vec![]) + .is_some_and(|count| count.saturating_sub(*offset) >= *n); + Ok(Some(( + Batch::try_new(output, rows)?, + (input, counts, memory, done), + ))) + } + }, + ) + .boxed_local()) + } + _ => unreachable!(), + } +} diff --git a/crates/asap-physical-operators/src/operators/mod.rs b/crates/asap-physical-operators/src/operators/mod.rs new file mode 100644 index 00000000..e4c6944f --- /dev/null +++ b/crates/asap-physical-operators/src/operators/mod.rs @@ -0,0 +1,316 @@ +//! Native physical operators. Each module owns its constructors and execution. +mod aligned_binary; +use crate::plan::{Boundedness, Emission, PhysicalOperator, PlanProperties}; +use crate::{ + runtime::{Cooperative, Input, OutputStream, Reservation, RunContext}, + values::{field, group_key, plain, Batch, Schema, Value}, + Error, +}; +use futures::StreamExt; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema, SummaryUpdate}, + pre_asap::{ColumnRef, DataType}, +}; +use std::{collections::BTreeMap, sync::Arc}; +pub(crate) mod common; +use crate::expressions::ordered; +pub use crate::expressions::Expression; +use common::*; +mod aggregate; +mod current_series; +mod filter; +mod joins; +mod limit; +mod panes; +mod projection; +mod sort; +mod source; +mod summary; +mod unchecked; +pub(crate) mod vector_binary; +pub(crate) mod vector_window; +pub use aggregate::Reduction; +pub use sort::SortKey; +pub use summary::ReadoutQuery; +#[derive(Clone, serde::Serialize, serde::Deserialize)] +enum Kind { + #[serde(skip)] + Source(Vec), + Constant { + value: Value, + dtype: DataType, + }, + PaneInput { + coordinate: usize, + layout: planner_types::post_asap::PaneLayout, + offset_ms: Option, + }, + ScopeTimestamp { + columns: Vec>, + }, + CurrentSeries { + identity: usize, + coordinate: usize, + value: usize, + lookback_ms: i64, + }, + Union, + VectorToScalar { + column: usize, + }, + VectorBinary { + operator: planner_types::post_asap::BinaryOperator, + return_bool: bool, + }, + AlignedBinary { + keys: Vec<(usize, usize)>, + values: (usize, usize), + operator: planner_types::post_asap::BinaryOperator, + }, + RangeWindow { + intent: Box>, + }, + HistogramQuantile, + Project(Vec), + Filter(Expression), + Limit { + n: u64, + offset: u64, + groups: Vec, + }, + Sort { + keys: Vec, + groups: Vec, + }, + Window { + intent: Box>, + coordinate: usize, + value: usize, + groups: Vec, + window: Option<(i64, i64)>, + }, + Aggregate { + groups: Vec, + measures: Vec, + }, + SemiJoin { + keys: Vec<(usize, usize)>, + require_complete_right: bool, + }, + Join { + kind: planner_types::pre_asap::JoinKind, + predicate: Box, + }, + SummaryBuild { + family: SummaryFamilyType, + value: usize, + time: Option, + groups: Vec, + }, + KeyedSummaryBuild { + family: SummaryFamilyType, + value: usize, + items: Vec, + groups: Vec, + }, + KeyedReadout { + state: usize, + k: usize, + }, + SummaryMerge { + state: usize, + groups: Vec, + }, + Readout { + state: usize, + query: ReadoutQuery, + }, +} +/// A bound operation has a fully checked input/output contract before execution. +#[derive(Clone, serde::Serialize, serde::Deserialize)] +#[serde(try_from = "unchecked::UncheckedOperator")] +pub struct Operator { + kind: Kind, + inputs: Vec, + output: Schema, +} +impl Operator { + /// Resolve a counter readout's logical lookback to this run's evaluation range. + pub(super) fn readout_range(&self, context: &RunContext) -> Result, Error> { + let Kind::Readout { + query: + ReadoutQuery::Exact(crate::summary_kernels::exact::ExactReadout { + lookback_ms: Some(lookback), + .. + }), + .. + } = &self.kind + else { + return Ok(None); + }; + let end = match context.scope { + crate::runtime::Scope::Query { + evaluation_time_ms, .. + } => evaluation_time_ms, + crate::runtime::Scope::Ingestion { window_end_ms, .. } => window_end_ms, + }; + let start = end + .checked_sub(*lookback) + .ok_or_else(|| invalid("counter window overflows Int64"))?; + if let crate::runtime::Scope::Ingestion { + window_start_ms, .. + } = context.scope + { + if window_start_ms != start { + return Err(invalid( + "maintenance window differs from logical counter window", + )); + } + } + Ok(Some((start, end))) + } + + pub(crate) fn with_output_schema(mut self, output: Schema) -> Result { + if self.output.fields.len() != output.fields.len() + || self + .output + .fields + .iter() + .zip(&output.fields) + .any(|(actual, declared)| { + actual.dtype != declared.dtype || (actual.nullable && !declared.nullable) + }) + { + return Err(invalid("native output type differs from Planner output")); + } + if output.time_index.is_some_and(|i| { + i >= output.fields.len() + || output.fields[i].dtype != SummaryFamilyType::Plain(DataType::Timestamp) + }) { + return Err(invalid("invalid output time column")); + } + self.output = output; + Ok(self) + } + pub fn schema(&self) -> Schema { + self.output.clone() + } +} +impl PhysicalOperator for Operator { + fn requires_bounded_input(&self) -> bool { + matches!( + self.kind, + Kind::Sort { .. } + | Kind::AlignedBinary { .. } + | Kind::VectorBinary { .. } + | Kind::RangeWindow { .. } + | Kind::HistogramQuantile + | Kind::CurrentSeries { .. } + | Kind::Aggregate { .. } + | Kind::Window { .. } + | Kind::Join { .. } + | Kind::SemiJoin { .. } + | Kind::SummaryBuild { .. } + | Kind::KeyedSummaryBuild { .. } + | Kind::SummaryMerge { .. } + | Kind::VectorToScalar { .. } + ) + } + fn properties(&self, inputs: &[PlanProperties]) -> PlanProperties { + let boundedness = match &self.kind { + Kind::Source(_) | Kind::Constant { .. } => Boundedness::Bounded, + Kind::Limit { groups, .. } if groups.is_empty() => Boundedness::Bounded, + _ => Boundedness::from_inputs(inputs), + }; + PlanProperties { + boundedness, + emission: if matches!( + self.kind, + Kind::PaneInput { .. } | Kind::ScopeTimestamp { .. } + ) { + inputs + .first() + .map_or(Emission::Unknown, |input| input.emission) + } else if self.requires_bounded_input() { + Emission::AfterInput + } else { + Emission::Incremental + }, + } + } + + fn name(&self) -> &str { + match self.kind { + Kind::Source(_) => "Source", + Kind::Constant { .. } => "Constant", + Kind::PaneInput { .. } => "PaneInput", + Kind::ScopeTimestamp { .. } => "ScopeTimestamp", + Kind::Union => "Union", + Kind::CurrentSeries { .. } => "CurrentSeries", + Kind::VectorToScalar { .. } => "VectorToScalar", + Kind::VectorBinary { .. } => "VectorBinary", + Kind::AlignedBinary { .. } => "AlignedBinary", + Kind::RangeWindow { .. } => "RangeWindow", + Kind::HistogramQuantile => "HistogramQuantile", + Kind::Project(_) => "Project", + Kind::Filter(_) => "Filter", + Kind::Limit { .. } => "Limit", + Kind::Sort { .. } => "Sort", + Kind::Aggregate { .. } => "Aggregate", + Kind::Window { .. } => "WindowAggregate", + Kind::SemiJoin { .. } => "SemiJoin", + Kind::Join { .. } => "RelationalJoin", + Kind::SummaryBuild { .. } | Kind::KeyedSummaryBuild { .. } => "SummaryAgg", + Kind::KeyedReadout { .. } => "SummaryEstimate", + Kind::SummaryMerge { .. } => "SummaryMerge", + Kind::Readout { .. } => "SummaryReadout", + } + } + fn validate_context(&self, context: &RunContext) -> Result<(), Error> { + panes::validate_context(self, context)?; + current_series::validate_context(self, context)?; + self.readout_range(context).map(|_| ()) + } + fn input_schemas(&self) -> Vec { + self.inputs.clone() + } + fn output_schema(&self) -> Schema { + self.output.clone() + } + fn output_bytes(&self, value: &Batch) -> usize { + value.bytes() + } + fn start<'a>( + &'a self, + inputs: Vec>, + context: RunContext, + ) -> Result, Error> { + match self.kind { + Kind::Source(_) | Kind::Constant { .. } | Kind::Union | Kind::VectorToScalar { .. } => { + source::execute(self, inputs, context) + } + Kind::VectorBinary { .. } => vector_binary::execute(self, inputs, context), + Kind::AlignedBinary { .. } => aligned_binary::execute(self, inputs, context), + Kind::RangeWindow { .. } | Kind::HistogramQuantile => { + vector_window::execute(self, inputs, context) + } + Kind::Project(_) => projection::execute(self, inputs, context), + Kind::CurrentSeries { .. } => current_series::execute(self, inputs, context), + Kind::PaneInput { .. } | Kind::ScopeTimestamp { .. } => { + panes::execute(self, inputs, context) + } + Kind::Filter(_) => filter::execute(self, inputs, context), + Kind::Limit { .. } => limit::execute(self, inputs, context), + Kind::Sort { .. } => sort::execute(self, inputs, context), + Kind::Window { .. } | Kind::Aggregate { .. } => { + aggregate::execute(self, inputs, context) + } + Kind::Join { .. } | Kind::SemiJoin { .. } => joins::execute(self, inputs, context), + Kind::SummaryMerge { .. } => summary::execute_merge(self, inputs, context), + Kind::SummaryBuild { .. } + | Kind::Readout { .. } + | Kind::KeyedSummaryBuild { .. } + | Kind::KeyedReadout { .. } => summary::execute(self, inputs, context), + } + } +} diff --git a/crates/asap-physical-operators/src/operators/panes.rs b/crates/asap-physical-operators/src/operators/panes.rs new file mode 100644 index 00000000..4748db09 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/panes.rs @@ -0,0 +1,228 @@ +//! Run-scoped pane population checks and timestamp restoration after reduction. +use super::*; +use crate::runtime::Scope; +use planner_types::post_asap::{validate_pane_coverage, PaneLayout, WindowEdgeCoverage}; + +impl Operator { + pub(crate) fn pane_input( + input: Schema, + coordinate: usize, + layout: PaneLayout, + offset_ms: Option, + ) -> Result { + if plain(&input, coordinate)? != (&DataType::Timestamp, false) { + return Err(invalid("pane input requires a non-null timestamp")); + } + if layout.pane_width_ms > i64::MAX as u64 { + return Err(invalid("pane width exceeds timestamp range")); + } + validate_pane_coverage( + &layout, + layout.pane_origin_ms, + &WindowEdgeCoverage::PaneAligned, + ) + .map_err(|error| Error::Invalid(format!("invalid pane layout: {error:?}")))?; + if offset_ms.is_some_and(|offset| offset < 0) { + return Err(invalid("negative pane offset")); + } + Ok(Self { + kind: Kind::PaneInput { + coordinate, + layout, + offset_ms, + }, + inputs: vec![input.clone()], + output: input, + }) + } + + pub(crate) fn scope_timestamp(input: Schema, output: Schema) -> Result { + crate::values::validate_schema(&output)?; + let coordinate = output + .time_index + .ok_or_else(|| invalid("temporal output requires a time index"))?; + if plain(&output, coordinate)? != (&DataType::Timestamp, false) { + return Err(invalid("temporal output requires a non-null timestamp")); + } + let mut columns = Vec::new(); + let mut used = std::collections::BTreeSet::new(); + for (index, field) in output.fields.iter().enumerate() { + if index == coordinate { + columns.push(None); + continue; + } + let matches: Vec<_> = input + .fields + .iter() + .enumerate() + .filter(|(_, candidate)| { + candidate.dtype == field.dtype + && candidate.nullable == field.nullable + && (candidate.name == field.name + || !matches!(field.dtype, SummaryFamilyType::Plain(_))) + }) + .map(|(index, _)| index) + .collect(); + let [column] = matches.as_slice() else { + return Err(invalid("temporal output column missing or ambiguous")); + }; + if !used.insert(*column) { + return Err(invalid("temporal output repeats an input column")); + } + columns.push(Some(*column)); + } + if used.len() != input.fields.len() { + return Err(invalid("temporal output drops an input column")); + } + Ok(Self { + kind: Kind::ScopeTimestamp { columns }, + inputs: vec![input], + output, + }) + } +} + +pub(super) fn validate_context(operator: &Operator, context: &RunContext) -> Result<(), Error> { + let Kind::PaneInput { + layout, offset_ms, .. + } = &operator.kind + else { + return Ok(()); + }; + let end = match (&context.scope, offset_ms) { + ( + Scope::Ingestion { + window_start_ms, + window_end_ms, + .. + }, + None, + ) => { + if window_end_ms.checked_sub(*window_start_ms) != Some(layout.pane_width_ms as i64) { + return Err(invalid("maintenance run must cover exactly one pane")); + } + *window_end_ms + } + ( + Scope::Query { + evaluation_time_ms, .. + }, + Some(offset), + ) => { + let end = evaluation_time_ms + .checked_sub(*offset) + .ok_or_else(|| invalid("query pane timestamp overflows"))?; + end.checked_sub(layout.pane_width_ms as i64) + .ok_or_else(|| invalid("query pane start overflows"))?; + end + } + _ => return Err(invalid("pane operator received the wrong execution scope")), + }; + validate_pane_coverage(layout, Some(end), &WindowEdgeCoverage::PaneAligned).map_err(|error| { + Error::Invalid(format!( + "query requires aligned panes or boundary residuals: {error:?}" + )) + }) +} + +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + validate_context(operator, &context)?; + let input = inputs.pop().ok_or_else(|| invalid("pane input missing"))?; + let output = operator.output.clone(); + let mut seen = std::collections::BTreeSet::new(); + let mut memory = context.reserve(0)?; + let mut key_bytes = 0; + Ok(input + .map(move |batch| { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + let batch = batch?; + match &operator.kind { + Kind::PaneInput { + coordinate, + offset_ms, + .. + } => { + let groups: Vec<_> = output + .fields + .iter() + .enumerate() + .filter(|(index, field)| { + *index != *coordinate + && matches!(field.dtype, SummaryFamilyType::Plain(_)) + }) + .map(|(index, _)| index) + .collect(); + for row in batch.rows() { + let Value::Timestamp(timestamp) = row[*coordinate] else { + return Err(invalid("pane timestamp type mismatch")); + }; + match (&context.scope, offset_ms) { + ( + Scope::Ingestion { + window_start_ms, + window_end_ms, + .. + }, + None, + ) if timestamp > *window_start_ms && timestamp <= *window_end_ms => {} + ( + Scope::Query { + evaluation_time_ms, .. + }, + Some(offset), + ) if timestamp + == evaluation_time_ms + .checked_sub(*offset) + .ok_or_else(|| invalid("pane timestamp overflows"))? => + { + let key = group_key(row, &groups)?; + if seen.contains(&key) { + return Err(invalid("duplicate entity state within a pane")); + } + key_bytes += key.iter().map(Vec::len).sum::() + + key.len() * std::mem::size_of::>() + + 64; + memory.resize(key_bytes)?; + seen.insert(key); + } + _ => { + return Err(invalid("input population differs from required pane")) + } + } + } + Ok(batch.value().clone()) + } + Kind::ScopeTimestamp { columns } => { + let timestamp = match context.scope { + Scope::Ingestion { window_end_ms, .. } => window_end_ms, + Scope::Query { + evaluation_time_ms, .. + } => evaluation_time_ms, + }; + let rows = batch + .rows() + .iter() + .map(|row| { + columns + .iter() + .map(|column| { + column.map_or(Value::Timestamp(timestamp), |column| { + row[column].clone() + }) + }) + .collect() + }) + .collect(); + Batch::try_new(output.clone(), rows) + } + _ => unreachable!(), + } + }) + .boxed_local()) +} diff --git a/crates/asap-physical-operators/src/operators/projection.rs b/crates/asap-physical-operators/src/operators/projection.rs new file mode 100644 index 00000000..4a174ed9 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/projection.rs @@ -0,0 +1,56 @@ +use super::*; +impl Operator { + pub fn project(input: Schema, columns: Vec<(String, Expression)>) -> Result { + let fields = columns + .iter() + .map(|(name, e)| { + if let Expression::Column(index) = e { + let mut field = input + .fields + .get(*index) + .ok_or_else(|| invalid("projection column out of range"))? + .clone(); + field.name = name.clone(); + return Ok(field); + } + let (t, n) = e.dtype(&input)?; + Ok(result_field(name, t, n)) + }) + .collect::>()?; + Ok(Self { + kind: Kind::Project(columns.into_iter().map(|(_, e)| e).collect()), + inputs: vec![input], + output: schema(fields), + }) + } +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + match &operator.kind { + Kind::Project(expressions) => Ok(input + .map(move |batch| { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + let batch = batch?; + let rows = batch + .rows() + .iter() + .map(|r| { + expressions + .iter() + .map(|e| e.evaluate(r)) + .collect::, _>>() + }) + .collect::, _>>()?; + Batch::try_new(output.clone(), rows) + }) + .boxed_local()), + _ => unreachable!(), + } +} diff --git a/crates/asap-physical-operators/src/operators/sort.rs b/crates/asap-physical-operators/src/operators/sort.rs new file mode 100644 index 00000000..71f81de0 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/sort.rs @@ -0,0 +1,172 @@ +use super::*; +impl Operator { + pub fn sort(input: Schema, keys: Vec, groups: Vec) -> Result { + validate_groups(&input, &groups)?; + for key in &keys { + if !ordered(plain(&input, key.column)?.0) { + return Err(invalid("unsupported sort type")); + } + } + Ok(Self { + kind: Kind::Sort { keys, groups }, + inputs: vec![input.clone()], + output: input, + }) + } +} +#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)] +pub struct SortKey { + pub column: usize, + pub descending: bool, + pub nulls_first: bool, +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + Ok(futures::stream::once(async move { + let (rows, _memory) = collect_rows(input, &context).await?; + let result = match &operator.kind { + Kind::Sort { keys, groups } => { + let mut grouped = BTreeMap::>, Vec>>::new(); + let mut work = Cooperative::new(&context); + let mut workspace = Workspace::new(&context)?; + for row in rows { + work.checkpoint().await?; + for key in keys { + if matches!(row[key.column], Value::Map(_)) && nested_nan(&row[key.column]) + { + return Err(invalid("NaN in collection sort key")); + } + } + let key = group_key(&row, groups)?; + workspace.grow(std::mem::size_of::>())?; + if !grouped.contains_key(&key) { + workspace.grow(key_bytes(&key))?; + } + grouped.entry(key).or_default().push(row); + } + let mut result = Vec::new(); + for rows in grouped.into_values() { + let rows = + cooperative_sort(rows, |a, b| compare_rows(a, b, keys), &context).await?; + result.extend(rows); + } + result + } + _ => unreachable!(), + }; + Batch::try_new(output, result) + }) + .boxed_local()) +} + +fn compare_rows(a: &[Value], b: &[Value], keys: &[SortKey]) -> std::cmp::Ordering { + use std::cmp::Ordering::*; + for key in keys { + let (a, b) = (&a[key.column], &b[key.column]); + let order = match (a, b) { + (Value::Null, Value::Null) => Equal, + (Value::Null, _) => { + if key.nulls_first { + Less + } else { + Greater + } + } + (_, Value::Null) => { + if key.nulls_first { + Greater + } else { + Less + } + } + (Value::Float64(a), Value::Float64(b)) if a.is_nan() || b.is_nan() => { + match (a.is_nan(), b.is_nan()) { + (true, true) => Equal, + (true, false) => Greater, + _ => Less, + } + } + _ => { + let order = a.compare(b).expect("bound ordered types"); + if key.descending { + order.reverse() + } else { + order + } + } + }; + if order != Equal { + return order; + } + } + Equal +} +fn nested_nan(value: &Value) -> bool { + match value { + Value::Float64(value) => value.is_nan(), + Value::Map(values) => values + .iter() + .any(|(key, value)| nested_nan(key) || nested_nan(value)), + Value::List(values) | Value::Struct(values) => values.iter().any(nested_nan), + _ => false, + } +} +/// Stable in-memory merge sort with bounded synchronous chunks. Scratch storage +/// is reserved before allocation; comparisons yield between merge steps. +pub(super) async fn cooperative_sort( + rows: Vec, + compare: impl Fn(&T, &T) -> std::cmp::Ordering, + context: &RunContext, +) -> Result, Error> { + use std::collections::VecDeque; + let bytes = rows + .len() + .checked_mul(std::mem::size_of::() + std::mem::size_of::>()) + .and_then(|n| n.checked_mul(3)) + .ok_or(Error::MemoryLimit)?; + let _scratch = context.reserve(bytes)?; + let mut work = Cooperative::new(context); + let mut rows = rows.into_iter(); + let mut runs = VecDeque::new(); + loop { + work.checkpoint().await?; + let mut chunk = rows.by_ref().take(256).collect::>(); + if chunk.is_empty() { + break; + } + chunk.sort_by(&compare); + runs.push_back(VecDeque::from(chunk)); + } + // Merge adjacent runs in rounds to preserve ties in original input order. + while runs.len() > 1 { + let mut next = VecDeque::new(); + while let Some(mut left) = runs.pop_front() { + let Some(mut right) = runs.pop_front() else { + next.push_back(left); + break; + }; + let mut merged = VecDeque::with_capacity(left.len() + right.len()); + while !left.is_empty() || !right.is_empty() { + work.checkpoint().await?; + let take_left = match (left.front(), right.front()) { + (Some(a), Some(b)) => !compare(a, b).is_gt(), + (Some(_), None) => true, + _ => false, + }; + merged.push_back(if take_left { + left.pop_front().unwrap() + } else { + right.pop_front().unwrap() + }); + } + next.push_back(merged); + } + runs = next; + } + Ok(runs.pop_front().unwrap_or_default().into()) +} diff --git a/crates/asap-physical-operators/src/operators/source.rs b/crates/asap-physical-operators/src/operators/source.rs new file mode 100644 index 00000000..42701e17 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/source.rs @@ -0,0 +1,96 @@ +use super::*; +impl Operator { + pub fn source(output: Schema, batches: Vec) -> Result { + crate::values::validate_schema(&output)?; + if batches.iter().any(|b| b.schema() != &output) { + return Err(invalid("source schema mismatch")); + } + Ok(Self { + kind: Kind::Source(batches), + inputs: vec![], + output, + }) + } + pub fn scalar(value: Value, dtype: DataType) -> Result { + let output = schema(vec![result_field( + "value", + dtype.clone(), + matches!(value, Value::Null), + )]); + Batch::try_new(output.clone(), vec![vec![value.clone()]])?; + Ok(Self { + kind: Kind::Constant { value, dtype }, + inputs: vec![], + output, + }) + } + pub fn vector_to_scalar(input: Schema, column: usize) -> Result { + if plain(&input, column)? != (&DataType::Float64, false) { + return Err(invalid("scalar conversion requires non-null Float64")); + } + Ok(Self { + kind: Kind::VectorToScalar { column }, + inputs: vec![input], + output: schema(vec![result_field("value", DataType::Float64, false)]), + }) + } + pub fn union(input: Schema, arity: usize) -> Result { + if arity == 0 { + return Err(invalid("union needs at least one input")); + } + Ok(Self { + kind: Kind::Union, + inputs: vec![input.clone(); arity], + output: input, + }) + } +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + if let Kind::Source(batches) = &operator.kind { + return Ok(futures::stream::iter(batches.iter().cloned().map(Ok)).boxed_local()); + } + if let Kind::Constant { value, .. } = &operator.kind { + return Ok(futures::stream::once(async move { + Batch::try_new(output, vec![vec![value.clone()]]) + }) + .boxed_local()); + } + if matches!(operator.kind, Kind::Union) { + return Ok(futures::stream::select_all(inputs) + .map(|batch| batch.map(|batch| batch.value().clone())) + .boxed_local()); + } + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + match &operator.kind { + Kind::VectorToScalar { column } => Ok(futures::stream::once(async move { + let mut input = input; + let mut work = Cooperative::new(&context); + let mut value = f64::NAN; + let mut count = 0usize; + while let Some(batch) = input.next().await { + for row in batch?.rows() { + work.checkpoint().await?; + count = count.saturating_add(1); + if let Value::Float64(v) = row[*column] { + value = v; + } + } + } + Batch::try_new( + output, + vec![vec![Value::Float64(if count == 1 { + value + } else { + f64::NAN + })]], + ) + }) + .boxed_local()), + _ => unreachable!(), + } +} diff --git a/crates/asap-physical-operators/src/operators/summary/mod.rs b/crates/asap-physical-operators/src/operators/summary/mod.rs new file mode 100644 index 00000000..cc87a75a --- /dev/null +++ b/crates/asap-physical-operators/src/operators/summary/mod.rs @@ -0,0 +1,567 @@ +use super::*; +/// A summary readout: a sketch query, or an exact readout with typed parameters. +#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] +pub enum ReadoutQuery { + Sketch(planner_types::post_asap::SketchQuery), + Exact(crate::summary_kernels::exact::ExactReadout), +} + +impl Operator { + pub fn keyed_summary_build( + input: Schema, + family: SummaryFamilyType, + value: usize, + items: Vec, + groups: Vec, + ) -> Result { + use crate::summary_kernels::weighted_frequency::WeightedFrequency; + crate::values::validate_family(&family)?; + let SummaryFamilyType::Sketch(kind, _) = &family else { + return Err(invalid("keyed sketch required")); + }; + WeightedFrequency::configuration(kind)?; + validate_groups(&input, &groups)?; + if items.is_empty() || plain(&input, value)? != (&DataType::Float64, false) { + return Err(invalid( + "keyed summary requires identities and non-null Float64 weights", + )); + } + for &item in &items { + if !matches!( + plain(&input, item)?.0, + DataType::Utf8 + | DataType::Timestamp + | DataType::Int64 + | DataType::Float64 + | DataType::Bool + | DataType::Null + ) { + return Err(invalid("unsupported keyed summary identity type")); + } + } + let mut fields = groups + .iter() + .map(|&i| input.fields[i].clone()) + .collect::>(); + fields.push(SummaryField { + name: "state".into(), + dtype: family.clone(), + nullable: false, + }); + Ok(Self { + kind: Kind::KeyedSummaryBuild { + family, + value, + items, + groups, + }, + inputs: vec![input], + output: schema(fields), + }) + } + pub fn keyed_readout( + input: Schema, + state: usize, + k: usize, + output: Schema, + ) -> Result { + use crate::summary_kernels::weighted_frequency::WeightedFrequency; + crate::values::validate_family(&field(&input, state)?.dtype)?; + let SummaryFamilyType::Sketch(kind, _) = &field(&input, state)?.dtype else { + return Err(invalid("keyed readout requires summary state")); + }; + let (_, _, _, capacity) = WeightedFrequency::configuration(kind)?; + if k > capacity || output.fields.len() <= input.fields.len() { + return Err(invalid("invalid keyed readout shape or capacity")); + } + if state + 1 != input.fields.len() + || output.fields[..state] != input.fields[..state] + || output.fields.last().unwrap().dtype != SummaryFamilyType::Plain(DataType::Float64) + { + return Err(invalid( + "keyed readout must preserve partitions and return a Float64 score", + )); + } + crate::values::validate_schema(&output)?; + Ok(Self { + kind: Kind::KeyedReadout { state, k }, + inputs: vec![input], + output, + }) + } + pub fn summary_build( + input: Schema, + family: SummaryFamilyType, + value: usize, + time: Option, + groups: Vec, + ) -> Result { + crate::values::validate_family(&family)?; + validate_groups(&input, &groups)?; + if plain(&input, value)?.0 != &DataType::Float64 { + return Err(invalid("summary numeric update requires Float64")); + } + if let Some(time) = time { + if plain(&input, time)? != (&DataType::Timestamp, false) { + return Err(invalid("summary time column must be a timestamp")); + } + } + if time.is_none() + && matches!( + family, + SummaryFamilyType::ExactAggregate( + planner_types::post_asap::ExactKind::Rate + | planner_types::post_asap::ExactKind::Increase, + _ + ) + ) + { + return Err(invalid("counter summary requires a timestamp column")); + } + crate::capability::validate_summary_kernel( + &family, + &SummaryUpdate::column(ColumnRef::SampleValue), + &Default::default(), + ) + .map_err(Error::Invalid)?; + let mut fields = groups + .iter() + .map(|&i| input.fields[i].clone()) + .collect::>(); + fields.push(SummaryField { + name: "state".into(), + dtype: family.clone(), + nullable: false, + }); + Ok(Self { + kind: Kind::SummaryBuild { + family, + value, + time, + groups, + }, + inputs: vec![input], + output: schema(fields), + }) + } + pub fn summary_merge(input: Schema, state: usize, groups: Vec) -> Result { + validate_groups(&input, &groups)?; + crate::values::validate_family(&field(&input, state)?.dtype)?; + if matches!(field(&input, state)?.dtype, SummaryFamilyType::Plain(_)) { + return Err(invalid("summary state required")); + } + let mut fields = groups + .iter() + .map(|&i| input.fields[i].clone()) + .collect::>(); + fields.push(input.fields[state].clone()); + Ok(Self { + kind: Kind::SummaryMerge { state, groups }, + inputs: vec![input], + output: schema(fields), + }) + } + pub fn readout(input: Schema, state: usize, query: ReadoutQuery) -> Result { + let family = &field(&input, state)?.dtype; + crate::values::validate_family(family)?; + match &query { + ReadoutQuery::Sketch(query) => { + crate::capability::validate_sketch_readout(family, query)? + } + ReadoutQuery::Exact(readout) => { + crate::capability::validate_exact_readout(family, readout)? + } + } + let mut fields = input.fields.clone(); + let result_type = if matches!( + fields[state].dtype, + SummaryFamilyType::ExactAggregate(planner_types::post_asap::ExactKind::Count, _) + ) { + DataType::Int64 + } else { + DataType::Float64 + }; + // A state-only row represents the global population. Its extrema may + // be empty, just like an ordinary ungrouped MIN/MAX aggregate. + let nullable = fields.len() == 1 + && matches!( + fields[state].dtype, + SummaryFamilyType::ExactAggregate( + planner_types::post_asap::ExactKind::Min + | planner_types::post_asap::ExactKind::Max, + _ + ) + ); + fields[state] = result_field("value", result_type, nullable); + Ok(Self { + kind: Kind::Readout { state, query }, + inputs: vec![input], + output: schema(fields), + }) + } +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let range_ms = operator.readout_range(&context)?; + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + match &operator.kind { + Kind::SummaryBuild { + family, + value, + time, + groups, + } => Ok(futures::stream::once(async move { + Batch::try_new( + output, + build_summary(input, family, *value, *time, groups, &context).await?, + ) + }) + .boxed_local()), + Kind::KeyedSummaryBuild { + family, + value, + items, + groups, + } => Ok(futures::stream::once(async move { + Batch::try_new( + output, + build_keyed_summary(input, family, *value, items, groups, &context).await?, + ) + }) + .boxed_local()), + Kind::KeyedReadout { state, k } => Ok(input + .map(move |batch| { + let batch = batch?; + let mut rows = Vec::new(); + for row in batch.rows() { + let Value::Summary { state: summary, .. } = &row[*state] else { + return Err(invalid("summary value required")); + }; + let summary = summary + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("weighted frequency typed state required"))?; + for items in summary.rows(*k) { + let mut values = row[..*state].to_vec(); + values.extend(items); + // The typed output schema restores epoch-millisecond + // timestamp keys from the kernel's Int64 representation. + for (value, field) in values.iter_mut().zip(&output.fields) { + if field.dtype == SummaryFamilyType::Plain(DataType::Timestamp) { + if let Value::Int64(time) = value { + *value = Value::Timestamp(*time); + } + } + } + rows.push(values); + } + } + Batch::try_new(output.clone(), rows) + }) + .boxed_local()), + Kind::Readout { state, query } => Ok(input + .map(move |batch| { + let batch = batch?; + let mut rows = batch.rows().to_vec(); + if let ReadoutQuery::Exact(readout) = query { + rows.retain(|row| !matches!(&row[*state], Value::Summary { state: summary, .. } + if crate::readout::insufficient_counter_samples(summary.as_ref(), readout.statistic))); + } + for row in &mut rows { + let Value::Summary { state: summary, .. } = &row[*state] else { + return Err(invalid("summary value required")); + }; + row[*state] = match query { + ReadoutQuery::Sketch(query) => Value::Float64( + summary + .estimate(query) + .map_err(|e| Error::Operator(e.to_string()))?, + ), + ReadoutQuery::Exact(readout) => { + let exact = summary + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("exact readout requires exact state"))?; + if output.fields[*state].dtype == SummaryFamilyType::Plain(DataType::Int64) { + let count = exact.count().ok_or_else(|| { + Error::Operator("exact count state lacks an integer count".into()) + })?; + Value::Int64(i64::try_from(count).map_err(|_| { + Error::Operator("exact count exceeds Int64".into()) + })?) + } else { + match exact + .readout(readout.statistic, range_ms, None) + .map_err(|e| Error::Operator(e.to_string()))? + { + Some(value) => Value::Float64(value), + None if output.fields[*state].nullable => Value::Null, + None => { + return Err(Error::Operator( + "empty exact population".into(), + )) + } + } + } + } + }; + } + Batch::try_new(output.clone(), rows) + }) + .boxed_local()), + _ => unreachable!(), + } +} +pub(super) fn execute_merge<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + Ok(futures::stream::once(async move { + let (rows, _memory) = collect_rows(input, &context).await?; + let result = match &operator.kind { + Kind::SummaryMerge { state, groups } => { + merge_summary(rows, *state, groups, &context).await? + } + _ => unreachable!(), + }; + Batch::try_new(output, result) + }) + .boxed_local()) +} + +async fn build_summary( + mut input: Input<'_, Batch>, + family: &SummaryFamilyType, + value: usize, + time: Option, + groups: &[usize], + context: &RunContext, +) -> Result>, Error> { + type State = ( + Vec, + Box, + Reservation, + usize, + Option, + ); + let create = |labels: Vec, key_bytes: usize| -> Result { + let updater = crate::factory::create_planner_accumulator( + family, + &SummaryUpdate::column(ColumnRef::SampleValue), + &Default::default(), + ) + .map_err(Error::Operator)?; + let overhead = labels.iter().map(Value::bytes).sum::() + key_bytes + 64; + let memory = context.reserve(updater.memory_usage_bytes() + overhead)?; + Ok((labels, updater, memory, overhead, None)) + }; + let mut work = Cooperative::new(context); + let mut states = BTreeMap::>, State>::new(); + if groups.is_empty() { + states.insert(vec![], create(vec![], 0)?); + } + let ordered_time = matches!( + family, + SummaryFamilyType::ExactAggregate( + planner_types::post_asap::ExactKind::Rate + | planner_types::post_asap::ExactKind::Increase, + _ + ) + ); + while let Some(batch) = input.next().await { + let batch = batch?; + for row in batch.rows() { + work.checkpoint().await?; + let key = group_key(row, groups)?; + if !states.contains_key(&key) { + let labels = groups.iter().map(|&i| row[i].clone()).collect(); + let state = create( + labels, + key.iter() + .map(|v| v.len() + std::mem::size_of::>()) + .sum(), + )?; + states.insert(key.clone(), state); + } + let (_, updater, memory, overhead, previous) = + states.get_mut(&key).expect("inserted group"); + // SQL aggregates ignore NULL samples while retaining the group. + // A missing counter sample also contributes no observation. + let value = match row[value] { + Value::Float64(value) => value, + Value::Null => continue, + _ => return Err(invalid("summary update type")), + }; + let timestamp = if let Some(time) = time { + let Value::Timestamp(time) = row[time] else { + return Err(invalid("summary time type")); + }; + time + } else { + 0 + }; + if ordered_time && previous.is_some_and(|prior| timestamp <= prior) { + return Err(Error::Operator( + "counter samples must have strictly increasing timestamps within each group" + .into(), + )); + } + updater + .validate_single_input(value) + .map_err(Error::Operator)?; + updater.update_single(value, timestamp); + *previous = Some(timestamp); + memory.resize(updater.memory_usage_bytes() + *overhead)?; + } + } + Ok(states + .into_values() + .map(|(mut labels, updater, _memory, _, _)| { + labels.push(Value::Summary { + family: family.clone(), + state: Arc::from(updater.into_accumulator()), + }); + labels + }) + .collect()) +} +async fn merge_summary( + rows: Vec>, + state_column: usize, + groups: &[usize], + context: &RunContext, +) -> Result>, Error> { + type GroupState = (Vec, SummaryFamilyType, Arc); + let mut states: BTreeMap>, GroupState> = BTreeMap::new(); + let mut work = Cooperative::new(context); + let mut memory = context.reserve(0)?; + let mut retained = 0usize; + for row in rows { + work.checkpoint().await?; + let Value::Summary { family, state } = &row[state_column] else { + return Err(invalid("summary state required")); + }; + let key = group_key(&row, groups)?; + if let Some((_, expected, existing)) = states.get_mut(&key) { + if expected != family { + return Err(invalid("incompatible summary family")); + } + let old_bytes = existing.approx_memory_bytes(); + // Reserve an estimate for the replacement while both input states remain live. + memory.resize( + retained + .checked_add(old_bytes) + .and_then(|n| n.checked_add(state.approx_memory_bytes())) + .ok_or(Error::MemoryLimit)?, + )?; + *existing = Arc::from( + existing + .merge_with(state.as_ref()) + .map_err(|e| Error::Operator(e.to_string()))?, + ); + retained = retained + .checked_sub(old_bytes) + .and_then(|n| n.checked_add(existing.approx_memory_bytes())) + .ok_or(Error::MemoryLimit)?; + memory.resize(retained)?; + } else { + retained = retained + .checked_add(key_bytes(&key) + row_bytes(&row)) + .ok_or(Error::MemoryLimit)?; + memory.resize(retained)?; + states.insert( + key, + ( + groups.iter().map(|&i| row[i].clone()).collect(), + family.clone(), + state.clone(), + ), + ); + } + } + Ok(states + .into_values() + .map(|(mut keys, family, state)| { + keys.push(Value::Summary { family, state }); + keys + }) + .collect()) +} + +async fn build_keyed_summary( + mut input: Input<'_, Batch>, + family: &SummaryFamilyType, + value: usize, + items: &[usize], + groups: &[usize], + context: &RunContext, +) -> Result>, Error> { + use crate::{summary_kernels::weighted_frequency::WeightedFrequency, AggregateCore}; + let SummaryFamilyType::Sketch(kind, _) = family else { + unreachable!() + }; + let (algorithm, width, depth, capacity) = WeightedFrequency::configuration(kind)?; + let mut work = Cooperative::new(context); + let mut states = + BTreeMap::>, (Vec, WeightedFrequency, Reservation, usize)>::new(); + while let Some(batch) = input.next().await { + let batch = batch?; + for row in batch.rows() { + work.checkpoint().await?; + let key = group_key(row, groups)?; + if !states.contains_key(&key) { + let labels = groups.iter().map(|&i| row[i].clone()).collect::>(); + let overhead = labels.iter().map(Value::bytes).sum::() + + key.iter().map(|v| v.len() + 24).sum::() + + 128; + let bytes = width + .checked_mul(depth) + .and_then(|n| n.checked_mul(8)) + .and_then(|n| n.checked_add(overhead)) + .ok_or_else(|| invalid("weighted frequency memory size overflow"))?; + let reservation = context.reserve(bytes)?; + states.insert( + key.clone(), + ( + labels, + WeightedFrequency::new(algorithm, width, depth, capacity)?, + reservation, + overhead, + ), + ); + } + let (_, summary, reservation, overhead) = states.get_mut(&key).unwrap(); + let Value::Float64(weight) = row[value] else { + return Err(invalid("weighted frequency weight type")); + }; + summary.update( + &items + .iter() + .map(|&i| match &row[i] { + Value::Timestamp(time) => Value::Int64(*time), + value => value.clone(), + }) + .collect::>(), + weight, + )?; + reservation.resize(summary.approx_memory_bytes() + *overhead)?; + } + } + Ok(states + .into_values() + .map(|(mut labels, summary, _, _)| { + labels.push(Value::Summary { + family: family.clone(), + state: Arc::new(summary), + }); + labels + }) + .collect()) +} diff --git a/crates/asap-physical-operators/src/operators/unchecked.rs b/crates/asap-physical-operators/src/operators/unchecked.rs new file mode 100644 index 00000000..2ecbaabe --- /dev/null +++ b/crates/asap-physical-operators/src/operators/unchecked.rs @@ -0,0 +1,140 @@ +//! Deserialized operators are validated before use, whatever the encoding. +use super::*; + +#[derive(serde::Deserialize)] +#[serde(deny_unknown_fields)] +pub(super) struct UncheckedOperator { + kind: Kind, + inputs: Vec, + output: Schema, +} +impl TryFrom for Operator { + type Error = Error; + fn try_from(unchecked: UncheckedOperator) -> Result { + let expected_kind = + serde_json::to_value(&unchecked.kind).map_err(|error| invalid(&error.to_string()))?; + let UncheckedOperator { + kind, + inputs, + output, + } = unchecked; + for schema in inputs.iter().chain(std::iter::once(&output)) { + crate::values::validate_schema(schema)?; + } + let input = |index| { + inputs + .get(index) + .cloned() + .ok_or_else(|| invalid("missing operator input")) + }; + let op = match kind { + Kind::Source(_) => return Err(invalid("physical plans cannot serialize live sources")), + Kind::Constant { value, dtype } => Operator::scalar(value, dtype)?, + Kind::PaneInput { + coordinate, + layout, + offset_ms, + } => Operator::pane_input(input(0)?, coordinate, layout, offset_ms)?, + Kind::ScopeTimestamp { .. } => Operator::scope_timestamp(input(0)?, output.clone())?, + Kind::Union => Operator::union(input(0)?, inputs.len())?, + Kind::CurrentSeries { + identity, + coordinate, + value, + lookback_ms, + } => Operator::current_series(input(0)?, identity, coordinate, value, lookback_ms)?, + Kind::VectorToScalar { column } => Operator::vector_to_scalar(input(0)?, column)?, + Kind::VectorBinary { + operator, + return_bool, + } => Operator::vector_binary(input(0)?, input(1)?, operator, return_bool)?, + Kind::AlignedBinary { + keys, + values, + operator, + } => Operator::aligned_binary(input(0)?, input(1)?, keys, values, operator)?, + Kind::RangeWindow { intent } => Operator::range_window(*intent)?, + Kind::HistogramQuantile => Operator::histogram_quantile(), + Kind::Project(expressions) => { + if expressions.len() != output.fields.len() { + return Err(invalid("projection width mismatch")); + } + Operator::project( + input(0)?, + output + .fields + .iter() + .zip(expressions) + .map(|(f, e)| (f.name.clone(), e)) + .collect(), + )? + } + Kind::Filter(expression) => Operator::filter(input(0)?, expression)?, + Kind::Limit { n, offset, groups } => Operator::limit(input(0)?, n, offset, groups)?, + Kind::Sort { keys, groups } => Operator::sort(input(0)?, keys, groups)?, + Kind::Window { + intent, + coordinate, + value, + groups, + window, + } => Operator::window(input(0)?, *intent, coordinate, value, groups, window)?, + Kind::Aggregate { groups, measures } => { + if groups.len() + measures.len() != output.fields.len() { + return Err(invalid("aggregate width mismatch")); + } + let names = output.fields[groups.len()..].iter().map(|f| f.name.clone()); + Operator::aggregate(input(0)?, groups, names.zip(measures).collect())? + } + Kind::SemiJoin { + keys, + require_complete_right, + } => { + let operator = Operator::semi_join(input(0)?, input(1)?, keys)?; + if require_complete_right { + operator.require_complete_right() + } else { + operator + } + } + Kind::Join { kind, predicate } => Operator::relational_join( + input(0)?, + input(1)?, + kind, + &planner_types::pre_asap::Predicate(std::rc::Rc::new( + predicate.expression().clone(), + )), + output.clone(), + )?, + Kind::SummaryBuild { + family, + value, + time, + groups, + } => Operator::summary_build(input(0)?, family, value, time, groups)?, + Kind::KeyedSummaryBuild { + family, + value, + items, + groups, + } => Operator::keyed_summary_build(input(0)?, family, value, items, groups)?, + Kind::KeyedReadout { state, k } => { + Operator::keyed_readout(input(0)?, state, k, output.clone())? + } + Kind::SummaryMerge { state, groups } => { + Operator::summary_merge(input(0)?, state, groups)? + } + Kind::Readout { state, query } => Operator::readout(input(0)?, state, query)?, + } + .with_output_schema(output)?; + if serde_json::to_value(&op.kind).map_err(|error| invalid(&error.to_string()))? + != expected_kind + { + return Err(invalid("operator contains inconsistent compiled fields")); + } + if op.inputs != inputs { + return Err(invalid("operator input contracts differ")); + } + Ok(op) + } +} diff --git a/crates/asap-physical-operators/src/operators/vector_binary.rs b/crates/asap-physical-operators/src/operators/vector_binary.rs new file mode 100644 index 00000000..cc60ff4c --- /dev/null +++ b/crates/asap-physical-operators/src/operators/vector_binary.rs @@ -0,0 +1,237 @@ +//! Label matching and scalar broadcasting are physical computation, not source binding. +use super::*; +use planner_types::{post_asap::BinaryOperator, pre_asap::BinaryOpKind}; + +pub(crate) fn value_schema(scalar: bool) -> Schema { + let mut fields = Vec::new(); + if !scalar { + fields.push(result_field( + "labels", + DataType::Map { + key: Box::new(DataType::Utf8), + value: Box::new(DataType::Utf8), + value_nullable: false, + }, + false, + )); + } + fields.push(result_field( + if scalar { "$promql_scalar" } else { "value" }, + DataType::Float64, + false, + )); + schema(fields) +} + +fn is_scalar(input: &Schema) -> Result { + for scalar in [true, false] { + let expected = value_schema(scalar); + if input.fields.len() == expected.fields.len() + && input + .fields + .iter() + .zip(&expected.fields) + .all(|(a, b)| a.dtype == b.dtype && !a.nullable) + { + return Ok(scalar); + } + } + Err(invalid( + "vector binary requires Float64 scalars or complete label-map vectors", + )) +} + +impl Operator { + pub fn vector_binary( + left: Schema, + right: Schema, + operator: BinaryOperator, + return_bool: bool, + ) -> Result { + let scalar = is_scalar(&left)? && is_scalar(&right)?; + is_scalar(&right)?; + let expression = Expression::Binary { + operator: operator.clone(), + left: Box::new(Expression::Column(0)), + right: Box::new(Expression::Column(1)), + }; + expression.dtype(&schema(vec![ + result_field("left", DataType::Float64, false), + result_field("right", DataType::Float64, false), + ]))?; + let comparison = matches!(operator.kind, BinaryOpKind::Compare(_)); + if (return_bool && !comparison) || (scalar && comparison && !return_bool) { + return Err(invalid("invalid scalar/vector comparison bool mode")); + } + Ok(Self { + inputs: vec![left, right], + output: value_schema(scalar), + kind: Kind::VectorBinary { + operator, + return_bool, + }, + }) + } +} + +type Labels = BTreeMap, Arc>; +fn labels(row: &[Value]) -> Result { + let Some(Value::Map(entries)) = row.first() else { + return Err(invalid("vector requires label map")); + }; + let mut result = BTreeMap::new(); + for (key, value) in entries.iter() { + let (Value::Utf8(key), Value::Utf8(value)) = (key, value) else { + return Err(invalid("labels must be Utf8")); + }; + if result.insert(key.clone(), value.clone()).is_some() { + return Err(invalid("duplicate label name")); + } + } + Ok(result) +} +fn identity(mut labels: Labels) -> Labels { + labels.remove("__name__"); + labels.retain(|_, value| !value.is_empty()); + labels +} +fn value(row: &[Value]) -> Result { + match row.last() { + Some(Value::Float64(value)) => Ok(*value), + _ => Err(invalid("binary value must be Float64")), + } +} +fn label_bytes(labels: &Labels) -> usize { + labels.iter().map(|(k, v)| 64 + k.len() + v.len()).sum() +} + +pub(super) fn execute<'a>( + op: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let Kind::VectorBinary { + operator, + return_bool, + } = &op.kind + else { + unreachable!() + }; + let left_scalar = is_scalar(&op.inputs[0])?; + let right_scalar = is_scalar(&op.inputs[1])?; + let right = inputs.pop().ok_or_else(|| invalid("missing right input"))?; + let left = inputs.pop().ok_or_else(|| invalid("missing left input"))?; + Ok(futures::stream::once(async move { + let ((left, _left_memory), (right, _right_memory)) = + futures::try_join!(collect_rows(left, &context), collect_rows(right, &context))?; + if (left_scalar && left.len() != 1) || (right_scalar && right.len() != 1) { + return Err(invalid("scalar input must contain exactly one value")); + } + let mut workspace = Workspace::new(&context)?; + let mut work = Cooperative::new(&context); + let mut rows = Vec::new(); + let mut result_identities = std::collections::BTreeSet::new(); + let mut emit = |labels: Labels, a: f64, b: f64| -> Result<(), Error> { + let arithmetic = matches!(operator.kind, BinaryOpKind::Arithmetic(_)); + let result = match crate::expressions::arithmetic::evaluate_binary(operator, a, b)? { + Value::Float64(value) => value, + Value::Bool(value) if *return_bool => { + if value { + 1. + } else { + 0. + } + } + Value::Bool(true) => { + if left_scalar { + b + } else { + a + } + } + Value::Bool(false) => return Ok(()), + _ => return Err(invalid("invalid binary result")), + }; + let mut row = Vec::new(); + if !left_scalar || !right_scalar { + let labels = if arithmetic || *return_bool { + identity(labels) + } else { + labels + }; + workspace.grow(label_bytes(&labels) + 64)?; + if !result_identities.insert(labels.clone()) { + return Err(invalid("duplicate vector result labels")); + } + workspace.grow( + label_bytes(&labels) + + std::mem::size_of::>() + + 2 * std::mem::size_of::(), + )?; + row.push(Value::Map( + labels + .into_iter() + .map(|(k, v)| (Value::Utf8(k), Value::Utf8(v))) + .collect::>() + .into(), + )); + } else { + workspace.grow(std::mem::size_of::>() + std::mem::size_of::())?; + } + row.push(Value::Float64(result)); + rows.push(row); + Ok(()) + }; + if left_scalar || right_scalar { + let vectors = if left_scalar { &right } else { &left }; + for row in vectors { + work.checkpoint().await?; + let labels = if left_scalar && right_scalar { + Labels::new() + } else { + labels(row)? + }; + emit( + labels, + if left_scalar { + value(&left[0])? + } else { + value(row)? + }, + if right_scalar { + value(&right[0])? + } else { + value(row)? + }, + )?; + } + } else { + let mut rhs = BTreeMap::new(); + // Keep matching workspace separate from the output reservation captured by emit. + let mut matching = Workspace::new(&context)?; + for row in &right { + work.checkpoint().await?; + let key = identity(labels(row)?); + matching.grow(label_bytes(&key) + 64)?; + if rhs.insert(key, value(row)?).is_some() { + return Err(invalid("duplicate vector matching labels")); + } + } + let mut seen = std::collections::BTreeSet::new(); + for row in &left { + work.checkpoint().await?; + let labels = labels(row)?; + let key = identity(labels.clone()); + matching.grow(label_bytes(&key) + 64)?; + if !seen.insert(key.clone()) { + return Err(invalid("duplicate vector matching labels")); + } + if let Some(b) = rhs.get(&key) { + emit(labels, value(row)?, *b)?; + } + } + } + Batch::try_new(op.output.clone(), rows) + }) + .boxed_local()) +} diff --git a/crates/asap-physical-operators/src/operators/vector_window.rs b/crates/asap-physical-operators/src/operators/vector_window.rs new file mode 100644 index 00000000..407f04ea --- /dev/null +++ b/crates/asap-physical-operators/src/operators/vector_window.rs @@ -0,0 +1,147 @@ +//! Window bounds are typed input data; aggregation and histogram semantics stay native. +use super::*; +use planner_types::pre_asap::AggIntent; + +pub(crate) fn matrix_schema() -> Schema { + let mut fields = vector_binary::value_schema(false).fields.clone(); + fields.insert(1, result_field("timestamp", DataType::Timestamp, false)); + fields.push(result_field("window_start", DataType::Timestamp, false)); + fields.push(result_field("window_end", DataType::Timestamp, false)); + Arc::new(SummarySchema { + fields, + time_index: Some(1), + }) +} + +impl Operator { + pub fn range_window(intent: AggIntent) -> Result { + // Reuse the window constructor's semantic admission, without fixing request time. + Self::window(matrix_schema(), intent.clone(), 1, 2, vec![0], Some((0, 1)))?; + Ok(Self { + kind: Kind::RangeWindow { + intent: Box::new(intent), + }, + inputs: vec![matrix_schema()], + output: vector_binary::value_schema(false), + }) + } + pub fn histogram_quantile() -> Self { + Self { + kind: Kind::HistogramQuantile, + inputs: vec![ + vector_binary::value_schema(true), + vector_binary::value_schema(false), + ], + output: vector_binary::value_schema(false), + } + } +} + +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + match &operator.kind { + Kind::RangeWindow { intent } => { + let input = inputs + .pop() + .ok_or_else(|| invalid("missing matrix input"))?; + Ok(futures::stream::once(async move { + let (rows, _memory) = collect_rows(input, &context).await?; + let mut window = None; + let mut work = Cooperative::new(&context); + for row in &rows { + work.checkpoint().await?; + let (Value::Timestamp(start), Value::Timestamp(end)) = (&row[3], &row[4]) + else { + return Err(invalid("missing matrix window bounds")); + }; + if start >= end || window.is_some_and(|bounds| bounds != (*start, *end)) { + return Err(invalid("matrix rows must share one nonempty window")); + } + window = Some((*start, *end)); + } + let mut rows = + aggregate::temporal::reduce(rows, intent, &[0], 1, 2, window, &context).await?; + for row in &mut rows { + work.checkpoint().await?; + row[1] = Expression::ExactFloat64(1).evaluate(row)?; + } + Batch::try_new(operator.output.clone(), rows) + }) + .boxed_local()) + } + Kind::HistogramQuantile => { + let buckets = inputs + .pop() + .ok_or_else(|| invalid("missing histogram buckets"))?; + let quantile = inputs.pop().ok_or_else(|| invalid("missing quantile"))?; + Ok(futures::stream::once(async move { + let ((quantile, _q_memory), (buckets, _bucket_memory)) = futures::try_join!( + collect_rows(quantile, &context), + collect_rows(buckets, &context) + )?; + let [row] = quantile.as_slice() else { + return Err(invalid("histogram quantile requires one scalar")); + }; + let [Value::Float64(q)] = row.as_slice() else { + return Err(invalid("invalid quantile scalar")); + }; + let mut rows = Vec::new(); + let mut work = Cooperative::new(&context); + let mut workspace = Workspace::new(&context)?; + for row in buckets { + work.checkpoint().await?; + let Value::Map(entries) = &row[0] else { + return Err(invalid("histogram buckets require labels")); + }; + let mut bound = None; + let mut labels = BTreeMap::new(); + let mut seen = std::collections::BTreeSet::new(); + for (key, value) in entries.iter() { + let (Value::Utf8(key), Value::Utf8(value)) = (key, value) else { + return Err(invalid("histogram labels must be Utf8")); + }; + if !seen.insert(key) { + return Err(invalid("duplicate histogram label")); + } + if key.as_ref() == "le" { + bound = value.parse::().ok(); + } else if key.as_ref() != "__name__" && !value.is_empty() { + labels.insert(key.clone(), value.clone()); + } + } + if let Some(bound) = bound { + let projected = vec![ + Value::Map( + labels + .into_iter() + .map(|(k, v)| (Value::Utf8(k), Value::Utf8(v))) + .collect::>() + .into(), + ), + Value::Float64(bound), + row[1].clone(), + ]; + workspace.grow(row_bytes(&projected))?; + rows.push(projected); + } + } + let result = aggregate::temporal::reduce( + rows, + &AggIntent::HistogramQuantile { q: *q }, + &[0], + 1, + 2, + None, + &context, + ) + .await?; + Batch::try_new(operator.output.clone(), result) + }) + .boxed_local()) + } + _ => unreachable!(), + } +} diff --git a/crates/asap-physical-operators/src/plan/mod.rs b/crates/asap-physical-operators/src/plan/mod.rs new file mode 100644 index 00000000..0dce4dd8 --- /dev/null +++ b/crates/asap-physical-operators/src/plan/mod.rs @@ -0,0 +1,160 @@ +//! Immutable physical graph, operator contracts and pre-execution validation. +use crate::{ + runtime::{Input, OutputStream, RunContext}, + Error, +}; +use std::{ + collections::{BTreeMap, BTreeSet}, + fmt::Debug, +}; +pub type NodeId = u64; +mod properties; +pub use properties::{Boundedness, Emission, PlanProperties}; +/// Operators own computation. The runtime provides already-connected inputs; +/// an operator must not recursively execute another plan node itself. +pub trait PhysicalOperator { + fn name(&self) -> &str; + /// Source implementations must explicitly declare finite input before feeding blocking operators. + fn properties(&self, inputs: &[PlanProperties]) -> PlanProperties { + PlanProperties { + boundedness: Boundedness::from_inputs(inputs), + emission: Emission::Unknown, + } + } + fn requires_bounded_input(&self) -> bool { + false + } + + /// Validate run-specific contracts before any source is opened. + fn validate_context(&self, _context: &RunContext) -> Result<(), Error> { + Ok(()) + } + fn input_schemas(&self) -> Vec; + fn output_schema(&self) -> S; + fn start<'a>( + &'a self, + inputs: Vec>, + context: RunContext, + ) -> Result, Error>; + fn output_bytes(&self, value: &V) -> usize; +} +pub(crate) struct Node<'a, V, S> { + pub(crate) inputs: Vec, + pub(crate) operator: Box + 'a>, +} +pub struct PhysicalDag<'a, V, S> { + pub(crate) nodes: BTreeMap>, +} +impl Default for PhysicalDag<'_, V, S> { + fn default() -> Self { + Self { + nodes: BTreeMap::new(), + } + } +} +impl<'a, V: 'a, S: Clone + PartialEq + Debug + 'a> PhysicalDag<'a, V, S> { + pub fn add( + &mut self, + id: NodeId, + inputs: Vec, + operator: impl PhysicalOperator + 'a, + ) -> Result<(), Error> { + self.add_boxed(id, inputs, Box::new(operator)) + } + pub fn add_boxed( + &mut self, + id: NodeId, + inputs: Vec, + operator: Box + 'a>, + ) -> Result<(), Error> { + if self.nodes.contains_key(&id) { + return Err(Error::Invalid(format!("duplicate node {id}"))); + } + self.nodes.insert(id, Node { inputs, operator }); + Ok(()) + } + pub fn validate(&self, roots: &[NodeId]) -> Result<(), Error> { + self.properties(roots).map(|_| ()) + } + /// Derive properties while checking topology and schemas, before starting sources. + pub fn properties(&self, roots: &[NodeId]) -> Result, Error> { + fn visit( + dag: &PhysicalDag<'_, V, S>, + id: NodeId, + active: &mut BTreeSet, + done: &mut BTreeMap, + ) -> Result { + if let Some((depth, _)) = done.get(&id) { + return Ok(*depth); + } + if active.len() >= 128 { + return Err(Error::Invalid( + "DAG exceeds the supported execution depth of 128".into(), + )); + } + if !active.insert(id) { + return Err(Error::Invalid(format!("cycle at node {id}"))); + } + let node = dag + .nodes + .get(&id) + .ok_or_else(|| Error::Invalid(format!("missing node {id}")))?; + let expected = node.operator.input_schemas(); + if expected.len() != node.inputs.len() { + return Err(Error::Invalid(format!("node {id} input arity mismatch"))); + } + let mut depth = 1; + let mut input_properties = Vec::new(); + for (input, schema) in node.inputs.iter().zip(expected) { + depth = depth.max(1 + visit(dag, *input, active, done)?); + input_properties.push(done[input].1); + let actual = dag.nodes[input].operator.output_schema(); + if actual != schema { + return Err(Error::Invalid(format!( + "node {id} input {input} schema mismatch: {actual:?} vs {schema:?}" + ))); + } + } + if depth > 128 { + return Err(Error::Invalid( + "DAG exceeds the supported execution depth of 128".into(), + )); + } + if node.operator.requires_bounded_input() + && input_properties + .iter() + .any(|p| p.boundedness != Boundedness::Bounded) + { + return Err(Error::Invalid(format!( + "node {id} ({}) requires bounded inputs", + node.operator.name() + ))); + } + let properties = node.operator.properties(&input_properties); + active.remove(&id); + done.insert(id, (depth, properties)); + Ok(depth) + } + if roots.is_empty() { + return Err(Error::Invalid("execution needs a root".into())); + } + let mut done = BTreeMap::new(); + for &root in roots { + visit(self, root, &mut BTreeSet::new(), &mut done)?; + } + Ok(done + .into_iter() + .map(|(id, (_, properties))| (id, properties)) + .collect()) + } + pub fn execute<'r>( + &'r self, + roots: &[NodeId], + context: RunContext, + ) -> Result>, Error> + where + 'a: 'r, + { + crate::runtime::execute(self, roots, context) + } +} diff --git a/crates/asap-physical-operators/src/plan/properties.rs b/crates/asap-physical-operators/src/plan/properties.rs new file mode 100644 index 00000000..7ec770ea --- /dev/null +++ b/crates/asap-physical-operators/src/plan/properties.rs @@ -0,0 +1,32 @@ +//! Execution facts used to reject operators that cannot finish on their inputs. +#[derive(serde::Serialize, serde::Deserialize, Clone, Copy, Debug, PartialEq, Eq)] +pub enum Boundedness { + /// The source or operator promises a finite result for this run. + Bounded, + Unbounded, + /// No finite-input guarantee has been supplied. + Unknown, +} +impl Boundedness { + pub fn from_inputs(inputs: &[PlanProperties]) -> Self { + if inputs.iter().any(|p| p.boundedness == Self::Unbounded) { + Self::Unbounded + } else if inputs.is_empty() || inputs.iter().any(|p| p.boundedness == Self::Unknown) { + Self::Unknown + } else { + Self::Bounded + } + } +} +#[derive(serde::Serialize, serde::Deserialize, Clone, Copy, Debug, PartialEq, Eq)] +pub enum Emission { + Incremental, + /// Produces its result only after all inputs end, even if accumulation is incremental. + AfterInput, + Unknown, +} +#[derive(serde::Serialize, serde::Deserialize, Clone, Copy, Debug, PartialEq, Eq)] +pub struct PlanProperties { + pub boundedness: Boundedness, + pub emission: Emission, +} diff --git a/crates/asap-physical-operators/src/readout.rs b/crates/asap-physical-operators/src/readout.rs new file mode 100644 index 00000000..e884ecc3 --- /dev/null +++ b/crates/asap-physical-operators/src/readout.rs @@ -0,0 +1,111 @@ +//! Readouts over merged exact summary states. +use crate::summary_kernels::exact::ExactAccumulator; +use crate::{AggregateCore, KeyByLabelValues, Statistic}; +use std::sync::Arc; + +fn merge_exact_states( + states: impl IntoIterator>, +) -> Result { + let mut states = states.into_iter(); + let exact = |state: &Arc| { + state + .as_any() + .downcast_ref::() + .cloned() + .ok_or_else(|| "readout requires Planner exact state".to_string()) + }; + let mut merged = exact(&states.next().ok_or("empty exact state input")?)?; + for state in states { + merged + .merge_from(&exact(&state)?) + .map_err(|error| error.to_string())?; + } + Ok(merged) +} + +/// PromQL counter readouts omit a series with fewer than two samples. Other +/// state/type/range failures remain errors rather than empty results. +pub fn insufficient_counter_samples(state: &dyn AggregateCore, statistic: Statistic) -> bool { + matches!(statistic, Statistic::Rate | Statistic::Increase) + && state + .as_any() + .downcast_ref::() + .is_some_and(|state| state.insufficient_counter_samples(statistic, &None)) +} + +/// Merge already selected exact panes and read one population. `None` means +/// the population is absent from the result: a counter with too few samples, +/// or an empty MIN/MAX. +pub fn exact_readout( + states: impl IntoIterator>, + statistic: Statistic, + range_ms: Option<(i64, i64)>, + key: Option<&KeyByLabelValues>, +) -> Result, String> { + let merged = merge_exact_states(states)?; + if merged.insufficient_counter_samples(statistic, &key.cloned()) { + return Ok(None); + } + merged + .readout(statistic, range_ms, key) + .map_err(|error| error.to_string()) +} + +#[cfg(test)] +mod counter_tests { + use super::*; + use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; + + fn counter(kind: ExactKind, params: ExactParams, keyed: bool) -> ExactAccumulator { + ExactAccumulator::new(SummaryFamilyType::ExactAggregate(kind, params), keyed).unwrap() + } + + // A counter population with a single sample is absent, keyed or not. + #[test] + fn planner_counter_population_omits_insufficient_samples() { + for (kind, params, statistic) in [ + (ExactKind::Rate, ExactParams::Rate, Statistic::Rate), + ( + ExactKind::Increase, + ExactParams::Increase, + Statistic::Increase, + ), + ] { + for keyed in [false, true] { + let mut state = counter(kind.clone(), params.clone(), keyed); + let key = keyed.then(|| KeyByLabelValues::new_with_labels(vec!["checkout".into()])); + state.update(key.as_ref(), 10., 10_000); + assert_eq!( + exact_readout( + [Arc::new(state) as Arc], + statistic, + None, + key.as_ref() + ) + .unwrap(), + None + ); + } + } + } + + // Two ordered samples read a rate; an inverted range and empty input fail. + #[test] + fn sparse_counter_is_absent_but_invalid_ranges_still_fail() { + let mut state = counter(ExactKind::Rate, ExactParams::Rate, false); + state.update(None, 10., 10_000); + let rate = Statistic::Rate; + let one = [Arc::new(state.clone()) as Arc]; + assert_eq!( + exact_readout(one, rate, Some((0, 60_000)), None).unwrap(), + None + ); + state.update(None, 20., 20_000); + let two = || [Arc::new(state.clone()) as Arc]; + assert!(exact_readout(two(), rate, Some((0, 60_000)), None) + .unwrap() + .is_some()); + assert!(exact_readout(two(), rate, Some((60_000, 0)), None).is_err()); + assert!(exact_readout([], rate, Some((0, 60_000)), None).is_err()); + } +} diff --git a/crates/asap-physical-operators/src/runtime/batch_execution.rs b/crates/asap-physical-operators/src/runtime/batch_execution.rs new file mode 100644 index 00000000..f6300440 --- /dev/null +++ b/crates/asap-physical-operators/src/runtime/batch_execution.rs @@ -0,0 +1,210 @@ +//! Execute a bounded in-memory batch through native operators. This is also the +//! bridge for deployments whose boundary values are not yet streaming batches. +use crate::{ + operators::Operator, + plan::PhysicalDag, + runtime::{RunContext, SharedValue}, + values::Batch, + Error, +}; +use futures::{FutureExt, StreamExt}; + +/// Every input is already in memory; the chain contains native operators only. +/// This deliberately does not enter a nested executor when called from a DAG +/// adapter. I/O belongs to source operators in the surrounding execution. +pub fn evaluate_batch( + input: Batch, + operators: Vec, + context: RunContext, +) -> Result>, Error> { + let mut graph = PhysicalDag::default(); + graph.add( + 0, + vec![], + Operator::source(input.schema().clone(), vec![input])?, + )?; + let mut root = 0; + for operator in operators { + graph.add(root + 1, vec![root], operator)?; + root += 1; + } + evaluate_graph(graph, root, context) +} + +/// Bind the ordered in-memory inputs of a native multi-input operator. +pub fn evaluate_inputs( + inputs: Vec, + operator: Operator, + context: RunContext, +) -> Result>, Error> { + let mut graph = PhysicalDag::default(); + let root = inputs.len() as u64; + for (id, input) in inputs.into_iter().enumerate() { + graph.add( + id as u64, + vec![], + Operator::source(input.schema().clone(), vec![input])?, + )?; + } + graph.add(root, (0..root).collect(), operator)?; + evaluate_graph(graph, root, context) +} + +/// Evaluate a native in-memory source, including scalar sources, in the caller's scope. +pub fn evaluate_source( + source: Operator, + context: RunContext, +) -> Result>, Error> { + let mut graph = PhysicalDag::default(); + graph.add(0, vec![], source)?; + evaluate_graph(graph, 0, context) +} + +fn evaluate_graph( + graph: PhysicalDag<'_, Batch, crate::values::Schema>, + root: crate::plan::NodeId, + context: RunContext, +) -> Result>, Error> { + let mut output = graph.execute(&[root], context)?.remove(0); + let mut batches = Vec::new(); + loop { + match output.next().now_or_never() { + Some(Some(Ok(batch))) => batches.push(batch), + Some(Some(Err(error))) => return Err(error), + Some(None) => return Ok(batches), + // Native operators have no I/O sources here. Pending is the + // shared runtime's cooperative yield after a batch quantum. + None => continue, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + operators::Expression, + runtime::{Limits, Scope}, + values::Value, + }; + use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::DataType, + }; + use std::sync::Arc; + + // Engine adapters can run the identical native chain from an outer executor. + #[test] + fn same_native_chain_inside_query_and_ingestion_execution() { + let schema = Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "value".into(), + dtype: SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }], + time_index: None, + }); + for scope in [ + Scope::Query { + evaluation_time_ms: 20, + revision: 1, + }, + Scope::Ingestion { + window_start_ms: 10, + window_end_ms: 20, + revision: 1, + }, + ] { + let batch = Batch::try_new(schema.clone(), vec![vec![Value::Float64(7.)]]).unwrap(); + let negate = Operator::project( + schema.clone(), + vec![( + "value".into(), + Expression::Negate(Box::new(Expression::Column(0))), + )], + ) + .unwrap(); + let context = RunContext::new(scope, Limits::default()).unwrap(); + let result = futures::executor::block_on(async { + evaluate_batch(batch, vec![negate], context.clone()) + }) + .unwrap(); + assert!(matches!(result[0].rows()[0][0], Value::Float64(-7.))); + let source = Operator::scalar(Value::Float64(9.), DataType::Float64).unwrap(); + let scalar = evaluate_source(source, context).unwrap(); + assert!(matches!(scalar[0].rows()[0][0], Value::Float64(9.))); + } + } + + // Native sources may cross the runtime's cooperative batch quantum. + #[test] + fn in_memory_source_drives_cooperative_yields() { + let schema = Arc::new(SummarySchema { + fields: vec![], + time_index: None, + }); + let batch = Batch::try_new(schema.clone(), vec![vec![]]).unwrap(); + let source = Operator::source(schema, vec![batch; 65]).unwrap(); + let context = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits::default(), + ) + .unwrap(); + assert_eq!(evaluate_source(source, context).unwrap().len(), 65); + } + + // An adapter-held output must retain its parent's reservation after execution. + #[test] + fn returned_batches_keep_their_resource_reservation() { + let schema = Arc::new(SummarySchema { + fields: vec![], + time_index: None, + }); + let batch = Batch::try_new(schema.clone(), vec![vec![]]).unwrap(); + let bytes = batch.bytes(); + let source = Operator::source(schema, vec![batch]).unwrap(); + let context = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits { + max_bytes: bytes, + max_buffered_batches: 1, + }, + ) + .unwrap(); + let held = evaluate_source(source.clone(), context.clone()).unwrap(); + assert_eq!(context.retained_bytes(), bytes); + assert!(evaluate_source(source.clone(), context.clone()).is_err()); + drop(held); + assert_eq!(context.retained_bytes(), 0); + assert!(evaluate_source(source, context).is_ok()); + } + + // A cancelled surrounding execution also prevents its native computation. + #[test] + fn cancellation_is_not_bypassed_by_in_memory_execution() { + let schema = Arc::new(SummarySchema { + fields: vec![], + time_index: None, + }); + let batch = Batch::try_new(schema, vec![vec![]]).unwrap(); + let context = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits::default(), + ) + .unwrap(); + context.cancel(); + assert!(matches!( + evaluate_batch(batch, vec![], context), + Err(Error::Cancelled) + )); + } +} diff --git a/crates/asap-physical-operators/src/runtime/context.rs b/crates/asap-physical-operators/src/runtime/context.rs new file mode 100644 index 00000000..c145a457 --- /dev/null +++ b/crates/asap-physical-operators/src/runtime/context.rs @@ -0,0 +1,133 @@ +use crate::Error; +use std::{ + cell::{Cell, RefCell}, + rc::Rc, + task::Waker, +}; +/// Scope is part of an execution instance, never mutable state in a reusable plan. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Scope { + Ingestion { + window_start_ms: i64, + window_end_ms: i64, + revision: u64, + }, + Query { + evaluation_time_ms: i64, + revision: u64, + }, +} +#[derive(Clone, Debug)] +pub struct Limits { + pub max_buffered_batches: usize, + pub max_bytes: usize, +} +impl Default for Limits { + fn default() -> Self { + Self { + max_buffered_batches: 8, + max_bytes: 64 * 1024 * 1024, + } + } +} +pub(super) struct Control { + cancelled: Cell, + bytes: Cell, + peak: Cell, + pub(super) limits: Limits, + waiters: RefCell>, +} +#[derive(Clone)] +pub struct RunContext { + pub scope: Scope, + pub(super) control: Rc, +} +impl RunContext { + pub fn new(scope: Scope, limits: Limits) -> Result { + if limits.max_buffered_batches == 0 || limits.max_bytes == 0 { + return Err(Error::Invalid("execution limits must be positive".into())); + } + if matches!(&scope, Scope::Ingestion { window_start_ms, window_end_ms, .. } if window_start_ms > window_end_ms) + { + return Err(Error::Invalid("inverted ingestion window".into())); + } + Ok(Self { + scope, + control: Rc::new(Control { + cancelled: Cell::new(false), + bytes: Cell::new(0), + peak: Cell::new(0), + limits, + waiters: RefCell::new(Vec::new()), + }), + }) + } + pub fn cancel(&self) { + self.control.cancelled.set(true); + for waiter in self.control.waiters.borrow_mut().drain(..) { + waiter.wake(); + } + } + pub fn is_cancelled(&self) -> bool { + self.control.cancelled.get() + } + pub fn retained_bytes(&self) -> usize { + self.control.bytes.get() + } + pub fn peak_bytes(&self) -> usize { + self.control.peak.get() + } + pub fn reserve(&self, bytes: usize) -> Result { + let total = self + .control + .bytes + .get() + .checked_add(bytes) + .ok_or(Error::MemoryLimit)?; + if total > self.control.limits.max_bytes { + return Err(Error::MemoryLimit); + } + self.control.bytes.set(total); + self.control.peak.set(self.control.peak.get().max(total)); + Ok(Reservation { + bytes, + control: Rc::clone(&self.control), + }) + } + pub(super) fn register(&self, waker: &Waker) { + let mut waiters = self.control.waiters.borrow_mut(); + if !waiters.iter().any(|old| old.will_wake(waker)) { + waiters.push(waker.clone()); + } + } +} +pub struct Reservation { + bytes: usize, + pub(super) control: Rc, +} +impl Reservation { + /// Adjust an operator-owned allocation without accumulating bookkeeping entries. + pub fn resize(&mut self, bytes: usize) -> Result<(), Error> { + let total = self + .control + .bytes + .get() + .checked_sub(self.bytes) + .and_then(|total| total.checked_add(bytes)) + .ok_or(Error::MemoryLimit)?; + if total > self.control.limits.max_bytes { + return Err(Error::MemoryLimit); + } + self.control.bytes.set(total); + self.control.peak.set(self.control.peak.get().max(total)); + self.bytes = bytes; + Ok(()) + } +} +impl Drop for Reservation { + fn drop(&mut self) { + self.control + .bytes + .set(self.control.bytes.get().saturating_sub(self.bytes)); + } +} diff --git a/crates/asap-physical-operators/src/runtime/cooperative.rs b/crates/asap-physical-operators/src/runtime/cooperative.rs new file mode 100644 index 00000000..fae869d6 --- /dev/null +++ b/crates/asap-physical-operators/src/runtime/cooperative.rs @@ -0,0 +1,40 @@ +//! Worker-local CPU loops yield so other consumers and cancellation can progress. +use super::RunContext; +use crate::Error; +use std::task::Poll; + +pub(crate) struct Cooperative { + context: RunContext, + remaining: usize, +} +impl Cooperative { + pub(crate) fn new(context: &RunContext) -> Self { + Self { + context: context.clone(), + remaining: 1024, + } + } + pub(crate) async fn checkpoint(&mut self) -> Result<(), Error> { + if self.context.is_cancelled() { + return Err(Error::Cancelled); + } + self.remaining -= 1; + if self.remaining == 0 { + self.remaining = 1024; + let mut yielded = false; + futures::future::poll_fn(|cx| { + if self.context.is_cancelled() { + return Poll::Ready(Err(Error::Cancelled)); + } + if yielded { + return Poll::Ready(Ok(())); + } + yielded = true; + cx.waker().wake_by_ref(); + Poll::Pending + }) + .await?; + } + Ok(()) + } +} diff --git a/crates/asap-physical-operators/src/runtime/mod.rs b/crates/asap-physical-operators/src/runtime/mod.rs new file mode 100644 index 00000000..f72a57ef --- /dev/null +++ b/crates/asap-physical-operators/src/runtime/mod.rs @@ -0,0 +1,272 @@ +//! Per-run producer sharing, streams, backpressure and resource ownership. +use crate::{ + plan::{NodeId, PhysicalDag}, + Error, +}; +use futures::{stream::LocalBoxStream, Stream}; +use std::{ + cell::RefCell, + collections::{BTreeMap, VecDeque}, + fmt::Debug, + pin::Pin, + rc::Rc, + sync::Arc, + task::{Context, Poll, Waker}, +}; +mod context; +pub use context::{Limits, Reservation, RunContext, Scope}; +pub type OutputStream<'a, V> = LocalBoxStream<'a, Result>; +/// An output owns its memory reservation even after it leaves the DAG's queue. +pub struct SharedValue { + value: Arc, + _reservation: Rc, +} +impl Clone for SharedValue { + fn clone(&self) -> Self { + Self { + value: Arc::clone(&self.value), + _reservation: Rc::clone(&self._reservation), + } + } +} +impl std::ops::Deref for SharedValue { + type Target = V; + fn deref(&self) -> &V { + &self.value + } +} +impl SharedValue { + pub fn value(&self) -> &V { + &self.value + } +} + +pub(crate) fn execute<'r, V: 'r, S: Clone + PartialEq + Debug + 'r>( + dag: &'r PhysicalDag<'_, V, S>, + roots: &[NodeId], + context: RunContext, +) -> Result>, Error> { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + dag.validate(roots)?; + let mut pending = roots.to_vec(); + let mut visited = std::collections::BTreeSet::new(); + while let Some(id) = pending.pop() { + if visited.insert(id) { + let node = &dag.nodes[&id]; + node.operator.validate_context(&context)?; + pending.extend(node.inputs.iter().copied()); + } + } + fn build<'r, V: 'r, S: 'r>( + dag: &'r PhysicalDag<'_, V, S>, + id: NodeId, + context: &RunContext, + states: &mut BTreeMap>>>, + ) -> Result>>, Error> { + if let Some(state) = states.get(&id) { + return Ok(Rc::clone(state)); + } + let node = &dag.nodes[&id]; + let mut inputs = Vec::new(); + for &child in &node.inputs { + inputs.push(Input::subscribe(build(dag, child, context, states)?)); + } + let stream = node + .operator + .start(inputs, context.clone()) + .map_err(|source| Error::AtNode { + node: id, + operation: node.operator.name().into(), + source: Box::new(source), + })?; + let op = node.operator.as_ref(); + let state = Rc::new(RefCell::new(Producer { + stream: Some(stream), + node: id, + operation: node.operator.name().into(), + size: Box::new(move |value| op.output_bytes(value)), + context: context.clone(), + queue: VecDeque::new(), + base: 0, + next_reader: 0, + batches_polled: 0, + readers: BTreeMap::new(), + waiters: BTreeMap::new(), + finished: false, + failure: None, + })); + states.insert(id, Rc::clone(&state)); + Ok(state) + } + let mut states = BTreeMap::new(); + roots + .iter() + .map(|&id| build(dag, id, &context, &mut states).map(Input::subscribe)) + .collect() +} +struct Producer<'a, V> { + node: NodeId, + operation: String, + stream: Option>, + size: Box usize + 'a>, + context: RunContext, + queue: VecDeque>, + base: u64, + next_reader: u64, + batches_polled: usize, + readers: BTreeMap, + waiters: BTreeMap, + finished: bool, + failure: Option, +} +impl Producer<'_, V> { + fn trim(&mut self) { + let minimum = self + .readers + .values() + .copied() + .min() + .unwrap_or(self.base + self.queue.len() as u64); + while self.base < minimum { + self.queue.pop_front(); + self.base += 1; + } + for (_, waker) in std::mem::take(&mut self.waiters) { + waker.wake(); + } + if self.readers.is_empty() { + self.stream = None; + self.queue.clear(); + } + } +} +pub struct Input<'a, V> { + producer: Rc>>, + reader: u64, + done: bool, +} +impl<'a, V> Input<'a, V> { + fn subscribe(producer: Rc>>) -> Self { + let reader = { + let mut state = producer.borrow_mut(); + let id = state.next_reader; + state.next_reader += 1; + let base = state.base; + state.readers.insert(id, base); + id + }; + Self { + producer, + reader, + done: false, + } + } +} +impl Drop for Input<'_, V> { + fn drop(&mut self) { + let mut state = self.producer.borrow_mut(); + state.readers.remove(&self.reader); + state.waiters.remove(&self.reader); + state.trim(); + } +} +impl Stream for Input<'_, V> { + type Item = Result, Error>; + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.done { + return Poll::Ready(None); + } + let mut state = this.producer.borrow_mut(); + state.context.register(cx.waker()); + if state.context.is_cancelled() { + state.failure = Some(Error::Cancelled); + state.finished = true; + state.stream = None; + state.queue.clear(); + } + let position = state.readers[&this.reader]; + let index = (position - state.base) as usize; + if let Some(value) = state.queue.get(index).cloned() { + state.readers.insert(this.reader, position + 1); + state.trim(); + return Poll::Ready(Some(Ok(value))); + } + if state.finished { + this.done = true; + state.readers.remove(&this.reader); + let failure = state.failure.clone(); + state.trim(); + return Poll::Ready(failure.map(Err)); + } + state.waiters.insert(this.reader, cx.waker().clone()); + if state.queue.len() >= state.context.control.limits.max_buffered_batches { + return Poll::Pending; + } + // Always-ready sources must still give cancellation and other roots a turn. + if state.batches_polled >= 32 { + state.batches_polled = 0; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + let polled = state + .stream + .as_mut() + .expect("unfinished producer") + .as_mut() + .poll_next(cx); + if matches!(&polled, Poll::Ready(Some(Ok(_)))) { + state.batches_polled += 1; + } + match polled { + Poll::Pending => Poll::Pending, + Poll::Ready(Some(Ok(value))) => match state.context.reserve((state.size)(&value)) { + Ok(reservation) => { + let value = SharedValue { + value: Arc::new(value), + _reservation: Rc::new(reservation), + }; + state.queue.push_back(value.clone()); + state.readers.insert(this.reader, position + 1); + state.trim(); + Poll::Ready(Some(Ok(value))) + } + Err(error) => { + state.failure = Some(error.clone()); + state.finished = true; + state.stream = None; + this.done = true; + state.readers.remove(&this.reader); + state.trim(); + Poll::Ready(Some(Err(error))) + } + }, + Poll::Ready(result) => { + let error = result.and_then(Result::err).map(|source| match source { + Error::AtNode { .. } | Error::Cancelled | Error::MemoryLimit => source, + source => Error::AtNode { + node: state.node, + operation: state.operation.clone(), + source: Box::new(source), + }, + }); + state.failure = error.clone(); + state.finished = true; + state.stream = None; + this.done = true; + state.readers.remove(&this.reader); + state.trim(); + Poll::Ready(error.map(Err)) + } + } + } +} + +pub mod batch_execution; +#[cfg(test)] +mod tests; + +mod cooperative; +pub(crate) use cooperative::Cooperative; diff --git a/crates/asap-physical-operators/src/runtime/tests.rs b/crates/asap-physical-operators/src/runtime/tests.rs new file mode 100644 index 00000000..f683041c --- /dev/null +++ b/crates/asap-physical-operators/src/runtime/tests.rs @@ -0,0 +1,262 @@ +use super::*; +use crate::plan::PhysicalOperator; +use futures::{executor::block_on, stream, StreamExt}; +use std::cell::Cell; + +struct Source { + starts: Rc>, + polls: Rc>, + fail: bool, + end: u64, +} +impl PhysicalOperator for Source { + fn name(&self) -> &str { + "CountingSource" + } + fn input_schemas(&self) -> Vec<()> { + vec![] + } + fn output_schema(&self) {} + fn output_bytes(&self, _: &u64) -> usize { + 8 + } + fn start<'a>( + &'a self, + _: Vec>, + _: RunContext, + ) -> Result, Error> { + self.starts.set(self.starts.get() + 1); + Ok(stream::iter(0..self.end) + .map(move |n| { + self.polls.set(self.polls.get() + 1); + if self.fail && n == 1 { + Err(Error::Operator("source failure".into())) + } else { + Ok(n) + } + }) + .boxed_local()) + } +} +struct Identity; +impl PhysicalOperator for Identity { + fn name(&self) -> &str { + "Identity" + } + fn input_schemas(&self) -> Vec<()> { + vec![()] + } + fn output_schema(&self) {} + fn output_bytes(&self, _: &u64) -> usize { + 8 + } + fn start<'a>( + &'a self, + mut inputs: Vec>, + _: RunContext, + ) -> Result, Error> { + Ok(inputs + .remove(0) + .map(|value| value.map(|v| *v)) + .boxed_local()) + } +} +fn context() -> RunContext { + RunContext::new( + Scope::Query { + evaluation_time_ms: 100, + revision: 1, + }, + Limits { + max_buffered_batches: 1, + max_bytes: 1024, + }, + ) + .unwrap() +} +fn source(fail: bool) -> (Source, Rc>, Rc>) { + let starts = Rc::new(Cell::new(0)); + let polls = Rc::new(Cell::new(0)); + ( + Source { + starts: starts.clone(), + polls: polls.clone(), + fail, + end: 4, + }, + starts, + polls, + ) +} + +// A shared producer runs once, and the slow reader bounds producer progress. +#[test] +fn shared_source_backpressure_and_reader_drop() { + let (source, starts, polls) = source(false); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + let context = context(); + let mut readers = dag.execute(&[0, 0], context.clone()).unwrap(); + let mut slow = readers.pop().unwrap(); + let mut fast = readers.pop().unwrap(); + assert_eq!(starts.get(), 1); + let first = block_on(fast.next()).unwrap().unwrap(); + assert_eq!(*first, 0); + let mut cx = Context::from_waker(futures::task::noop_waker_ref()); + assert!(Pin::new(&mut fast).poll_next(&mut cx).is_pending()); + assert_eq!(polls.get(), 1); + let same = block_on(slow.next()).unwrap().unwrap(); + assert!(Arc::ptr_eq(&first.value, &same.value)); + drop(same); + drop(first); + assert_eq!(context.retained_bytes(), 0); + assert_eq!(*block_on(fast.next()).unwrap().unwrap(), 1); + drop(slow); + assert_eq!(*block_on(fast.next()).unwrap().unwrap(), 2); + assert_eq!(*block_on(fast.next()).unwrap().unwrap(), 3); + assert!(block_on(fast.next()).is_none()); + assert_eq!(polls.get(), 4); + drop(fast); + assert_eq!(context.retained_bytes(), 0); +} + +// Independent branches consume a common node concurrently without duplicate work. +#[test] +fn diamond_and_run_isolation() { + let (source, starts, polls) = source(false); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + dag.add(1, vec![0], Identity).unwrap(); + dag.add(2, vec![0], Identity).unwrap(); + for _ in 0..2 { + let mut outputs = dag.execute(&[1, 2], context()).unwrap(); + let a = outputs.pop().unwrap(); + let b = outputs.pop().unwrap(); + let (a, b) = + block_on(async { futures::join!(a.collect::>(), b.collect::>()) }); + assert_eq!( + a.iter().map(|v| **v.as_ref().unwrap()).collect::>(), + vec![0, 1, 2, 3] + ); + assert_eq!( + b.iter().map(|v| **v.as_ref().unwrap()).collect::>(), + vec![0, 1, 2, 3] + ); + } + assert_eq!(starts.get(), 2); + assert_eq!(polls.get(), 8); +} + +// Failure reaches every subscriber; cancellation stops further producer work. +#[test] +fn broadcast_error_and_cancel() { + let (source, _, polls) = source(true); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + let mut outputs = dag.execute(&[0, 0], context()).unwrap(); + let a = outputs.pop().unwrap(); + let b = outputs.pop().unwrap(); + let (a, b) = block_on(async { futures::join!(a.collect::>(), b.collect::>()) }); + for values in [a, b] { + assert_eq!(values.len(), 2); + assert!(matches!(values[1], Err(Error::AtNode { node: 0, .. }))); + } + assert_eq!(polls.get(), 2); + let run = context(); + let mut output = dag.execute(&[0], run.clone()).unwrap().remove(0); + run.cancel(); + assert!(matches!( + block_on(output.next()), + Some(Err(Error::Cancelled)) + )); + assert!(block_on(output.next()).is_none()); + assert_eq!(polls.get(), 2); +} + +// Retaining a consumer output retains its budget lease after queue eviction. +#[test] +fn retained_outputs_count_against_budget() { + let (source, _, _) = source(false); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + let run = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits { + max_buffered_batches: 1, + max_bytes: 8, + }, + ) + .unwrap(); + let mut input = dag.execute(&[0], run.clone()).unwrap().remove(0); + let held = block_on(input.next()).unwrap().unwrap(); + assert_eq!(run.retained_bytes(), 8); + assert!(matches!( + block_on(input.next()), + Some(Err(Error::MemoryLimit)) + )); + drop(input); + assert_eq!(run.retained_bytes(), 8); + drop(held); + assert_eq!(run.retained_bytes(), 0); +} + +// Invalid graphs fail before even starting a source. +#[test] +fn invalid_graphs_do_not_start_sources() { + let (source, starts, _) = source(false); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + dag.add(1, vec![2], Identity).unwrap(); + dag.add(2, vec![1], Identity).unwrap(); + assert!(dag.execute(&[0, 1], context()).is_err()); + assert_eq!(starts.get(), 0); + let mut missing = PhysicalDag::default(); + missing.add(1, vec![9], Identity).unwrap(); + assert!(missing.validate(&[1]).is_err()); + let mut arity = PhysicalDag::default(); + arity.add(1, vec![], Identity).unwrap(); + assert!(arity.validate(&[1]).is_err()); +} + +// An always-ready source must yield so cancellation can be polled on this worker. +#[test] +fn ready_sources_cooperate_with_cancellation() { + let (mut source, _, polls) = source(false); + source.end = 10_000; + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + let context = context(); + let mut input = dag.execute(&[0], context.clone()).unwrap().remove(0); + block_on(async { + let drain = async { + while let Some(result) = input.next().await { + if let Err(error) = result { + assert_eq!(error, Error::Cancelled); + return; + } + } + panic!("source completed without yielding"); + }; + let cancel = async { + context.cancel(); + }; + futures::join!(drain, cancel); + }); + assert_eq!(polls.get(), 32); + assert_eq!(context.retained_bytes(), 0); +} + +// Cached shorter paths must not hide an over-deep path through shared nodes. +#[test] +fn depth_limit_covers_shared_paths() { + let (source, _, _) = source(false); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + for id in 1..129 { + dag.add(id, vec![id - 1], Identity).unwrap(); + } + assert!(dag.validate(&(0..129).collect::>()).is_err()); +} diff --git a/crates/asap-physical-operators/src/sources/memory.rs b/crates/asap-physical-operators/src/sources/memory.rs new file mode 100644 index 00000000..7f77c54f --- /dev/null +++ b/crates/asap-physical-operators/src/sources/memory.rs @@ -0,0 +1,44 @@ +use super::*; +/// Immutable in-memory raw data. The connector owns the resident input; each +/// cursor clones only the next requested batch, not the entire data set. +pub struct MemorySource { + schema: Schema, + batches: Vec, +} +impl MemorySource { + pub fn new(schema: Schema, batches: Vec) -> Result { + crate::values::validate_schema(&schema)?; + if schema + .fields + .iter() + .any(|f| !matches!(f.dtype, SummaryFamilyType::Plain(_))) + { + return Err(Error::Invalid( + "raw source cannot contain summary states".into(), + )); + } + if batches.iter().any(|batch| batch.schema() != &schema) { + return Err(Error::Invalid("memory source batch schema mismatch".into())); + } + Ok(Self { schema, batches }) + } +} +impl RawSource for MemorySource { + fn boundedness(&self) -> crate::plan::Boundedness { + crate::plan::Boundedness::Bounded + } + fn schema(&self) -> Schema { + self.schema.clone() + } + fn scan(&self, context: RunContext) -> Result, Error> { + Ok(stream::iter(self.batches.iter()) + .map(move |batch| { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + let _allocation = context.reserve(batch.bytes())?; + Ok(batch.clone()) + }) + .boxed_local()) + } +} diff --git a/crates/asap-physical-operators/src/sources/mod.rs b/crates/asap-physical-operators/src/sources/mod.rs new file mode 100644 index 00000000..4da778b0 --- /dev/null +++ b/crates/asap-physical-operators/src/sources/mod.rs @@ -0,0 +1,187 @@ +//! Raw data access. Connectors provide rows; Scan owns Planner predicate semantics. +use crate::{ + expressions::CompiledExpression, + plan::PhysicalOperator, + runtime::{Input, OutputStream, RunContext}, + values::{Batch, Schema, Value}, + Error, +}; +use futures::{stream, StreamExt}; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::{DataType, QueryExpr, Source}, +}; +use std::sync::Arc; + +/// A bound data source. Metadata must be stable for the lifetime of the binding. +/// Each scan opens an independent cursor. Connectors return raw, unfiltered rows +/// and must honor cancellation and bound their own I/O buffers. Dropping a cursor +/// must release its resources. A connector error is never an empty successful scan. +pub trait RawSource { + fn schema(&self) -> Schema; + /// Declare a finite snapshot/window explicitly; execution scope alone does not bound a cursor. + fn boundedness(&self) -> crate::plan::Boundedness { + crate::plan::Boundedness::Unknown + } + fn scan(&self, context: RunContext) -> Result, Error>; +} + +/// Explicit source identities; no implicit network discovery or fallback. +#[derive(Default)] +pub struct DataSources { + sources: Vec<(Source, Arc)>, +} +impl DataSources { + pub fn register(&mut self, identity: Source, source: Arc) -> Result<(), Error> { + if self.sources.iter().any(|(key, _)| key == &identity) { + return Err(Error::Invalid("duplicate data source".into())); + } + crate::values::validate_schema(&source.schema())?; + self.sources.push((identity, source)); + Ok(()) + } + pub fn bind(&self, expression: &QueryExpr) -> Result { + let QueryExpr::Scan { + source, + predicates, + schema, + } = expression + else { + return Err(Error::Invalid( + "raw Scan requires a Planner Scan leaf".into(), + )); + }; + let output = Arc::new(SummarySchema { + fields: schema + .columns + .iter() + .map(|column| SummaryField { + name: column.name.clone(), + dtype: SummaryFamilyType::Plain(column.dtype.clone()), + nullable: column.nullable, + }) + .collect(), + time_index: schema.time_index, + }); + crate::values::validate_schema(&output)?; + let reader = self + .sources + .iter() + .find(|(key, _)| key == source) + .map(|(_, reader)| reader.clone()) + .ok_or_else(|| Error::Invalid(format!("unbound raw source: {source:?}")))?; + if reader.schema() != output { + return Err(Error::Invalid( + "raw source differs from Planner Scan schema".into(), + )); + } + let predicates = predicates + .iter() + .map(|predicate| { + let predicate = CompiledExpression::compile(&predicate.0, &output)?; + if predicate.dtype().0 != DataType::Bool { + return Err(Error::Invalid("Scan predicate must be boolean".into())); + } + Ok(predicate) + }) + .collect::, Error>>()?; + Ok(Scan { + reader, + output, + predicates, + }) + } +} + +pub struct Scan { + reader: Arc, + output: Schema, + predicates: Vec, +} +impl PhysicalOperator for Scan { + fn properties(&self, _: &[crate::plan::PlanProperties]) -> crate::plan::PlanProperties { + crate::plan::PlanProperties { + boundedness: self.reader.boundedness(), + emission: crate::plan::Emission::Incremental, + } + } + + fn name(&self) -> &str { + "Scan" + } + fn input_schemas(&self) -> Vec { + vec![] + } + fn output_schema(&self) -> Schema { + self.output.clone() + } + fn output_bytes(&self, batch: &Batch) -> usize { + batch.bytes() + } + fn start<'a>( + &'a self, + inputs: Vec>, + context: RunContext, + ) -> Result, Error> { + if !inputs.is_empty() { + return Err(Error::Invalid("Scan cannot have inputs".into())); + } + if context.is_cancelled() { + return Err(Error::Cancelled); + } + // Opening is lazy: validation and construction of a run perform no I/O. + let opening = context.clone(); + let stream = stream::once(async move { + if opening.is_cancelled() { + return Err(Error::Cancelled); + } + self.reader.scan(opening) + }); + use futures::TryStreamExt; + Ok(stream + .try_flatten() + .map(move |batch| { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + let batch = batch?; + if batch.schema() != &self.output { + return Err(Error::Invalid( + "connector returned a different Scan schema".into(), + )); + } + if self.predicates.is_empty() { + return Ok(batch); + } + let _workspace = + context.reserve(batch.bytes().checked_mul(2).ok_or(Error::MemoryLimit)?)?; + let mut rows = Vec::new(); + for row in batch.rows() { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + let mut keep = true; + for predicate in &self.predicates { + match predicate.evaluate(row)? { + Value::Bool(true) => {} + Value::Bool(false) | Value::Null => { + keep = false; + break; + } + _ => { + return Err(Error::Invalid("Scan predicate is not boolean".into())) + } + } + } + if keep { + rows.push(row.clone()); + } + } + Batch::try_new(self.output.clone(), rows) + }) + .boxed_local()) + } +} + +mod memory; +pub use memory::MemorySource; diff --git a/crates/asap-physical-operators/src/summary_kernels/exact.rs b/crates/asap-physical-operators/src/summary_kernels/exact.rs index 3343e875..cb9d6099 100644 --- a/crates/asap-physical-operators/src/summary_kernels/exact.rs +++ b/crates/asap-physical-operators/src/summary_kernels/exact.rs @@ -115,6 +115,28 @@ impl ExactAccumulator { pub fn family(&self) -> &SummaryFamilyType { &self.family } + pub(crate) fn insufficient_counter_samples( + &self, + statistic: Statistic, + key: &Option, + ) -> bool { + if statistic != self.statistic() { + return false; + } + let state = match (&self.keyed, key) { + (Some(states), Some(key)) => states.get(key), + (None, None) => Some(&self.scalar), + _ => None, + }; + match state { + Some(ScalarState::Counter(None)) => true, + Some(ScalarState::Counter(Some(counter))) => { + counter.sample_count < 2 + || counter.last_seen_timestamp == counter.starting_timestamp + } + _ => false, + } + } pub fn is_keyed(&self) -> bool { self.keyed.is_some() } diff --git a/crates/asap-physical-operators/src/values.rs b/crates/asap-physical-operators/src/values.rs index 6e761c11..56d43633 100644 --- a/crates/asap-physical-operators/src/values.rs +++ b/crates/asap-physical-operators/src/values.rs @@ -2,7 +2,7 @@ use crate::AggregateCore; use crate::Error; use planner_types::{ - post_asap::{SummaryFamilyType, SummarySchema}, + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, pre_asap::DataType, }; use std::{cmp::Ordering, sync::Arc}; @@ -247,6 +247,16 @@ impl Batch { .sum::() } } +pub(crate) fn group_key(row: &[Value], columns: &[usize]) -> Result>, Error> { + columns + .iter() + .map(|&i| { + row.get(i) + .ok_or_else(|| Error::Invalid("group column out of range".into()))? + .key() + }) + .collect() +} pub(crate) use crate::capability::validate_native_family as validate_family; @@ -328,6 +338,20 @@ pub(crate) fn validate_schema(schema: &Schema) -> Result<(), Error> { Ok(()) } +pub(crate) fn field(schema: &Schema, column: usize) -> Result<&SummaryField, Error> { + schema + .fields + .get(column) + .ok_or_else(|| Error::Invalid("column out of range".into())) +} +pub(crate) fn plain(schema: &Schema, column: usize) -> Result<(&DataType, bool), Error> { + let f = field(schema, column)?; + let SummaryFamilyType::Plain(dtype) = &f.dtype else { + return Err(Error::Invalid("plain value required".into())); + }; + Ok((dtype, f.nullable)) +} + #[cfg(test)] mod weighted_state_tests { use super::*; diff --git a/crates/asap-physical-operators/tests/blocking_resources.rs b/crates/asap-physical-operators/tests/blocking_resources.rs new file mode 100644 index 00000000..7a312892 --- /dev/null +++ b/crates/asap-physical-operators/tests/blocking_resources.rs @@ -0,0 +1,242 @@ +//! Blocking operators enforce resources before returning their first batch. +use asap_physical_operators::{ + operators::Operator, + plan::{PhysicalDag, PhysicalOperator}, + runtime::{Limits, RunContext, Scope}, + values::{Batch, Schema, Value}, + Error, +}; +use futures::{executor::block_on, FutureExt, StreamExt}; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::{DataType, JoinKind, Predicate, QueryExpr, ScalarValue}, +}; +use std::sync::Arc; + +fn schema(width: usize) -> Schema { + Arc::new(SummarySchema { + fields: (0..width) + .map(|i| SummaryField { + name: format!("v{i}"), + dtype: SummaryFamilyType::Plain(DataType::Int64), + nullable: false, + }) + .collect(), + time_index: None, + }) +} +fn context(max_bytes: usize) -> RunContext { + RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits { + max_bytes, + ..Limits::default() + }, + ) + .unwrap() +} +fn source(n: usize) -> PhysicalDag<'static, Batch, Schema> { + let mut dag = PhysicalDag::default(); + dag.add( + 0, + vec![], + Operator::source( + schema(1), + vec![Batch::try_new(schema(1), vec![vec![Value::Int64(1)]; n]).unwrap()], + ) + .unwrap(), + ) + .unwrap(); + dag +} +fn cross_join() -> Operator { + Operator::relational_join( + schema(1), + schema(1), + JoinKind::Cross, + &Predicate(std::rc::Rc::new(QueryExpr::Literal(ScalarValue::Boolean( + true, + )))), + schema(2), + ) + .unwrap() +} + +// Even callers starting an operator directly cannot bypass its workspace budget. +#[test] +fn join_reserves_result_growth_before_returning_output() { + let sources = source(64); + let run = context(32 * 1024); + let inputs = sources.execute(&[0, 0], run.clone()).unwrap(); + let join = cross_join(); + let mut output = join.start(inputs, run.clone()).unwrap(); + assert!(matches!( + block_on(output.next()), + Some(Err(Error::MemoryLimit)) + )); + drop(output); + assert_eq!(run.retained_bytes(), 0); +} + +// A single large input batch must not monopolize the worker during a join. +#[test] +fn join_yields_during_computation_and_observes_cancellation() { + let sources = source(64); + let run = context(16 * 1024 * 1024); + let inputs = sources.execute(&[0, 0], run.clone()).unwrap(); + let join = cross_join(); + let mut output = join.start(inputs, run.clone()).unwrap(); + assert!( + output.next().now_or_never().is_none(), + "join should yield before producing all 4096 rows" + ); + run.cancel(); + assert!(matches!( + block_on(output.next()), + Some(Err(Error::Cancelled)) + )); + drop(output); + assert_eq!(run.retained_bytes(), 0); +} + +// Sorting and grouping yield even for one large batch. +#[test] +fn blocking_reductions_yield_and_release_memory_on_cancellation() { + use asap_physical_operators::{ + operators::{Reduction, SortKey}, + plan::PhysicalOperator, + }; + let operators = vec![ + Operator::sort( + schema(1), + vec![SortKey { + column: 0, + descending: false, + nulls_first: false, + }], + vec![], + ) + .unwrap(), + Operator::aggregate(schema(1), vec![], vec![("sum".into(), Reduction::Sum(0))]).unwrap(), + ]; + for operator in operators { + let sources = source(768); + let run = context(16 * 1024 * 1024); + let inputs = sources.execute(&[0], run.clone()).unwrap(); + let mut output = operator.start(inputs, run.clone()).unwrap(); + assert!(output.next().now_or_never().is_none()); + run.cancel(); + assert!(matches!( + block_on(output.next()), + Some(Err(Error::Cancelled)) + )); + drop(output); + assert_eq!(run.retained_bytes(), 0); + } +} + +// Merge-sort rounds preserve input order for tied keys across chunk boundaries. +#[test] +fn cooperative_sort_preserves_ties_across_chunks() { + use asap_physical_operators::operators::SortKey; + let batch = Batch::try_new( + schema(2), + (0..1025) + .rev() + .map(|i| vec![Value::Int64(i % 3), Value::Int64(i)]) + .collect(), + ) + .unwrap(); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], Operator::source(schema(2), vec![batch]).unwrap()) + .unwrap(); + dag.add( + 1, + vec![0], + Operator::sort( + schema(2), + vec![SortKey { + column: 0, + descending: false, + nulls_first: false, + }], + vec![], + ) + .unwrap(), + ) + .unwrap(); + let mut output = dag + .execute(&[1], context(16 * 1024 * 1024)) + .unwrap() + .remove(0); + let batch = block_on(output.next()).unwrap().unwrap(); + let expected = (0..3) + .flat_map(|key| (0..1025).rev().filter(move |i| i % 3 == key)) + .collect::>(); + for (row, expected) in batch.rows().iter().zip(expected) { + assert!(matches!(row[1], Value::Int64(i) if i == expected)); + } + assert_eq!(batch.rows().len(), 1025); +} + +// The integrated weighted-summary path obeys the same cooperative cancellation contract. +#[test] +fn weighted_summary_build_yields_within_a_batch() { + use planner_types::post_asap::{SketchAlgorithm, SketchKind, SketchParams}; + let input = Arc::new(SummarySchema { + fields: vec![ + SummaryField { + name: "item".into(), + dtype: SummaryFamilyType::Plain(DataType::Int64), + nullable: false, + }, + SummaryField { + name: "weight".into(), + dtype: SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }, + ], + time_index: None, + }); + let mut sources = PhysicalDag::default(); + let batch = Batch::try_new( + input.clone(), + (0..1500) + .map(|i| vec![Value::Int64(i % 8), Value::Float64(0.25)]) + .collect(), + ) + .unwrap(); + sources + .add( + 0, + vec![], + Operator::source(input.clone(), vec![batch]).unwrap(), + ) + .unwrap(); + let family = SummaryFamilyType::Sketch( + SketchKind::new( + SketchAlgorithm::CmsWithHeap, + SketchParams::CmsWithHeap { + width: 64, + depth: 3, + heap_size: 8, + }, + ), + Default::default(), + ); + let operator = Operator::keyed_summary_build(input, family, 1, vec![0], vec![]).unwrap(); + let run = context(16 * 1024 * 1024); + let inputs = sources.execute(&[0], run.clone()).unwrap(); + let mut output = operator.start(inputs, run.clone()).unwrap(); + assert!(output.next().now_or_never().is_none()); + run.cancel(); + assert!(matches!( + block_on(output.next()), + Some(Err(Error::Cancelled)) + )); + drop(output); + assert_eq!(run.retained_bytes(), 0); +} diff --git a/crates/asap-physical-operators/tests/plan_properties.rs b/crates/asap-physical-operators/tests/plan_properties.rs new file mode 100644 index 00000000..f9a509c2 --- /dev/null +++ b/crates/asap-physical-operators/tests/plan_properties.rs @@ -0,0 +1,151 @@ +//! Finite-input contracts are validated before source execution. +use asap_physical_operators::{ + operators::{Operator, SortKey}, + plan::{Boundedness, Emission, PhysicalDag}, + runtime::{Limits, OutputStream, RunContext, Scope}, + sources::{DataSources, RawSource}, + values::{Batch, Schema}, + Error, +}; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::{Column, DataType, QueryExpr, Schema as LogicalSchema, Source}, +}; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, +}; +struct DeclaredSource { + schema: Schema, + boundedness: Boundedness, + opens: Arc, +} +impl RawSource for DeclaredSource { + fn schema(&self) -> Schema { + self.schema.clone() + } + fn boundedness(&self) -> Boundedness { + self.boundedness + } + fn scan(&self, _: RunContext) -> Result, Error> { + self.opens.fetch_add(1, Ordering::SeqCst); + Ok(Box::pin(futures::stream::empty())) + } +} +// A blocking parent must reject unknown and unbounded Scan inputs without opening a reader. +#[test] +fn blocking_inputs_require_an_explicit_finite_source() { + let schema = Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "v".into(), + dtype: SummaryFamilyType::Plain(DataType::Int64), + nullable: false, + }], + time_index: None, + }); + for boundedness in [ + Boundedness::Unknown, + Boundedness::Unbounded, + Boundedness::Bounded, + ] { + let opens = Arc::new(AtomicUsize::new(0)); + let mut registry = DataSources::default(); + let identity = Source::Table { + table_ref: "t".into(), + }; + registry + .register( + identity.clone(), + Arc::new(DeclaredSource { + schema: schema.clone(), + boundedness, + opens: opens.clone(), + }), + ) + .unwrap(); + let scan = registry + .bind(&QueryExpr::Scan { + source: identity, + schema: LogicalSchema::new(vec![Column::new("v", DataType::Int64, false)]), + predicates: vec![], + }) + .unwrap(); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], scan).unwrap(); + dag.add( + 1, + vec![0], + Operator::sort( + schema.clone(), + vec![SortKey { + column: 0, + descending: false, + nulls_first: false, + }], + vec![], + ) + .unwrap(), + ) + .unwrap(); + let run = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits::default(), + ) + .unwrap(); + if boundedness == Boundedness::Bounded { + let properties = dag.properties(&[1]).unwrap(); + assert_eq!(properties[&1].emission, Emission::AfterInput); + assert_eq!(properties[&1].boundedness, Boundedness::Bounded); + assert!(dag.execute(&[1], run).is_ok()); + } else { + assert!( + matches!(dag.execute(&[1], run), Err(Error::Invalid(message)) if message.contains("requires bounded inputs")) + ); + } + assert_eq!(opens.load(Ordering::SeqCst), 0); + } +} + +// Kernel support must not be mistaken for executable native state/readout support. +#[test] +fn summary_capability_levels_are_distinct() { + use asap_physical_operators::{ + capability::{validate_native_family, validate_sketch_readout, validate_summary_kernel}, + planner::post_asap::SketchQuery, + }; + use planner_types::{ + post_asap::{GroupingStrategy, SketchAlgorithm, SketchKind, SketchParams, SummaryUpdate}, + pre_asap::ColumnRef, + }; + let grouping = GroupingStrategy::default(); + let cms = SummaryFamilyType::Sketch( + SketchKind::new( + SketchAlgorithm::Cms, + SketchParams::Cms { + width: 64, + depth: 4, + }, + ), + grouping.clone(), + ); + let update = SummaryUpdate { + item: Some(planner_types::post_asap::SummaryInputExpr::Column( + ColumnRef::Named("host".into()), + )), + weight: planner_types::post_asap::SummaryInputExpr::Constant(1.0), + weight_domain: Default::default(), + }; + assert!(validate_summary_kernel(&cms, &update, &grouping).is_ok()); + assert!(validate_native_family(&cms).is_err()); + let kll = SummaryFamilyType::Sketch( + SketchKind::new(SketchAlgorithm::Kll, SketchParams::Kll { k: 128 }), + grouping, + ); + assert!(validate_native_family(&kll).is_ok()); + assert!(validate_sketch_readout(&kll, &SketchQuery::Quantile { q: 1.5 }).is_err()); + assert!(validate_sketch_readout(&kll, &SketchQuery::Cardinality).is_err()); + assert!(validate_sketch_readout(&kll, &SketchQuery::Quantile { q: 0.5 }).is_ok()); +}