use std::num::{NonZeroU32, NonZeroU64};
use haematite_query::stage::Stage;
use super::super::{matching, refused, rerank, window};
use super::{Executed, WindowContext};
use crate::serve::protocol::Metric;
use crate::serve::query::parser::{self, Plan, Question, Source};
use crate::serve::query::search::Search;
use crate::serve::query::source::Scope;
use crate::serve::storage::Failure;
pub(super) fn mapped(
plan: &Plan,
memory: Option<NonZeroU64>,
window: Option<&WindowContext<'_, '_>>,
) -> Result<Executed, Failure> {
let mut executed = Executed {
scope: Scope::All,
roots: Vec::new(),
search: None,
matching: None,
reranks: Vec::new(),
reading: "pipeline",
stages: Vec::new(),
names: Vec::new(),
};
if let Source::Collection { name, arguments } = &plan.source {
executed.scope = Scope::Collection(name.clone());
if !arguments.is_empty() {
let arguments = arguments
.iter()
.map(|argument| {
Ok((
argument.field.clone(),
serde_json::value::to_raw_value(&argument.value)
.map_err(Failure::storage)?,
))
})
.collect::<Result<_, Failure>>()?;
executed.stages.push(Stage::Filter(arguments));
executed.names.push("filter");
}
}
if let Source::Shaped(roots) = &plan.source {
executed.roots.clone_from(roots);
executed.names.push("{");
}
for (index, stage) in plan.stages.iter().enumerate() {
let Some((value, name)) = mapped_stage(&mut executed, &stage.value, index, memory, window)?
else {
continue;
};
let value = if matches!(value, Stage::Jq)
&& plan.stages[index + 1..]
.iter()
.any(|stage| matches!(stage.value, parser::Stage::Similar { .. }))
{
Stage::RecordJq
} else {
value
};
executed.stages.push(value);
executed.names.push(name);
}
Ok(executed)
}
fn mapped_stage(
executed: &mut Executed,
stage: &parser::Stage,
index: usize,
memory: Option<NonZeroU64>,
window: Option<&WindowContext<'_, '_>>,
) -> Result<Option<(Stage, &'static str)>, Failure> {
let mapped = match stage {
parser::Stage::Similar {
field,
question,
measure,
} => {
if index == 0 && executed.roots.is_empty() {
similar(executed, index, field, question, *measure)?;
return Ok(None);
}
let memory = memory.ok_or_else(|| {
Failure::new(
"query_protocol",
"similar mapping needs query memory settings",
)
})?;
let query = scored(executed, field, question, *measure, "similar", Some(memory))?;
(
Stage::Similar {
index: query,
memory,
},
"similar",
)
}
parser::Stage::Rerank {
field,
question,
measure,
} => (
mapped_rerank(executed, field, question, *measure)?,
"rerank",
),
parser::Stage::Sort { field, descending } => (
Stage::Sort {
field: field.clone(),
descending: *descending,
},
"sort",
),
parser::Stage::Limit(limit) => (Stage::Limit(NonZeroU64::from(*limit)), "limit"),
parser::Stage::Count => (Stage::Count, "count"),
parser::Stage::Percentile { field, percentile } => (
Stage::Percentile(
haematite_query::stage::percentile::Percentile::compile(
field.clone(),
percentile.clone(),
)
.map_err(|error| refused(&error))?,
),
"percentile",
),
parser::Stage::GroupBy { field } => (
Stage::GroupBy {
field: field.clone(),
},
"group_by",
),
parser::Stage::Jq(_) => (Stage::Jq, "jq"),
parser::Stage::Shape(selection) => (Stage::Shape(selection.clone()), "{"),
parser::Stage::Rg {
field,
pattern,
flags,
} => {
let memory = memory.ok_or_else(|| {
Failure::new("query_protocol", "rg mapping needs query memory settings")
})?;
let matcher =
haematite_query::stage::rg::Rg::compile(field.clone(), pattern, flags, memory)
.map_err(|error| refused(&error))?;
(Stage::Rg(matcher), "rg")
}
parser::Stage::Match { nodes, relations } => {
let (request, stage) = matching::mapped(nodes, relations, &mut executed.scope)?;
executed.matching = Some(request);
(stage, "match")
}
parser::Stage::Window(asked) => {
let window = window.ok_or_else(|| {
Failure::new(
"query_protocol",
"window mapping needs its captured history",
)
})?;
(
Stage::Window(window::resolve(
window.context,
window.at,
window.now,
window.scan,
asked,
)?),
"window",
)
}
};
Ok(Some(mapped))
}
fn similar(
executed: &mut Executed,
index: usize,
field: &str,
question: &Question,
measure: Metric,
) -> Result<(), Failure> {
if index != 0 || matches!(executed.scope, Scope::All) {
return Err(Failure::new(
"query_stage_not_available",
"similar runs only first, after a collection",
));
}
executed.search = Some(Search {
field: field.to_owned(),
question: copy_question(question),
measure,
limit: NonZeroU32::MAX,
});
executed.names.push("similar");
Ok(())
}
fn mapped_rerank(
executed: &mut Executed,
field: &str,
question: &Question,
measure: Metric,
) -> Result<Stage, Failure> {
let index = scored(executed, field, question, measure, "rerank", None)?;
Ok(Stage::Rerank { index })
}
fn scored(
executed: &mut Executed,
field: &str,
question: &Question,
measure: Metric,
name: &'static str,
memory: Option<NonZeroU64>,
) -> Result<usize, Failure> {
let collection = match &executed.scope {
Scope::Collection(collection) => collection.clone(),
Scope::All if name == "similar" && executed.roots.len() == 1 => {
executed.roots[0].name.clone()
}
Scope::All => {
return Err(Failure::new(
"query_stage_not_available",
format!("{name} requires a collection"),
));
}
};
let index = executed.reranks.len();
executed.reranks.push(rerank::Request {
field: field.to_owned(),
question: copy_question(question),
measure,
collection,
name,
memory,
});
Ok(index)
}
fn copy_question(question: &Question) -> Question {
match question {
Question::Record(record) => Question::Record(record.clone()),
Question::Text(text) => Question::Text(text.clone()),
}
}