use alloc::boxed::Box;
use alloc::string::String;
use alloc::vec::Vec;
use crate::ast::{
CollectionMode, CreateCollectionStmt, CreateShardKeyStmt, SparseVectorDef, Stmt, VectorDef,
VectorDistance,
};
use crate::error::QqlError;
use crate::token::TokenKind;
use super::{AstLowerer, ascii_equal};
impl<'a> AstLowerer<'a> {
pub fn parse_create(&mut self) -> Result<Stmt, QqlError> {
self.advance()?;
let tok = self.peek()?;
if tok.kind == TokenKind::Index {
return self.parse_create_index();
}
if tok.kind == TokenKind::Shard {
return self.parse_create_shard_key();
}
self.expect(TokenKind::Collection)?;
let collection = self.parse_identifier()?;
let mut hybrid = false;
let mut rerank = false;
let mut model: Option<String> = None;
let mut dense_vector: Option<String> = None;
let mut sparse_vector: Option<String> = None;
let mut explicit_vectors: Vec<VectorDef> = Vec::new();
let mut explicit_sparse_vectors: Vec<SparseVectorDef> = Vec::new();
if self.peek()?.kind == TokenKind::Hybrid {
self.advance()?;
hybrid = true;
if self.peek()?.kind == TokenKind::Rerank {
self.advance()?;
rerank = true;
} else {
while self.peek()?.kind == TokenKind::Dense
|| self.peek()?.kind == TokenKind::Sparse
{
let mode = self.advance()?.kind;
let tok = self.peek()?;
if tok.kind == TokenKind::Vector
|| (tok.kind == TokenKind::Identifier && ascii_equal(tok.text, "VECTOR"))
{
self.advance()?;
let v = self.parse_identifier()?;
if mode == TokenKind::Dense {
dense_vector = Some(v);
} else {
sparse_vector = Some(v);
}
} else {
return Err(QqlError::parse(
"QQL-PARSE-SYNTAX",
"expected VECTOR after DENSE/SPARSE",
self.peek()?.span,
));
}
}
}
} else if self.peek()?.kind == TokenKind::Using {
self.advance()?;
if self.peek()?.kind == TokenKind::Hybrid {
self.advance()?;
hybrid = true;
if self.peek()?.kind == TokenKind::Dense {
return Err(QqlError::validation(
"QQL-VALIDATION-CREATE-MODEL",
"HYBRID does not accept a single dense MODEL; configure the executor model or provide explicit vector dimensions",
Some(self.peek()?.span),
));
}
} else {
if self.peek()?.kind == TokenKind::Dense {
self.advance()?;
}
model = Some(self.parse_required_model_string()?);
}
}
if self.peek()?.kind == TokenKind::Lparen {
self.advance()?;
while self.peek()?.kind != TokenKind::Rparen && self.peek()?.kind != TokenKind::Eof {
let name = self.parse_identifier()?;
if self.peek()?.kind == TokenKind::Vector {
self.advance()?;
self.expect(TokenKind::Lparen)?;
let size_tok = self.peek()?;
let size = self.parse_numeric_literal()?;
if size <= 0.0 || size != (size as u64) as f64 {
return Err(QqlError::parse(
"QQL-PARSE-SYNTAX",
"vector size must be a positive integer",
size_tok.span,
));
}
if size > 65536.0 {
return Err(QqlError::parse(
"QQL-PARSE-VECTOR-SIZE",
"vector size must not exceed 65536",
size_tok.span,
));
}
self.expect(TokenKind::Comma)?;
let dist_tok = self.peek()?;
let distance = match dist_tok.kind {
TokenKind::Cosine => VectorDistance::Cosine,
TokenKind::Dot => VectorDistance::Dot,
TokenKind::Euclid => VectorDistance::Euclid,
TokenKind::Manhattan => VectorDistance::Manhattan,
_ => {
return Err(QqlError::parse(
"QQL-PARSE-SYNTAX",
"expected distance metric (COSINE, DOT, EUCLID, MANHATTAN)",
dist_tok.span,
));
}
};
self.advance()?;
self.expect(TokenKind::Rparen)?;
let mut hnsw = None;
let mut quant = None;
let mut multiv = None;
let mut vec_cfg = None;
while self.peek()?.kind == TokenKind::With {
self.advance()?;
if self.peek()?.kind == TokenKind::Hnsw {
self.advance()?;
hnsw = self.parse_hnsw_config_block()?.hnsw;
} else if self.peek()?.kind == TokenKind::Quantization
|| (self.peek()?.is_keyword_or_identifier()
&& ascii_equal(self.peek()?.text, "QUANTIZATION"))
{
self.advance()?;
quant = self.parse_quantization_config_block()?.quantization;
} else if self.peek()?.kind == TokenKind::Multivector
|| (self.peek()?.is_keyword_or_identifier()
&& ascii_equal(self.peek()?.text, "MULTIVECTOR"))
{
self.advance()?;
multiv = Some(self.parse_multivector_config_block()?);
} else if self.peek()?.kind == TokenKind::Vector {
self.advance()?;
vec_cfg = self.parse_vectors_config_block()?.vectors;
} else {
return Err(QqlError::parse(
"QQL-PARSE-SYNTAX",
"expected HNSW, QUANTIZATION, MULTIVECTOR, or VECTOR after WITH for vector configuration",
self.peek()?.span,
));
}
}
explicit_vectors.push(VectorDef {
name,
size: size as u64,
distance,
hnsw,
quantization: quant,
multivector: multiv,
vectors: vec_cfg,
});
} else if self.peek()?.kind == TokenKind::Sparse {
self.advance()?;
let mut index: Option<Box<crate::ast::SparseIndexConfig>> = None;
let mut modifier = None;
while self.peek()?.kind == TokenKind::With {
self.advance()?;
if self.peek()?.kind == TokenKind::Sparse
|| self.peek()?.kind == TokenKind::Index
{
self.advance()?;
let (idx, mod_val) = self.parse_sparse_config_block()?;
if let Some(new_idx) = idx {
if let Some(ref mut existing) = index {
if new_idx.full_scan_threshold.is_some() {
existing.full_scan_threshold = new_idx.full_scan_threshold;
}
if new_idx.on_disk.is_some() {
existing.on_disk = new_idx.on_disk;
}
if new_idx.datatype.is_some() {
existing.datatype = new_idx.datatype;
}
if new_idx.memory.is_some() {
existing.memory = new_idx.memory;
}
} else {
index = Some(new_idx);
}
}
if mod_val.is_some() {
modifier = mod_val;
}
} else {
return Err(QqlError::parse(
"QQL-PARSE-SYNTAX",
"expected SPARSE or INDEX after WITH for sparse vector configuration",
self.peek()?.span,
));
}
}
explicit_sparse_vectors.push(SparseVectorDef {
name,
index,
modifier,
});
} else {
return Err(QqlError::parse(
"QQL-PARSE-SYNTAX",
"expected VECTOR or SPARSE after vector name",
self.peek()?.span,
));
}
if self.peek()?.kind == TokenKind::Comma {
self.advance()?;
} else if self.peek()?.kind != TokenKind::Rparen {
return Err(QqlError::parse(
"QQL-PARSE-SYNTAX",
"expected comma or )",
self.peek()?.span,
));
}
}
self.expect(TokenKind::Rparen)?;
}
let config = self.parse_collection_config_blocks(false)?;
let mode = if rerank {
CollectionMode::Rerank
} else if hybrid {
CollectionMode::Hybrid {
dense_vector,
sparse_vector,
}
} else {
CollectionMode::Dense { model }
};
Ok(Stmt::CreateCollection(Box::new(CreateCollectionStmt {
collection,
mode,
vectors: explicit_vectors,
sparse_vectors: explicit_sparse_vectors,
config,
})))
}
pub fn parse_create_shard_key(&mut self) -> Result<Stmt, QqlError> {
self.expect(TokenKind::Shard)?;
self.expect(TokenKind::Key)?;
let shard_name = self.parse_shard_key_atom()?;
self.expect(TokenKind::On)?;
self.expect(TokenKind::Collection)?;
let collection = self.parse_identifier()?;
let mut shards_number = None;
let mut replication_factor = None;
let mut placement = None;
let mut initial_state = None;
if self.peek()?.kind == TokenKind::With {
self.advance()?;
let opts = self.parse_config_block()?;
for (key, value) in &opts {
let key_lower = key.to_ascii_lowercase();
match key_lower.as_str() {
"shards_number" => {
shards_number = Some(self.positive_shard_key_u64("shards_number", value)?);
}
"replication_factor" => {
replication_factor =
Some(self.positive_shard_key_u64("replication_factor", value)?);
}
"placement" => {
placement = Some(self.shard_key_placement(value)?);
}
"initial_state" => {
initial_state = Some(self.shard_key_initial_state(value)?);
}
_ => {
return Err(QqlError::parse(
"QQL-PARSE-SHARD-KEY-CONFIG",
alloc::format!(
"unknown CREATE SHARD KEY parameter '{}'. Expected: shards_number, replication_factor, placement, initial_state",
key
),
self.peek()?.span,
));
}
}
}
}
Ok(Stmt::CreateShardKey(Box::new(CreateShardKeyStmt {
collection,
shard_key: shard_name,
shards_number,
replication_factor,
placement,
initial_state,
})))
}
fn positive_shard_key_u64(
&mut self,
name: &str,
value: &crate::ast::Value,
) -> Result<u64, QqlError> {
match value {
crate::ast::Value::Int(n) if *n > 0 => Ok(*n as u64),
_ => Err(QqlError::parse(
"QQL-PARSE-SHARD-KEY-CONFIG",
alloc::format!("{} must be a positive integer", name),
self.peek()?.span,
)),
}
}
fn shard_key_placement(&mut self, value: &crate::ast::Value) -> Result<Vec<u64>, QqlError> {
match value {
crate::ast::Value::List(items) if !items.is_empty() => items
.iter()
.map(|item| match item {
crate::ast::Value::Int(n) if *n >= 0 => Ok(*n as u64),
_ => Err(QqlError::parse(
"QQL-PARSE-SHARD-KEY-CONFIG",
"placement entries must be non-negative integers (peer ids)",
self.peek()?.span,
)),
})
.collect(),
_ => Err(QqlError::parse(
"QQL-PARSE-SHARD-KEY-CONFIG",
"placement must be a non-empty list of peer ids",
self.peek()?.span,
)),
}
}
fn shard_key_initial_state(&mut self, value: &crate::ast::Value) -> Result<String, QqlError> {
let span = self.peek()?.span;
match value {
crate::ast::Value::Str(raw) => canonical_replica_state(raw).ok_or_else(|| {
QqlError::parse(
"QQL-PARSE-SHARD-KEY-CONFIG",
alloc::format!("unknown replica state '{raw}'"),
span,
)
}),
_ => Err(QqlError::parse(
"QQL-PARSE-SHARD-KEY-CONFIG",
"initial_state must be a string",
span,
)),
}
}
}
const REPLICA_STATES: &[&str] = &[
"Active",
"Dead",
"Partial",
"Initializing",
"Listener",
"PartialSnapshot",
"Recovery",
"Resharding",
"ReshardingScaleDown",
"ActiveRead",
"ManualRecovery",
];
fn canonical_replica_state(raw: &str) -> Option<String> {
REPLICA_STATES
.iter()
.find(|state| state.eq_ignore_ascii_case(raw))
.map(|state| state.to_string())
}