use serde::{Deserialize, Serialize};
use std::fmt;
use super::aggregation::{AggregateFunction, GroupByClause, HavingClause};
use super::condition::Condition;
use super::fusion::FusionClause;
use super::join::JoinClause;
use super::values::VectorExpr;
use super::with_clause::WithClause;
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
#[non_exhaustive]
pub enum DistinctMode {
#[default]
None,
All,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SelectStatement {
#[serde(default)]
pub distinct: DistinctMode,
pub columns: SelectColumns,
pub from: String,
#[serde(default)]
pub from_alias: Vec<String>,
#[serde(default)]
pub joins: Vec<JoinClause>,
pub where_clause: Option<Condition>,
pub order_by: Option<Vec<SelectOrderBy>>,
pub limit: Option<u64>,
pub offset: Option<u64>,
pub with_clause: Option<WithClause>,
#[serde(default)]
pub group_by: Option<GroupByClause>,
#[serde(default)]
pub having: Option<HavingClause>,
#[serde(default)]
pub fusion_clause: Option<FusionClause>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum SelectColumns {
All,
Columns(Vec<Column>),
Aggregations(Vec<AggregateFunction>),
Mixed {
columns: Vec<Column>,
aggregations: Vec<AggregateFunction>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
similarity_scores: Vec<SimilarityScoreExpr>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
qualified_wildcards: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
window_functions: Vec<super::window::WindowFunction>,
},
SimilarityScore(SimilarityScoreExpr),
QualifiedWildcard(String),
}
impl SelectColumns {
#[must_use]
pub fn to_display_names(&self) -> Vec<String> {
match self {
Self::All => vec!["*".to_string()],
Self::Columns(cols) => cols.iter().map(|c| c.name.clone()).collect(),
Self::Aggregations(aggs) => aggs
.iter()
.map(|a| format!("{:?}", a.function_type))
.collect(),
Self::Mixed {
columns,
aggregations,
similarity_scores,
qualified_wildcards,
window_functions,
} => {
let mut result: Vec<String> = columns.iter().map(|c| c.name.clone()).collect();
result.extend(
aggregations
.iter()
.map(|a| format!("{:?}", a.function_type)),
);
result.extend(similarity_scores.iter().map(|expr| {
expr.alias
.clone()
.unwrap_or_else(|| "similarity".to_string())
}));
result.extend(qualified_wildcards.iter().map(|alias| format!("{alias}.*")));
result.extend(window_functions.iter().map(|wf| {
wf.alias
.clone()
.unwrap_or_else(|| wf.function_type.default_alias().to_string())
}));
result
}
Self::SimilarityScore(expr) => {
vec![expr
.alias
.clone()
.unwrap_or_else(|| "similarity".to_string())]
}
Self::QualifiedWildcard(alias) => vec![format!("{alias}.*")],
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SimilarityScoreExpr {
pub alias: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Column {
pub name: String,
pub alias: Option<String>,
}
impl Column {
#[must_use]
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
alias: None,
}
}
#[must_use]
pub fn with_alias(name: impl Into<String>, alias: impl Into<String>) -> Self {
Self {
name: name.into(),
alias: Some(alias.into()),
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SelectOrderBy {
pub expr: OrderByExpr,
pub descending: bool,
}
impl SelectOrderBy {
#[must_use]
pub fn to_display_pair(&self) -> (String, String) {
let dir = if self.descending { "DESC" } else { "ASC" };
let col = match &self.expr {
OrderByExpr::Field(f) => f.clone(),
OrderByExpr::Similarity(_) | OrderByExpr::SimilarityBare => "similarity()".to_string(),
OrderByExpr::Aggregate(agg) => format!("{:?}", agg.function_type),
OrderByExpr::Arithmetic(expr) => format!("{expr}"),
};
(col, dir.to_string())
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum OrderByExpr {
Field(String),
Similarity(SimilarityOrderBy),
SimilarityBare,
Aggregate(AggregateFunction),
Arithmetic(ArithmeticExpr),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct LetBinding {
pub name: String,
pub expr: ArithmeticExpr,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum ArithmeticExpr {
Literal(f64),
Variable(String),
Similarity(Box<OrderByExpr>),
BinaryOp {
left: Box<ArithmeticExpr>,
op: ArithmeticOp,
right: Box<ArithmeticExpr>,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum ArithmeticOp {
Add,
Sub,
Mul,
Div,
}
impl fmt::Display for ArithmeticOp {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Add => write!(f, "+"),
Self::Sub => write!(f, "-"),
Self::Mul => write!(f, "*"),
Self::Div => write!(f, "/"),
}
}
}
impl fmt::Display for ArithmeticExpr {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Literal(v) => write!(f, "{v}"),
Self::Variable(name) => write!(f, "{name}"),
Self::Similarity(inner) => match inner.as_ref() {
OrderByExpr::Similarity(sim) => {
let vec_str = match &sim.vector {
VectorExpr::Parameter(name) => format!("${name}"),
VectorExpr::Literal(vals) => format!("{vals:?}"),
};
write!(f, "similarity({}, {vec_str})", sim.field)
}
_ => write!(f, "similarity()"),
},
Self::BinaryOp { left, op, right } => write!(f, "({left} {op} {right})"),
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SimilarityOrderBy {
pub field: String,
pub vector: VectorExpr,
}
impl SelectStatement {
#[must_use]
pub fn empty() -> Self {
Self {
distinct: DistinctMode::None,
columns: SelectColumns::All,
from: String::new(),
from_alias: Vec::new(),
joins: Vec::new(),
where_clause: None,
order_by: None,
limit: None,
offset: None,
with_clause: None,
group_by: None,
having: None,
fusion_clause: None,
}
}
}