use super::*;
use crate::pipeline::BindingBatchSource;
use crate::transform as executor_transform;
struct PreparedTransformSource<'a> {
context: BatchReadContext<'a>,
}
pub(super) fn stream_host_scoring_batches(
input: &PhysicalPlan,
scoring: &hawdb_plan_cypher::HostScoringPlan,
reference_time_millis: u64,
scorer: &mut dyn crate::scoring::HostScorer,
context: BatchReadContext<'_>,
execution_limit: ExecutionLimit,
emit: &mut dyn FnMut(BindingBatch) -> Result<BatchControl>,
) -> Result<BatchControl> {
let mut source = PreparedTransformSource { context };
executor_transform::stream_host_scoring_batches(
input,
executor_transform::HostScoringOptions {
score_column: scoring.score_column(),
max_candidate_rows: scoring.max_candidate_rows(),
reference_time_millis,
limit: scoring.limit(),
expected_identity: Some(crate::scoring::HostScorerDescriptor::new(
scoring.name(),
scoring.version(),
scoring.cpu_units_per_row(),
)?),
rank_policy: scoring.rank_policy(),
},
scorer,
&mut source,
context.kernel_context(),
execution_limit,
emit,
)
}
impl BindingBatchSource for PreparedTransformSource<'_> {
fn execute(
&mut self,
input: &PhysicalPlan,
execution_limit: ExecutionLimit,
emit: &mut dyn FnMut(BindingBatch) -> Result<BatchControl>,
) -> Result<BatchControl> {
execute_prepared_binding_batches(
BatchPlanRef::descendant(input),
self.context,
execution_limit,
emit,
)
}
}
pub(super) fn stream_filter_batches(
predicate: &Predicate,
input: &PhysicalPlan,
context: BatchReadContext<'_>,
execution_limit: ExecutionLimit,
emit: &mut dyn FnMut(BindingBatch) -> Result<BatchControl>,
) -> Result<BatchControl> {
if let PhysicalPlan::SeqNodeScan { variable, label } = input {
let filter = property_filter_from_predicate(predicate).ok();
return stream_node_scan_batches(
variable,
label,
Some((predicate, filter.as_ref())),
context,
execution_limit,
emit,
);
}
if let PhysicalPlan::AdjacencyExpandExec {
rel_variable: Some(rel_variable),
input: expand_input,
..
} = input
&& let Some(filter) = exact_relationship_scan_filter_from_predicate(predicate, rel_variable)
{
return stream_filtered_adjacency_expand_batches(
input,
expand_input,
predicate,
context,
execution_limit,
AdjacencyExpandFilters {
relationship_scan_filter: Some(&filter),
target_scan_filter: None,
},
emit,
);
}
if let PhysicalPlan::AdjacencyExpandExec {
target_variable,
input: expand_input,
..
} = input
&& predicate_references_only_variable(predicate, target_variable)
&& let Ok(filter) = property_filter_from_predicate(predicate)
{
return stream_filtered_adjacency_expand_batches(
input,
expand_input,
predicate,
context,
execution_limit,
AdjacencyExpandFilters {
relationship_scan_filter: None,
target_scan_filter: Some(&filter),
},
emit,
);
}
let predicate_account = context
.kernel_context()
.source_account("FilterExec relationship predicate");
let source_memory = ExecutionMemoryConfig {
batch_rows: if matches!(
input,
PhysicalPlan::NodeProjectionScanExec { .. }
| PhysicalPlan::IndexNodeSeek { .. }
| PhysicalPlan::IndexNodeMultiSeek { .. }
| PhysicalPlan::IndexNodeUnionSeek { .. }
| PhysicalPlan::IndexNodeCompositeSeek { .. }
| PhysicalPlan::IndexNodeCompositeRangeSeek { .. }
| PhysicalPlan::IndexNodeRangeSeek { .. }
| PhysicalPlan::IndexNodeTextSeek { .. }
) {
NonZeroUsize::MIN
} else {
context.memory.batch_rows
},
..context.memory.clone()
};
let mut source = PreparedTransformSource {
context: BatchReadContext {
memory: &source_memory,
..context
},
};
executor_transform::stream_filter_batches(
input,
&mut source,
context.kernel_context(),
execution_limit,
&mut |binding| {
evaluate_predicate_observed(
predicate,
context.catalog,
context.store,
binding,
context.observer,
crate::store::AdjacencyReadMemory {
budget_bytes: context.memory.query_memory_bytes.get(),
account: Some(&predicate_account),
},
context.task_context,
)
},
emit,
)
}
pub(super) fn stream_projection_batches(
items: &[Projection],
input: &PhysicalPlan,
context: BatchReadContext<'_>,
execution_limit: ExecutionLimit,
emit: &mut dyn FnMut(BindingBatch) -> Result<BatchControl>,
) -> Result<BatchControl> {
if context.observer.seed_graph_scoring_input().is_none()
&& let Some(result) =
try_stream_columnar_projection_batches(items, input, context, execution_limit, emit)
{
return result;
}
let mut source = PreparedTransformSource { context };
executor_transform::stream_projection_batches(
items,
input,
&mut source,
context.kernel_context(),
execution_limit,
emit,
)
}
pub(super) fn stream_limit_batches(
offset: usize,
limit: Option<usize>,
input: &PhysicalPlan,
context: BatchReadContext<'_>,
execution_limit: ExecutionLimit,
emit: &mut dyn FnMut(BindingBatch) -> Result<BatchControl>,
) -> Result<BatchControl> {
let mut source = PreparedTransformSource { context };
executor_transform::stream_limit_batches(
offset,
limit,
input,
&mut source,
context.kernel_context(),
execution_limit,
emit,
)
}
pub(super) fn stream_scoring_rerank_batches(
input: &PhysicalPlan,
score_column: &str,
spec: &hawdb_core::graph_rag::ScoringSpec,
limit: usize,
context: BatchReadContext<'_>,
execution_limit: ExecutionLimit,
emit: &mut dyn FnMut(BindingBatch) -> Result<BatchControl>,
) -> Result<BatchControl> {
let mut source = PreparedTransformSource { context };
executor_transform::stream_scoring_rerank_batches(
input,
score_column,
spec,
limit,
&mut source,
context.kernel_context(),
execution_limit,
emit,
)
}
#[allow(clippy::too_many_arguments)]
pub(super) fn stream_scoring_program_batches(
input: &PhysicalPlan,
score_column: &str,
program: &hawdb_core::graph_rag::ScoringProgram,
reference_time_millis: u64,
limit: usize,
context: BatchReadContext<'_>,
execution_limit: ExecutionLimit,
emit: &mut dyn FnMut(BindingBatch) -> Result<BatchControl>,
) -> Result<BatchControl> {
let mut source = PreparedTransformSource { context };
executor_transform::stream_scoring_program_batches(
input,
score_column,
program,
reference_time_millis,
limit,
&mut source,
context.kernel_context(),
execution_limit,
emit,
)
}