use std::cmp::Ordering;
use std::ops::Range;
use vibesql_ast::{
Expression, FrameBound, FrameExclude, FrameUnit, NullsOrder, OrderByItem, OrderDirection,
UnaryOperator, WindowFrame,
};
use vibesql_types::SqlValue;
use vibesql_storage::Row;
use super::partitioning::Partition;
use super::sorting::{compare_values, compare_values_with_collation};
pub fn validate_frame(frame_spec: &Option<WindowFrame>) -> Result<(), String> {
let frame = match frame_spec {
Some(f) => f,
None => return Ok(()),
};
validate_frame_bound(&frame.start, true)?;
if let Some(ref end_bound) = frame.end {
validate_frame_bound(end_bound, false)?;
}
Ok(())
}
pub fn validate_range_order_by(
frame_spec: &Option<WindowFrame>,
order_by: &Option<Vec<OrderByItem>>,
) -> Result<(), String> {
let frame = match frame_spec {
Some(f) => f,
None => return Ok(()),
};
if !matches!(frame.unit, FrameUnit::Range) {
return Ok(());
}
let has_offset_bound =
matches!(frame.start, FrameBound::Preceding(_) | FrameBound::Following(_))
|| frame.end.as_ref().is_some_and(|end| {
matches!(end, FrameBound::Preceding(_) | FrameBound::Following(_))
});
if !has_offset_bound {
return Ok(());
}
let order_by_count = order_by.as_ref().map_or(0, |items| items.len());
if order_by_count != 1 {
return Err(
"RANGE with offset PRECEDING/FOLLOWING requires one ORDER BY expression".to_string()
);
}
Ok(())
}
fn validate_frame_bound(bound: &FrameBound, is_start: bool) -> Result<(), String> {
match bound {
FrameBound::UnboundedPreceding
| FrameBound::UnboundedFollowing
| FrameBound::CurrentRow => Ok(()),
FrameBound::Preceding(offset_expr) | FrameBound::Following(offset_expr) => {
let bound_type = if is_start { "starting" } else { "ending" };
match evaluate_offset_expr(offset_expr) {
Err(_) => {
Err(format!("frame {} offset must be a non-negative integer", bound_type))
}
Ok(Some(offset)) if offset < 0.0 => {
Err(format!("frame {} offset must be a non-negative number", bound_type))
}
Ok(None) => {
Err(format!("frame {} offset must be a non-negative number", bound_type))
}
Ok(Some(_)) => Ok(()), }
}
}
}
fn evaluate_offset_expr(expr: &Expression) -> Result<Option<f64>, String> {
match expr {
Expression::Literal(SqlValue::Integer(n)) => Ok(Some(*n as f64)),
Expression::Literal(SqlValue::Smallint(n)) => Ok(Some(*n as f64)),
Expression::Literal(SqlValue::Bigint(n)) => Ok(Some(*n as f64)),
Expression::Literal(SqlValue::Unsigned(n)) => Ok(Some(*n as f64)),
Expression::Literal(SqlValue::Float(f)) => Ok(Some(*f as f64)),
Expression::Literal(SqlValue::Real(f)) => Ok(Some(*f)),
Expression::Literal(SqlValue::Double(f)) => Ok(Some(*f)),
Expression::Literal(SqlValue::Numeric(f)) => Ok(Some(*f)),
Expression::Literal(SqlValue::Null) => {
Ok(None)
}
Expression::Literal(SqlValue::Character(s)) | Expression::Literal(SqlValue::Varchar(s)) => {
let s_str = s.as_str();
if s_str.is_empty() {
Ok(None)
} else if let Ok(n) = s_str.parse::<f64>() {
Ok(Some(n))
} else {
Ok(None)
}
}
Expression::Literal(SqlValue::Blob(_)) => {
Ok(None)
}
Expression::UnaryOp { op: UnaryOperator::Minus, expr } => {
match evaluate_offset_expr(expr)? {
Some(inner) => Ok(Some(-inner)),
None => Ok(None),
}
}
Expression::UnaryOp { op: UnaryOperator::Plus, expr } => evaluate_offset_expr(expr),
_ => Err("frame offset must be a constant expression".to_string()),
}
}
pub fn calculate_frame<F>(
partition: &Partition,
current_row_idx: usize,
order_by: &Option<Vec<OrderByItem>>,
frame_spec: &Option<WindowFrame>,
eval_fn: &F,
) -> Range<usize>
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
let partition_size = partition.len();
let frame = match frame_spec {
Some(f) => f,
None => {
let has_order_by = order_by.as_ref().is_some_and(|items| !items.is_empty());
if has_order_by {
return calculate_range_frame(
partition,
current_row_idx,
order_by,
&FrameBound::UnboundedPreceding,
&Some(FrameBound::CurrentRow),
eval_fn,
);
} else {
return 0..partition_size;
}
}
};
match frame.unit {
FrameUnit::Rows => {
calculate_rows_frame(partition, current_row_idx, &frame.start, &frame.end)
}
FrameUnit::Range => calculate_range_frame(
partition,
current_row_idx,
order_by,
&frame.start,
&frame.end,
eval_fn,
),
FrameUnit::Groups => calculate_groups_frame(
partition,
current_row_idx,
order_by,
&frame.start,
&frame.end,
eval_fn,
),
}
}
fn calculate_rows_frame(
partition: &Partition,
current_row_idx: usize,
start: &FrameBound,
end: &Option<FrameBound>,
) -> Range<usize> {
let partition_size = partition.len();
let start_idx = calculate_rows_boundary(start, current_row_idx, partition_size, true);
let end_idx = match end {
Some(end_bound) => {
calculate_rows_boundary(end_bound, current_row_idx, partition_size, false)
}
None => current_row_idx + 1, };
let start = start_idx.min(partition_size);
let end = end_idx.min(partition_size).max(start);
start..end
}
fn calculate_range_frame<F>(
partition: &Partition,
current_row_idx: usize,
order_by: &Option<Vec<OrderByItem>>,
start: &FrameBound,
end: &Option<FrameBound>,
eval_fn: &F,
) -> Range<usize>
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
let partition_size = partition.len();
if partition_size == 0 || current_row_idx >= partition_size {
return 0..0;
}
let order_items = match order_by {
Some(items) if !items.is_empty() => items,
_ => return 0..partition_size,
};
let current_row = &partition.rows[current_row_idx];
let current_value = eval_fn(&order_items[0].expr, current_row).unwrap_or(SqlValue::Null);
let start_idx = calculate_range_boundary(
partition,
current_row_idx,
order_items,
¤t_value,
start,
true,
eval_fn,
);
let end_idx = match end {
Some(end_bound) => calculate_range_boundary(
partition,
current_row_idx,
order_items,
¤t_value,
end_bound,
false,
eval_fn,
),
None => {
find_last_peer(partition, current_row_idx, order_items, eval_fn) + 1
}
};
let start = start_idx.min(partition_size);
let end = end_idx.min(partition_size).max(start);
start..end
}
fn calculate_range_boundary<F>(
partition: &Partition,
current_row_idx: usize,
order_items: &[OrderByItem],
current_value: &SqlValue,
bound: &FrameBound,
is_start: bool,
eval_fn: &F,
) -> usize
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
use vibesql_ast::OrderDirection;
let partition_size = partition.len();
let is_desc =
order_items.first().is_some_and(|item| matches!(item.direction, OrderDirection::Desc));
match bound {
FrameBound::UnboundedPreceding => 0,
FrameBound::UnboundedFollowing => partition_size,
FrameBound::CurrentRow => {
if is_start {
find_first_peer(partition, current_row_idx, order_items, eval_fn)
} else {
find_last_peer(partition, current_row_idx, order_items, eval_fn) + 1
}
}
FrameBound::Preceding(offset_expr) => {
let offset = get_numeric_offset(offset_expr);
let target_value = if is_desc {
add_to_value(current_value, offset)
} else {
subtract_from_value(current_value, offset)
};
if is_start {
find_first_row_at_or_after(partition, order_items, &target_value, eval_fn)
} else {
find_last_row_at_or_before(partition, order_items, &target_value, eval_fn)
.map_or(0, |i| i + 1)
}
}
FrameBound::Following(offset_expr) => {
let offset = get_numeric_offset(offset_expr);
let target_value = if is_desc {
subtract_from_value(current_value, offset)
} else {
add_to_value(current_value, offset)
};
if is_start {
find_first_row_at_or_after(partition, order_items, &target_value, eval_fn)
} else {
find_last_row_at_or_before(partition, order_items, &target_value, eval_fn)
.map_or(0, |i| i + 1)
}
}
}
}
fn nulls_sort_first(item: &OrderByItem) -> bool {
match item.nulls_order {
Some(NullsOrder::First) => true,
Some(NullsOrder::Last) => false,
None => matches!(item.direction, OrderDirection::Asc),
}
}
fn explicit_collation(expr: &Expression) -> Option<&str> {
match expr {
Expression::Collate { collation, .. } => Some(collation),
_ => None,
}
}
fn compare_keys_window_order(
a: &SqlValue,
b: &SqlValue,
item: &OrderByItem,
collation: Option<&str>,
) -> Ordering {
let a_null = matches!(a, SqlValue::Null);
let b_null = matches!(b, SqlValue::Null);
match (a_null, b_null) {
(true, true) => Ordering::Equal,
(true, false) => {
if nulls_sort_first(item) {
Ordering::Less
} else {
Ordering::Greater
}
}
(false, true) => {
if nulls_sort_first(item) {
Ordering::Greater
} else {
Ordering::Less
}
}
(false, false) => {
let cmp = compare_values_with_collation(a, b, collation);
match item.direction {
OrderDirection::Asc => cmp,
OrderDirection::Desc => cmp.reverse(),
}
}
}
}
fn calculate_groups_frame<F>(
partition: &Partition,
current_row_idx: usize,
order_by: &Option<Vec<OrderByItem>>,
start: &FrameBound,
end: &Option<FrameBound>,
eval_fn: &F,
) -> Range<usize>
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
let partition_size = partition.len();
if partition_size == 0 || current_row_idx >= partition_size {
return 0..0;
}
let order_items = match order_by {
Some(items) if !items.is_empty() => items,
_ => return 0..partition_size,
};
let group_boundaries = build_group_boundaries(partition, order_items, eval_fn);
let current_group = find_group_for_row(&group_boundaries, current_row_idx);
let start_idx =
calculate_groups_boundary(&group_boundaries, current_group, start, true, partition_size);
let end_idx = match end {
Some(end_bound) => calculate_groups_boundary(
&group_boundaries,
current_group,
end_bound,
false,
partition_size,
),
None => {
if current_group + 1 < group_boundaries.len() {
group_boundaries[current_group + 1]
} else {
partition_size
}
}
};
let start = start_idx.min(partition_size);
let end = end_idx.min(partition_size).max(start);
start..end
}
fn calculate_groups_boundary(
group_boundaries: &[usize],
current_group: usize,
bound: &FrameBound,
is_start: bool,
partition_size: usize,
) -> usize {
let num_groups = group_boundaries.len();
match bound {
FrameBound::UnboundedPreceding => 0,
FrameBound::UnboundedFollowing => partition_size,
FrameBound::CurrentRow => {
if is_start {
group_boundaries[current_group]
} else {
if current_group + 1 < num_groups {
group_boundaries[current_group + 1]
} else {
partition_size
}
}
}
FrameBound::Preceding(offset_expr) => {
let offset = get_numeric_offset(offset_expr) as usize;
if offset > current_group {
if is_start {
0
} else {
0
}
} else {
let target_group = current_group - offset;
if is_start {
group_boundaries[target_group]
} else {
if target_group + 1 < num_groups {
group_boundaries[target_group + 1]
} else {
partition_size
}
}
}
}
FrameBound::Following(offset_expr) => {
let offset = get_numeric_offset(offset_expr) as usize;
let target_group = current_group + offset;
if target_group >= num_groups {
if is_start {
partition_size
} else {
partition_size
}
} else {
if is_start {
group_boundaries[target_group]
} else {
if target_group + 1 < num_groups {
group_boundaries[target_group + 1]
} else {
partition_size
}
}
}
}
}
}
fn build_group_boundaries<F>(
partition: &Partition,
order_items: &[OrderByItem],
eval_fn: &F,
) -> Vec<usize>
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
let mut boundaries = vec![0];
for i in 1..partition.len() {
let prev_row = &partition.rows[i - 1];
let curr_row = &partition.rows[i];
let mut is_new_group = false;
for item in order_items {
let prev_val = eval_fn(&item.expr, prev_row).unwrap_or(SqlValue::Null);
let curr_val = eval_fn(&item.expr, curr_row).unwrap_or(SqlValue::Null);
if compare_values(&prev_val, &curr_val) != Ordering::Equal {
is_new_group = true;
break;
}
}
if is_new_group {
boundaries.push(i);
}
}
boundaries
}
fn find_group_for_row(group_boundaries: &[usize], row_idx: usize) -> usize {
for (i, &boundary) in group_boundaries.iter().enumerate().rev() {
if row_idx >= boundary {
return i;
}
}
0
}
fn find_first_peer<F>(
partition: &Partition,
current_row_idx: usize,
order_items: &[OrderByItem],
eval_fn: &F,
) -> usize
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
let current_row = &partition.rows[current_row_idx];
for i in (0..current_row_idx).rev() {
let row = &partition.rows[i];
let mut is_peer = true;
for item in order_items {
let curr_val = eval_fn(&item.expr, current_row).unwrap_or(SqlValue::Null);
let row_val = eval_fn(&item.expr, row).unwrap_or(SqlValue::Null);
if compare_values(&curr_val, &row_val) != Ordering::Equal {
is_peer = false;
break;
}
}
if !is_peer {
return i + 1;
}
}
0
}
fn find_last_peer<F>(
partition: &Partition,
current_row_idx: usize,
order_items: &[OrderByItem],
eval_fn: &F,
) -> usize
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
let current_row = &partition.rows[current_row_idx];
for i in (current_row_idx + 1)..partition.len() {
let row = &partition.rows[i];
let mut is_peer = true;
for item in order_items {
let curr_val = eval_fn(&item.expr, current_row).unwrap_or(SqlValue::Null);
let row_val = eval_fn(&item.expr, row).unwrap_or(SqlValue::Null);
if compare_values(&curr_val, &row_val) != Ordering::Equal {
is_peer = false;
break;
}
}
if !is_peer {
return i - 1;
}
}
partition.len() - 1
}
fn find_first_row_at_or_after<F>(
partition: &Partition,
order_items: &[OrderByItem],
target: &SqlValue,
eval_fn: &F,
) -> usize
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
let item = &order_items[0];
let collation = explicit_collation(&item.expr);
for (i, row) in partition.rows.iter().enumerate() {
let val = eval_fn(&item.expr, row).unwrap_or(SqlValue::Null);
if compare_keys_window_order(&val, target, item, collation) != Ordering::Less {
return i;
}
}
partition.len()
}
fn find_last_row_at_or_before<F>(
partition: &Partition,
order_items: &[OrderByItem],
target: &SqlValue,
eval_fn: &F,
) -> Option<usize>
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
let item = &order_items[0];
let collation = explicit_collation(&item.expr);
for i in (0..partition.len()).rev() {
let val = eval_fn(&item.expr, &partition.rows[i]).unwrap_or(SqlValue::Null);
if compare_keys_window_order(&val, target, item, collation) != Ordering::Greater {
return Some(i);
}
}
None
}
fn get_numeric_offset(expr: &Expression) -> f64 {
match evaluate_offset_expr(expr) {
Ok(Some(n)) => n,
_ => 0.0,
}
}
fn subtract_from_value(value: &SqlValue, offset: f64) -> SqlValue {
match value {
SqlValue::Integer(n) => SqlValue::Real(*n as f64 - offset),
SqlValue::Smallint(n) => SqlValue::Real(*n as f64 - offset),
SqlValue::Bigint(n) => SqlValue::Real(*n as f64 - offset),
SqlValue::Unsigned(n) => SqlValue::Real(*n as f64 - offset),
SqlValue::Float(f) => SqlValue::Real(*f as f64 - offset),
SqlValue::Real(f) => SqlValue::Real(*f - offset),
SqlValue::Double(f) => SqlValue::Real(*f - offset),
SqlValue::Numeric(f) => SqlValue::Real(*f - offset),
_ => value.clone(), }
}
fn add_to_value(value: &SqlValue, offset: f64) -> SqlValue {
match value {
SqlValue::Integer(n) => SqlValue::Real(*n as f64 + offset),
SqlValue::Smallint(n) => SqlValue::Real(*n as f64 + offset),
SqlValue::Bigint(n) => SqlValue::Real(*n as f64 + offset),
SqlValue::Unsigned(n) => SqlValue::Real(*n as f64 + offset),
SqlValue::Float(f) => SqlValue::Real(*f as f64 + offset),
SqlValue::Real(f) => SqlValue::Real(*f + offset),
SqlValue::Double(f) => SqlValue::Real(*f + offset),
SqlValue::Numeric(f) => SqlValue::Real(*f + offset),
_ => value.clone(), }
}
fn calculate_rows_boundary(
bound: &FrameBound,
current_row_idx: usize,
partition_size: usize,
is_start: bool,
) -> usize {
match bound {
FrameBound::UnboundedPreceding => 0,
FrameBound::UnboundedFollowing => partition_size,
FrameBound::CurrentRow => {
if is_start {
current_row_idx
} else {
current_row_idx + 1 }
}
FrameBound::Preceding(offset_expr) => {
let offset = get_numeric_offset(offset_expr) as usize;
if is_start {
current_row_idx.saturating_sub(offset)
} else {
if offset > current_row_idx {
0
} else {
current_row_idx - offset + 1
}
}
}
FrameBound::Following(offset_expr) => {
let offset = get_numeric_offset(offset_expr) as usize;
let result = current_row_idx + offset;
if is_start {
result.min(partition_size)
} else {
(result + 1).min(partition_size) }
}
}
}
#[derive(Debug, Clone)]
pub struct FrameResult {
pub range: Range<usize>,
pub exclude: Option<FrameExclude>,
pub current_row_idx: usize,
}
impl FrameResult {
pub fn includes<F>(
&self,
row_idx: usize,
partition: &Partition,
order_by: &Option<Vec<OrderByItem>>,
eval_fn: &F,
) -> bool
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
if !self.range.contains(&row_idx) {
return false;
}
match self.exclude {
None | Some(FrameExclude::NoOthers) => true,
Some(FrameExclude::CurrentRow) => row_idx != self.current_row_idx,
Some(FrameExclude::Group) => {
!is_peer(row_idx, self.current_row_idx, partition, order_by, eval_fn)
}
Some(FrameExclude::Ties) => {
row_idx == self.current_row_idx
|| !is_peer(row_idx, self.current_row_idx, partition, order_by, eval_fn)
}
}
}
pub fn included_indices<'a, F>(
&'a self,
partition: &'a Partition,
order_by: &'a Option<Vec<OrderByItem>>,
eval_fn: &'a F,
) -> impl Iterator<Item = usize> + 'a
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
self.range.clone().filter(move |&idx| self.includes(idx, partition, order_by, eval_fn))
}
}
pub fn calculate_frame_with_exclusion<F>(
partition: &Partition,
current_row_idx: usize,
order_by: &Option<Vec<OrderByItem>>,
frame_spec: &Option<WindowFrame>,
eval_fn: &F,
) -> FrameResult
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
let range = calculate_frame(partition, current_row_idx, order_by, frame_spec, eval_fn);
let exclude = frame_spec.as_ref().and_then(|f| f.exclude);
FrameResult { range, exclude, current_row_idx }
}
fn is_peer<F>(
row_idx_a: usize,
row_idx_b: usize,
partition: &Partition,
order_by: &Option<Vec<OrderByItem>>,
eval_fn: &F,
) -> bool
where
F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
{
if row_idx_a == row_idx_b {
return true;
}
let order_items = match order_by {
Some(items) if !items.is_empty() => items,
_ => return true,
};
if row_idx_a >= partition.len() || row_idx_b >= partition.len() {
return false;
}
let row_a = &partition.rows[row_idx_a];
let row_b = &partition.rows[row_idx_b];
for order_item in order_items {
let val_a = eval_fn(&order_item.expr, row_a).unwrap_or(SqlValue::Null);
let val_b = eval_fn(&order_item.expr, row_b).unwrap_or(SqlValue::Null);
if compare_values(&val_a, &val_b) != Ordering::Equal {
return false;
}
}
true
}