use crate::ast::expr::{BinaryOp, Literal, UnaryOp};
use crate::ast::span::Span;
use crate::planner::logical_plan::LogicalPlan;
use crate::planner::types::ResolvedType;
#[derive(Debug, Clone)]
pub struct TypedExpr {
pub kind: TypedExprKind,
pub resolved_type: ResolvedType,
pub span: Span,
}
#[derive(Debug, Clone)]
pub enum TypedExprKind {
Literal(Literal),
ColumnRef {
table: String,
column: String,
column_index: usize,
},
BinaryOp {
left: Box<TypedExpr>,
op: BinaryOp,
right: Box<TypedExpr>,
},
UnaryOp {
op: UnaryOp,
operand: Box<TypedExpr>,
},
FunctionCall {
name: String,
args: Vec<TypedExpr>,
distinct: bool,
star: bool,
},
Cast {
expr: Box<TypedExpr>,
target_type: ResolvedType,
},
Between {
expr: Box<TypedExpr>,
low: Box<TypedExpr>,
high: Box<TypedExpr>,
negated: bool,
},
Like {
expr: Box<TypedExpr>,
pattern: Box<TypedExpr>,
escape: Option<Box<TypedExpr>>,
negated: bool,
},
InList {
expr: Box<TypedExpr>,
list: Vec<TypedExpr>,
negated: bool,
},
IsNull {
expr: Box<TypedExpr>,
negated: bool,
},
VectorLiteral(Vec<f64>),
ScalarSubquery(Box<LogicalPlan>),
InSubquery {
expr: Box<TypedExpr>,
subquery: Box<LogicalPlan>,
negated: bool,
},
Exists {
subquery: Box<LogicalPlan>,
negated: bool,
},
Quantified {
expr: Box<TypedExpr>,
op: BinaryOp,
quantifier: Quantifier,
subquery: Box<LogicalPlan>,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Quantifier {
Any,
All,
}
#[derive(Debug, Clone)]
pub struct SortExpr {
pub expr: TypedExpr,
pub asc: bool,
pub nulls_first: bool,
}
#[derive(Debug, Clone)]
pub struct TypedAssignment {
pub column: String,
pub column_index: usize,
pub value: TypedExpr,
}
#[derive(Debug, Clone)]
pub struct ProjectedColumn {
pub expr: TypedExpr,
pub alias: Option<String>,
}
#[derive(Debug, Clone)]
pub enum Projection {
All(Vec<String>),
Columns(Vec<ProjectedColumn>),
}
impl TypedExpr {
pub fn new(kind: TypedExprKind, resolved_type: ResolvedType, span: Span) -> Self {
Self {
kind,
resolved_type,
span,
}
}
pub fn literal(lit: Literal, resolved_type: ResolvedType, span: Span) -> Self {
Self::new(TypedExprKind::Literal(lit), resolved_type, span)
}
pub fn column_ref(
table: String,
column: String,
column_index: usize,
resolved_type: ResolvedType,
span: Span,
) -> Self {
Self::new(
TypedExprKind::ColumnRef {
table,
column,
column_index,
},
resolved_type,
span,
)
}
pub fn binary_op(
left: TypedExpr,
op: BinaryOp,
right: TypedExpr,
resolved_type: ResolvedType,
span: Span,
) -> Self {
Self::new(
TypedExprKind::BinaryOp {
left: Box::new(left),
op,
right: Box::new(right),
},
resolved_type,
span,
)
}
pub fn unary_op(
op: UnaryOp,
operand: TypedExpr,
resolved_type: ResolvedType,
span: Span,
) -> Self {
Self::new(
TypedExprKind::UnaryOp {
op,
operand: Box::new(operand),
},
resolved_type,
span,
)
}
pub fn function_call(
name: String,
args: Vec<TypedExpr>,
distinct: bool,
star: bool,
resolved_type: ResolvedType,
span: Span,
) -> Self {
Self::new(
TypedExprKind::FunctionCall {
name,
args,
distinct,
star,
},
resolved_type,
span,
)
}
pub fn cast(expr: TypedExpr, target_type: ResolvedType, span: Span) -> Self {
Self::new(
TypedExprKind::Cast {
expr: Box::new(expr),
target_type: target_type.clone(),
},
target_type,
span,
)
}
pub fn vector_literal(values: Vec<f64>, dimension: u32, span: Span) -> Self {
use crate::ast::ddl::VectorMetric;
Self::new(
TypedExprKind::VectorLiteral(values),
ResolvedType::Vector {
dimension,
metric: VectorMetric::Cosine,
},
span,
)
}
}
impl SortExpr {
pub fn asc(expr: TypedExpr) -> Self {
Self {
expr,
asc: true,
nulls_first: false,
}
}
pub fn desc(expr: TypedExpr) -> Self {
Self {
expr,
asc: false,
nulls_first: false,
}
}
pub fn new(expr: TypedExpr, asc: bool, nulls_first: bool) -> Self {
Self {
expr,
asc,
nulls_first,
}
}
}
impl TypedAssignment {
pub fn new(column: String, column_index: usize, value: TypedExpr) -> Self {
Self {
column,
column_index,
value,
}
}
}
impl ProjectedColumn {
pub fn new(expr: TypedExpr) -> Self {
Self { expr, alias: None }
}
pub fn with_alias(expr: TypedExpr, alias: String) -> Self {
Self {
expr,
alias: Some(alias),
}
}
pub fn output_name(&self) -> Option<&str> {
if let Some(ref alias) = self.alias {
return Some(alias);
}
if let TypedExprKind::ColumnRef { ref column, .. } = self.expr.kind {
return Some(column);
}
None
}
}
impl Projection {
pub fn len(&self) -> usize {
match self {
Projection::All(cols) => cols.len(),
Projection::Columns(cols) => cols.len(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn column_names(&self) -> Vec<Option<&str>> {
match self {
Projection::All(cols) => cols.iter().map(|s| Some(s.as_str())).collect(),
Projection::Columns(cols) => cols.iter().map(|c| c.output_name()).collect(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ast::ddl::VectorMetric;
#[test]
fn test_typed_expr_literal() {
let expr = TypedExpr::literal(
Literal::Number("42".to_string()),
ResolvedType::Integer,
Span::default(),
);
assert!(matches!(
expr.kind,
TypedExprKind::Literal(Literal::Number(_))
));
assert_eq!(expr.resolved_type, ResolvedType::Integer);
}
#[test]
fn test_typed_expr_column_ref() {
let expr = TypedExpr::column_ref(
"users".to_string(),
"id".to_string(),
0,
ResolvedType::Integer,
Span::default(),
);
if let TypedExprKind::ColumnRef {
table,
column,
column_index,
} = &expr.kind
{
assert_eq!(table, "users");
assert_eq!(column, "id");
assert_eq!(*column_index, 0);
} else {
panic!("Expected ColumnRef");
}
}
#[test]
fn test_typed_expr_binary_op() {
let left = TypedExpr::literal(
Literal::Number("1".to_string()),
ResolvedType::Integer,
Span::default(),
);
let right = TypedExpr::literal(
Literal::Number("2".to_string()),
ResolvedType::Integer,
Span::default(),
);
let expr = TypedExpr::binary_op(
left,
BinaryOp::Add,
right,
ResolvedType::Integer,
Span::default(),
);
assert!(matches!(expr.kind, TypedExprKind::BinaryOp { .. }));
assert_eq!(expr.resolved_type, ResolvedType::Integer);
}
#[test]
fn test_typed_expr_vector_literal() {
let values = vec![1.0, 2.0, 3.0];
let expr = TypedExpr::vector_literal(values.clone(), 3, Span::default());
if let TypedExprKind::VectorLiteral(v) = &expr.kind {
assert_eq!(v, &values);
} else {
panic!("Expected VectorLiteral");
}
if let ResolvedType::Vector { dimension, metric } = &expr.resolved_type {
assert_eq!(*dimension, 3);
assert_eq!(*metric, VectorMetric::Cosine);
} else {
panic!("Expected Vector type");
}
}
#[test]
fn test_sort_expr_asc() {
let col = TypedExpr::column_ref(
"users".to_string(),
"name".to_string(),
1,
ResolvedType::Text,
Span::default(),
);
let sort = SortExpr::asc(col);
assert!(sort.asc);
assert!(!sort.nulls_first);
}
#[test]
fn test_sort_expr_desc() {
let col = TypedExpr::column_ref(
"users".to_string(),
"name".to_string(),
1,
ResolvedType::Text,
Span::default(),
);
let sort = SortExpr::desc(col);
assert!(!sort.asc);
assert!(!sort.nulls_first);
}
#[test]
fn test_typed_assignment() {
let value = TypedExpr::literal(
Literal::String("Alice".to_string()),
ResolvedType::Text,
Span::default(),
);
let assignment = TypedAssignment::new("name".to_string(), 1, value);
assert_eq!(assignment.column, "name");
assert_eq!(assignment.column_index, 1);
}
#[test]
fn test_projected_column_output_name() {
let col = TypedExpr::column_ref(
"users".to_string(),
"name".to_string(),
1,
ResolvedType::Text,
Span::default(),
);
let proj1 = ProjectedColumn::new(col.clone());
assert_eq!(proj1.output_name(), Some("name"));
let proj2 = ProjectedColumn::with_alias(col, "user_name".to_string());
assert_eq!(proj2.output_name(), Some("user_name"));
}
#[test]
fn test_projection_all() {
let columns = vec!["id".to_string(), "name".to_string(), "email".to_string()];
let proj = Projection::All(columns);
assert_eq!(proj.len(), 3);
assert!(!proj.is_empty());
let names: Vec<_> = proj.column_names();
assert_eq!(names, vec![Some("id"), Some("name"), Some("email")]);
}
#[test]
fn test_projection_columns() {
let col1 = ProjectedColumn::new(TypedExpr::column_ref(
"users".to_string(),
"id".to_string(),
0,
ResolvedType::Integer,
Span::default(),
));
let col2 = ProjectedColumn::with_alias(
TypedExpr::column_ref(
"users".to_string(),
"name".to_string(),
1,
ResolvedType::Text,
Span::default(),
),
"user_name".to_string(),
);
let proj = Projection::Columns(vec![col1, col2]);
assert_eq!(proj.len(), 2);
let names: Vec<_> = proj.column_names();
assert_eq!(names, vec![Some("id"), Some("user_name")]);
}
#[test]
fn test_typed_expr_cast() {
let inner = TypedExpr::literal(
Literal::Number("42".to_string()),
ResolvedType::Integer,
Span::default(),
);
let expr = TypedExpr::cast(inner, ResolvedType::Double, Span::default());
assert!(matches!(expr.kind, TypedExprKind::Cast { .. }));
assert_eq!(expr.resolved_type, ResolvedType::Double);
}
#[test]
fn test_typed_expr_kind_between() {
let expr_kind = TypedExprKind::Between {
expr: Box::new(TypedExpr::column_ref(
"t".to_string(),
"x".to_string(),
0,
ResolvedType::Integer,
Span::default(),
)),
low: Box::new(TypedExpr::literal(
Literal::Number("1".to_string()),
ResolvedType::Integer,
Span::default(),
)),
high: Box::new(TypedExpr::literal(
Literal::Number("10".to_string()),
ResolvedType::Integer,
Span::default(),
)),
negated: false,
};
assert!(matches!(
expr_kind,
TypedExprKind::Between { negated: false, .. }
));
}
#[test]
fn test_typed_expr_kind_is_null() {
let expr_kind = TypedExprKind::IsNull {
expr: Box::new(TypedExpr::column_ref(
"t".to_string(),
"x".to_string(),
0,
ResolvedType::Integer,
Span::default(),
)),
negated: true,
};
assert!(matches!(
expr_kind,
TypedExprKind::IsNull { negated: true, .. }
));
}
}