use std::io::{Read, Write};
use vibesql_ast::{
FrameBound, FrameExclude, FrameUnit, WindowFrame, WindowFunctionSpec, WindowSpec,
};
use super::super::io::*;
use crate::StorageError;
pub(super) fn write_window_function_spec<W: Write>(
writer: &mut W,
spec: &WindowFunctionSpec,
) -> Result<(), StorageError> {
match spec {
WindowFunctionSpec::Aggregate { name, args, filter } => {
writer
.write_all(&[0u8])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
write_string(writer, name.canonical())?;
write_u32(writer, args.len() as u32)?;
for arg in args {
super::write_expression(writer, arg)?;
}
if let Some(f) = filter {
writer
.write_all(&[1u8])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
super::write_expression(writer, f)?;
} else {
writer
.write_all(&[0u8])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
}
}
WindowFunctionSpec::Ranking { name, args } => {
writer
.write_all(&[1u8])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
write_string(writer, name.canonical())?;
write_u32(writer, args.len() as u32)?;
for arg in args {
super::write_expression(writer, arg)?;
}
}
WindowFunctionSpec::Value { name, args } => {
writer
.write_all(&[2u8])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
write_string(writer, name.canonical())?;
write_u32(writer, args.len() as u32)?;
for arg in args {
super::write_expression(writer, arg)?;
}
}
}
Ok(())
}
pub(super) fn read_window_function_spec<R: Read>(
reader: &mut R,
) -> Result<WindowFunctionSpec, StorageError> {
let tag = read_u8(reader)?;
match tag {
0 => {
let name_str = read_string(reader)?;
let name = vibesql_ast::FunctionIdentifier::new(&name_str);
let arg_count = read_u32(reader)?;
let mut args = Vec::new();
for _ in 0..arg_count {
args.push(super::read_expression(reader)?);
}
let has_filter = read_u8(reader)?;
let filter = if has_filter == 1 {
Some(Box::new(super::read_expression(reader)?))
} else {
None
};
Ok(WindowFunctionSpec::Aggregate { name, args, filter })
}
1 => {
let name_str = read_string(reader)?;
let name = vibesql_ast::FunctionIdentifier::new(&name_str);
let arg_count = read_u32(reader)?;
let mut args = Vec::new();
for _ in 0..arg_count {
args.push(super::read_expression(reader)?);
}
Ok(WindowFunctionSpec::Ranking { name, args })
}
2 => {
let name_str = read_string(reader)?;
let name = vibesql_ast::FunctionIdentifier::new(&name_str);
let arg_count = read_u32(reader)?;
let mut args = Vec::new();
for _ in 0..arg_count {
args.push(super::read_expression(reader)?);
}
Ok(WindowFunctionSpec::Value { name, args })
}
_ => {
Err(StorageError::NotImplemented(format!("Unknown window function spec tag: {}", tag)))
}
}
}
pub(super) fn write_window_spec<W: Write>(
writer: &mut W,
spec: &WindowSpec,
) -> Result<(), StorageError> {
write_bool(writer, spec.partition_by.is_some())?;
if let Some(partition_exprs) = &spec.partition_by {
write_u32(writer, partition_exprs.len() as u32)?;
for expr in partition_exprs {
super::write_expression(writer, expr)?;
}
}
write_bool(writer, false)?;
write_bool(writer, spec.frame.is_some())?;
if let Some(frame) = &spec.frame {
write_window_frame(writer, frame)?;
}
Ok(())
}
pub(super) fn read_window_spec<R: Read>(reader: &mut R) -> Result<WindowSpec, StorageError> {
let has_partition_by = read_bool(reader)?;
let partition_by = if has_partition_by {
let count = read_u32(reader)?;
let mut exprs = Vec::new();
for _ in 0..count {
exprs.push(super::read_expression(reader)?);
}
Some(exprs)
} else {
None
};
let has_order_by = read_bool(reader)?;
if has_order_by {
return Err(StorageError::NotImplemented(
"OrderBy in window spec not yet supported".to_string(),
));
}
let has_frame = read_bool(reader)?;
let frame = if has_frame { Some(read_window_frame(reader)?) } else { None };
Ok(WindowSpec { base_window_name: None, partition_by, order_by: None, frame })
}
fn write_window_frame<W: Write>(writer: &mut W, frame: &WindowFrame) -> Result<(), StorageError> {
let unit_tag = match frame.unit {
FrameUnit::Rows => 0u8,
FrameUnit::Range => 1,
FrameUnit::Groups => 2,
};
writer
.write_all(&[unit_tag])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
write_frame_bound(writer, &frame.start)?;
write_bool(writer, frame.end.is_some())?;
if let Some(end) = &frame.end {
write_frame_bound(writer, end)?;
}
write_bool(writer, frame.exclude.is_some())?;
if let Some(exclude) = &frame.exclude {
let exclude_tag = match exclude {
FrameExclude::NoOthers => 0u8,
FrameExclude::CurrentRow => 1,
FrameExclude::Group => 2,
FrameExclude::Ties => 3,
};
writer
.write_all(&[exclude_tag])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
}
Ok(())
}
fn read_window_frame<R: Read>(reader: &mut R) -> Result<WindowFrame, StorageError> {
let unit_tag = read_u8(reader)?;
let unit = match unit_tag {
0 => FrameUnit::Rows,
1 => FrameUnit::Range,
2 => FrameUnit::Groups,
_ => {
return Err(StorageError::NotImplemented(format!(
"Unknown frame unit tag: {}",
unit_tag
)))
}
};
let start = read_frame_bound(reader)?;
let has_end = read_bool(reader)?;
let end = if has_end { Some(read_frame_bound(reader)?) } else { None };
let has_exclude = read_bool(reader)?;
let exclude = if has_exclude {
let exclude_tag = read_u8(reader)?;
let exclude_mode = match exclude_tag {
0 => FrameExclude::NoOthers,
1 => FrameExclude::CurrentRow,
2 => FrameExclude::Group,
3 => FrameExclude::Ties,
_ => {
return Err(StorageError::NotImplemented(format!(
"Unknown frame exclude tag: {}",
exclude_tag
)))
}
};
Some(exclude_mode)
} else {
None
};
Ok(WindowFrame { unit, start, end, exclude })
}
fn write_frame_bound<W: Write>(writer: &mut W, bound: &FrameBound) -> Result<(), StorageError> {
match bound {
FrameBound::UnboundedPreceding => {
writer
.write_all(&[0u8])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
}
FrameBound::Preceding(expr) => {
writer
.write_all(&[1u8])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
super::write_expression(writer, expr)?;
}
FrameBound::CurrentRow => {
writer
.write_all(&[2u8])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
}
FrameBound::Following(expr) => {
writer
.write_all(&[3u8])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
super::write_expression(writer, expr)?;
}
FrameBound::UnboundedFollowing => {
writer
.write_all(&[4u8])
.map_err(|e| StorageError::NotImplemented(format!("Write error: {}", e)))?;
}
}
Ok(())
}
fn read_frame_bound<R: Read>(reader: &mut R) -> Result<FrameBound, StorageError> {
let tag = read_u8(reader)?;
match tag {
0 => Ok(FrameBound::UnboundedPreceding),
1 => {
let expr = Box::new(super::read_expression(reader)?);
Ok(FrameBound::Preceding(expr))
}
2 => Ok(FrameBound::CurrentRow),
3 => {
let expr = Box::new(super::read_expression(reader)?);
Ok(FrameBound::Following(expr))
}
4 => Ok(FrameBound::UnboundedFollowing),
_ => Err(StorageError::NotImplemented(format!("Unknown frame bound tag: {}", tag))),
}
}