use std::{collections::BTreeMap, fmt, sync::Arc};
use super::{
AggregateNode, EdgeNormalizeNode, FilterNode, GraphAlgorithmNode, JoinNode, LimitNode, MlAlgorithmNode, NodeRemap,
PlanFacts, ProjectionNode, ResultSink, SchemaRef, SortNode, SortSpec, SourceNode, UnionNode, VectorAlgorithmNode,
WindowNode, compact_reachable_nodes,
graph::{check_node_order, validated_plan_parts},
};
use crate::{
dynamic_filter::{DynamicFilterConsumer, DynamicFilterProducer, DynamicFilterRoute},
error::{QueryError, QueryErrorCode, Result},
};
#[derive(Debug, Clone, PartialEq)]
pub struct QueryPlan {
pub(super) nodes: Vec<QueryNode>,
pub(super) sink: ResultSink,
pub(super) annotations: Annotations,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub(crate) struct Annotations {
pub(crate) dynamic_filter_routes: Vec<DynamicFilterRoute>,
pub(crate) advisory: Vec<super::AdvisoryOutputBinding>,
}
pub(crate) struct RewrittenNodes {
nodes: Vec<QueryNode>,
root: NodeId,
origin: RewriteOrigin,
}
enum RewriteOrigin {
InPlace,
RefinedSources,
Rebuilt(NodeRemap),
}
impl RewrittenNodes {
pub(crate) fn refined_sources(nodes: Vec<QueryNode>, root: NodeId) -> Self {
Self { nodes, root, origin: RewriteOrigin::RefinedSources }
}
pub(crate) fn in_place(nodes: Vec<QueryNode>, root: NodeId) -> Self {
Self { nodes, root, origin: RewriteOrigin::InPlace }
}
pub(crate) fn rebuilt(nodes: Vec<QueryNode>, root: NodeId, origin: BTreeMap<NodeId, NodeId>) -> Self {
let mut sources_per_image = BTreeMap::<NodeId, usize>::new();
for image in origin.values() {
*sources_per_image.entry(*image).or_default() += 1;
}
let origin = origin.into_iter().filter(|(_, image)| sources_per_image[image] == 1).collect();
Self { nodes, root, origin: RewriteOrigin::Rebuilt(NodeRemap::new(origin)) }
}
}
impl QueryPlan {
pub(crate) fn from_parts(nodes: Vec<QueryNode>, sink: ResultSink, annotations: Annotations) -> Result<Self> {
let nodes = validated_plan_parts(nodes, &sink)?;
Ok(Self { nodes, sink, annotations })
}
pub fn root_id(&self) -> NodeId {
self.sink.input()
}
pub fn root_node(&self) -> &QueryNode {
&self.nodes[self.root_id().index()]
}
pub fn sink(&self) -> &ResultSink {
&self.sink
}
pub fn dynamic_filter_routes(&self) -> &[DynamicFilterRoute] {
&self.annotations.dynamic_filter_routes
}
pub(crate) fn with_dynamic_filter_routes(mut self, routes: Vec<DynamicFilterRoute>) -> Self {
self.annotations.dynamic_filter_routes = routes;
self
}
#[cfg(any(test, feature = "test-hooks"))]
#[doc(hidden)]
pub fn with_dynamic_filter_routes_for_test(self, routes: Vec<DynamicFilterRoute>) -> Self {
self.with_dynamic_filter_routes(routes)
}
pub(crate) fn rewrite(&self, rewritten: RewrittenNodes) -> Result<Self> {
let RewrittenNodes { nodes, root, origin } = rewritten;
let sink = |root| self.sink.with_input(root);
match origin {
RewriteOrigin::InPlace => self.rebuilt(nodes, root, &NodeRemap::identity(self.node_count()), sink),
RewriteOrigin::RefinedSources => {
debug_assert_eq!(nodes.len(), self.node_count(), "refining sources keeps every node");
let mut refined = self.rebuilt(nodes, root, &NodeRemap::identity(self.node_count()), sink)?;
refined.annotations.dynamic_filter_routes.clone_from(&self.annotations.dynamic_filter_routes);
Ok(refined)
},
RewriteOrigin::Rebuilt(origin) => self.rebuilt(nodes, root, &origin, sink),
}
}
pub(crate) fn resident_subplan(&self, root: NodeId) -> Result<Self> {
let origin = NodeRemap::identity(self.node_count());
self.rebuilt(self.nodes.clone(), root, &origin, |input| ResultSink::GpuResident { input })
}
fn rebuilt(
&self,
nodes: Vec<QueryNode>,
root: NodeId,
origin: &NodeRemap,
sink: impl FnOnce(NodeId) -> ResultSink,
) -> Result<Self> {
check_node_order(&nodes, root)?;
let (nodes, root, compaction) = compact_reachable_nodes(nodes, root);
let remap = origin.then(&compaction);
let mut rewritten = Self::from_parts(nodes, sink(root), Annotations::default())?;
rewritten.annotations.dynamic_filter_routes = self
.annotations
.dynamic_filter_routes
.iter()
.filter_map(|route| {
let producer = remap_route_producer(&route.producer, &remap, self, &rewritten)?;
let consumer = remap_route_consumer(&route.consumer, &remap, self, &rewritten)?;
Some(DynamicFilterRoute { producer, consumer, ..route.clone() })
})
.collect();
rewritten.annotations.advisory = self.remapped_advisory_outputs(&rewritten, &remap);
Ok(rewritten)
}
fn remapped_advisory_outputs(&self, rewritten: &Self, remap: &NodeRemap) -> Vec<super::AdvisoryOutputBinding> {
let mut unchanged = vec![false; rewritten.node_count()];
for node in rewritten.nodes() {
unchanged[node.id.index()] =
remap.old_id(node.id).and_then(|old_id| self.node(old_id)).is_some_and(|old_node| {
remap.try_translate_inputs(&old_node.kind).as_ref() == Some(&node.kind)
&& old_node.facts == node.facts
&& node.input_ids().iter().all(|input| unchanged[input.index()])
});
}
self.annotations
.advisory
.iter()
.filter_map(|binding| {
let node = rewritten.node(remap.remap(binding.node)?)?;
let mut binding = binding.clone();
binding.node = node.id;
if !unchanged[node.id.index()] {
binding.invalidate(super::AdvisoryUnknownReason::RewrittenSubtree);
binding.schema = Arc::clone(node.output_schema());
}
Some(binding)
})
.collect()
}
#[cfg(any(test, feature = "test-hooks"))]
#[doc(hidden)]
pub fn with_nodes_for_test(&self, nodes: Vec<QueryNode>) -> Result<Self> {
self.rewrite(RewrittenNodes::in_place(nodes, self.root_id()))
}
pub fn clone_with_gpu_resident_sink(&self) -> Self {
let mut plan = self.clone();
plan.sink = ResultSink::GpuResident { input: self.root_id() };
plan
}
pub fn clone_with_deferred_host_materialization_sink(&self) -> Self {
let mut plan = self.clone();
if let ResultSink::HostArrow { input, schema } = &plan.sink {
plan.sink = ResultSink::DeferredHostArrow { input: *input, schema: Arc::clone(schema) };
}
plan
}
pub fn with_host_sink_schema(mut self, schema: SchemaRef) -> Result<Self> {
match &mut self.sink {
ResultSink::HostArrow { schema: host_schema, .. }
| ResultSink::DeferredHostArrow { schema: host_schema, .. } => {
if host_schema.fields() != schema.fields() {
return Err(QueryError::new(
QueryErrorCode::InvalidPlan,
"query_engine.plan.host_sink_schema",
"host sink schema replacement must preserve the existing host sink fields",
));
}
*host_schema = schema;
},
ResultSink::GpuResident { .. } => {
return Err(QueryError::new(
QueryErrorCode::InvalidPlan,
"query_engine.plan.host_sink_schema",
"a GPU-resident plan has no host sink schema to replace",
));
},
}
Ok(self)
}
pub fn with_appended_sort(mut self, spec: SortSpec) -> Self {
let input = self.root_id();
let output_schema = Arc::clone(self.root_node().output_schema());
let root = NodeId(self.nodes.len());
let output_ordering = spec.expressions.clone();
self.nodes.push(QueryNode {
id: root,
kind: PlanNode::Sort(SortNode { input, spec }),
facts: PlanFacts::new(output_schema).with_output_ordering(output_ordering),
});
self.sink = self.sink.with_input(root);
self
}
pub fn advisory_outputs(&self) -> &[super::AdvisoryOutputBinding] {
&self.annotations.advisory
}
pub fn nodes(&self) -> &[QueryNode] {
&self.nodes
}
pub fn node(&self, id: NodeId) -> Option<&QueryNode> {
self.nodes.get(id.index())
}
pub fn output_schema(&self) -> &SchemaRef {
match &self.sink {
ResultSink::GpuResident { input } | ResultSink::DeferredHostArrow { input, .. } => {
self.nodes[input.index()].output_schema()
},
ResultSink::HostArrow { schema, .. } => schema,
}
}
pub fn caller_visible_output_schema(&self) -> &SchemaRef {
self.sink.caller_visible_host_schema().unwrap_or_else(|| self.output_schema())
}
pub fn node_count(&self) -> usize {
self.nodes.len()
}
pub fn contains_join(&self) -> bool {
self.nodes.iter().any(|node| matches!(node.kind, PlanNode::Join(_)))
}
pub(crate) fn contains_graph_algorithm(&self) -> bool {
self.nodes.iter().any(|node| matches!(node.kind, PlanNode::GraphAlgorithm(_)))
}
pub fn contains_edge_normalize(&self) -> bool {
self.nodes.iter().any(|node| matches!(node.kind, PlanNode::EdgeNormalize(_)))
}
pub fn contains_vector_algorithm(&self) -> bool {
self.nodes.iter().any(|node| matches!(node.kind, PlanNode::VectorAlgorithm(_)))
}
pub fn contains_ml_algorithm(&self) -> bool {
self.nodes.iter().any(|node| matches!(node.kind, PlanNode::MlAlgorithm(_)))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct NodeId(pub(crate) usize);
impl NodeId {
pub const fn index(self) -> usize {
self.0
}
#[cfg(any(test, feature = "test-hooks"))]
#[doc(hidden)]
pub const fn from_index_for_test(index: usize) -> Self {
Self(index)
}
}
impl fmt::Display for NodeId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "n{}", self.0)
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct QueryNode {
pub id: NodeId,
pub kind: PlanNode,
pub facts: PlanFacts,
}
impl QueryNode {
pub fn output_schema(&self) -> &SchemaRef {
&self.facts.output_schema
}
pub fn input_ids(&self) -> &[NodeId] {
self.kind.input_ids()
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum PlanNode {
Source(Box<SourceNode>),
Filter(FilterNode),
Projection(ProjectionNode),
Aggregate(AggregateNode),
Window(WindowNode),
Sort(SortNode),
Limit(LimitNode),
Join(JoinNode),
Union(UnionNode),
EdgeNormalize(EdgeNormalizeNode),
GraphAlgorithm(GraphAlgorithmNode),
VectorAlgorithm(VectorAlgorithmNode),
MlAlgorithm(MlAlgorithmNode),
}
impl PlanNode {
pub(crate) fn input_ids(&self) -> &[NodeId] {
match self {
Self::Source(_) => &[],
Self::Filter(node) => std::slice::from_ref(&node.input),
Self::Projection(node) => std::slice::from_ref(&node.input),
Self::Aggregate(node) => std::slice::from_ref(&node.input),
Self::Window(node) => std::slice::from_ref(&node.input),
Self::Sort(node) => std::slice::from_ref(&node.input),
Self::Limit(node) => std::slice::from_ref(&node.input),
Self::Join(node) => &node.inputs,
Self::Union(node) => &node.inputs,
Self::EdgeNormalize(node) => std::slice::from_ref(&node.input),
Self::GraphAlgorithm(node) => &node.inputs,
Self::VectorAlgorithm(node) => &node.inputs,
Self::MlAlgorithm(node) => &node.inputs,
}
}
}
fn remapped_route_endpoint(
old: NodeId,
remap: &NodeRemap,
source: &QueryPlan,
rewritten: &QueryPlan,
) -> Option<NodeId> {
let new = remap.remap(old)?;
let (Some(old_node), Some(new_node)) = (source.node(old), rewritten.node(new)) else {
return None;
};
(remap.try_translate_inputs(&old_node.kind).as_ref() == Some(&new_node.kind) && old_node.facts == new_node.facts)
.then_some(new)
}
fn remap_route_producer(
producer: &DynamicFilterProducer,
remap: &NodeRemap,
source: &QueryPlan,
rewritten: &QueryPlan,
) -> Option<DynamicFilterProducer> {
match producer {
DynamicFilterProducer::HashJoinBuild { join_node_id, build_side, build_input } => {
Some(DynamicFilterProducer::HashJoinBuild {
join_node_id: remapped_route_endpoint(*join_node_id, remap, source, rewritten)?,
build_side: *build_side,
build_input: remapped_route_endpoint(*build_input, remap, source, rewritten)?,
})
},
DynamicFilterProducer::AggregateGroupKeys { aggregate_node_id } => {
Some(DynamicFilterProducer::AggregateGroupKeys {
aggregate_node_id: remapped_route_endpoint(*aggregate_node_id, remap, source, rewritten)?,
})
},
}
}
fn remap_route_consumer(
consumer: &DynamicFilterConsumer,
remap: &NodeRemap,
source: &QueryPlan,
rewritten: &QueryPlan,
) -> Option<DynamicFilterConsumer> {
match consumer {
DynamicFilterConsumer::LocalParquetSource { source_node_id } => {
Some(DynamicFilterConsumer::LocalParquetSource {
source_node_id: remapped_route_endpoint(*source_node_id, remap, source, rewritten)?,
})
},
DynamicFilterConsumer::HashJoinProbe { join_node_id, probe_side } => {
Some(DynamicFilterConsumer::HashJoinProbe {
join_node_id: remapped_route_endpoint(*join_node_id, remap, source, rewritten)?,
probe_side: *probe_side,
})
},
DynamicFilterConsumer::Filter { filter_node_id } => Some(DynamicFilterConsumer::Filter {
filter_node_id: remapped_route_endpoint(*filter_node_id, remap, source, rewritten)?,
}),
}
}