#![cfg_attr(not(feature = "gql"), allow(dead_code))]
use std::sync::Arc;
use futures::StreamExt;
use tracing::debug;
use crate::exec::operators::scan::fetch::{FetchFieldStateCache, resolve_with_field_state};
use crate::exec::operators::scan::graph_keys::{
EdgeTableSpec, compute_graph_ranges, decode_graph_edge,
};
use crate::exec::{
AccessMode, CardinalityHint, ContextLevel, ControlFlowExt, EvalContext, ExecOperator,
ExecutionContext, FlowResult, OperatorMetrics, PhysicalExpr, ValueBatch, ValueBatchStream,
buffer_stream, monitor_stream,
};
use crate::expr::{ControlFlow, Dir};
use crate::idx::planner::ScanDirection;
use crate::val::{Object, RecordId, TableName, Value};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExpandDir {
Out,
In,
}
impl ExpandDir {
fn as_dir(self) -> Dir {
match self {
ExpandDir::Out => Dir::Out,
ExpandDir::In => Dir::In,
}
}
fn target_field(self) -> &'static str {
match self {
ExpandDir::Out => "out",
ExpandDir::In => "in",
}
}
fn arrow(self) -> &'static str {
match self {
ExpandDir::Out => "->",
ExpandDir::In => "<-",
}
}
}
#[derive(Debug, Clone)]
pub enum EdgeBinding {
Full(String),
IdOnly(String),
}
impl EdgeBinding {
fn name(&self) -> &str {
match self {
EdgeBinding::Full(n) | EdgeBinding::IdOnly(n) => n,
}
}
fn fetch_full(&self) -> bool {
matches!(self, EdgeBinding::Full(_))
}
}
#[derive(Debug, Clone)]
pub struct Expand {
pub(crate) input: Arc<dyn ExecOperator>,
pub(crate) source: String,
pub(crate) direction: ExpandDir,
pub(crate) edge_tables: Vec<TableName>,
pub(crate) edge_binding: EdgeBinding,
pub(crate) target_binding: String,
pub(crate) target_label: Option<TableName>,
pub(crate) predicate: Option<Arc<dyn PhysicalExpr>>,
pub(crate) optional: bool,
pub(crate) metrics: Arc<OperatorMetrics>,
}
impl Expand {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
input: Arc<dyn ExecOperator>,
source: String,
direction: ExpandDir,
edge_tables: Vec<TableName>,
edge_binding: EdgeBinding,
target_binding: String,
target_label: Option<TableName>,
predicate: Option<Arc<dyn PhysicalExpr>>,
optional: bool,
) -> Self {
Self {
input,
source,
direction,
edge_tables,
edge_binding,
target_binding,
target_label,
predicate,
optional,
metrics: Arc::new(OperatorMetrics::new()),
}
}
fn edge_specs(&self) -> Vec<EdgeTableSpec> {
self.edge_tables
.iter()
.cloned()
.map(|table| EdgeTableSpec {
table,
range_start: std::ops::Bound::Unbounded,
range_end: std::ops::Bound::Unbounded,
})
.collect()
}
}
impl ExecOperator for Expand {
fn name(&self) -> &'static str {
if self.optional {
"OptionalExpand"
} else {
"Expand"
}
}
fn attrs(&self) -> Vec<(String, String)> {
let tables = if self.edge_tables.is_empty() {
"*".to_string()
} else {
self.edge_tables.iter().map(|t| t.as_str()).collect::<Vec<_>>().join(", ")
};
let mut attrs = vec![
("source".to_string(), self.source.clone()),
("direction".to_string(), self.direction.arrow().to_string()),
("tables".to_string(), tables),
("edge".to_string(), self.edge_binding.name().to_string()),
("target".to_string(), self.target_binding.clone()),
];
if let Some(label) = self.target_label.as_ref() {
attrs.push(("target_label".to_string(), label.as_str().to_string()));
}
if let Some(predicate) = self.predicate.as_ref() {
attrs.push(("predicate".to_string(), predicate.to_sql()));
}
attrs
}
fn required_context(&self) -> ContextLevel {
let mut level = self.input.required_context().max(ContextLevel::Database);
if let Some(predicate) = self.predicate.as_ref() {
level = level.max(predicate.required_context());
}
level
}
fn access_mode(&self) -> AccessMode {
let mut mode = self.input.access_mode();
if let Some(predicate) = self.predicate.as_ref() {
mode = mode.combine(predicate.access_mode());
}
mode
}
fn cardinality_hint(&self) -> CardinalityHint {
CardinalityHint::Unbounded
}
fn children(&self) -> Vec<&Arc<dyn ExecOperator>> {
vec![&self.input]
}
fn metrics(&self) -> Option<&OperatorMetrics> {
Some(&self.metrics)
}
fn expressions(&self) -> Vec<(&str, &Arc<dyn PhysicalExpr>)> {
match self.predicate.as_ref() {
Some(predicate) => vec![("predicate", predicate)],
None => vec![],
}
}
fn execute(&self, ctx: &ExecutionContext) -> FlowResult<ValueBatchStream> {
let db_ctx = ctx.database()?.clone();
let input_stream = buffer_stream(
self.input.execute(ctx)?,
self.input.access_mode(),
self.input.cardinality_hint(),
ctx.root().ctx.config.operator_buffer_size,
);
let source = self.source.clone();
let direction = self.direction;
let edge_specs = self.edge_specs();
let edge_binding = self.edge_binding.clone();
let target_binding = self.target_binding.clone();
let target_label = self.target_label.clone();
let predicate = self.predicate.clone();
let optional = self.optional;
let scan_batch_size = ctx.root().ctx.config.scan_batch_size;
let max_output_rows = ctx.root().ctx.config.gql_max_output_rows;
let ctx = ctx.clone();
let stream = async_stream::try_stream! {
let txn = ctx.txn();
let ns_id = db_ctx.ns_ctx.ns.namespace_id;
let db_id = db_ctx.db.database_id;
let version = ctx.version_stamp();
let dir = direction.as_dir();
let fetch_full_edge = edge_binding.fetch_full();
let edge_name = edge_binding.name().to_string();
let mut edge_cache = FetchFieldStateCache::new();
let mut target_cache = FetchFieldStateCache::new();
let mut emitted: usize = 0;
futures::pin_mut!(input_stream);
while let Some(batch_result) = input_stream.next().await {
crate::exec::operators::check_cancelled(&ctx)?;
let batch = batch_result?;
let mut out: Vec<Value> = Vec::with_capacity(batch.values.len());
for row in batch.values {
match extract_source_id(&row, &source) {
SourceId::Found(rid) => {
expand_row(
&ctx,
&txn,
ns_id,
db_id,
version,
&row,
&rid,
dir,
direction,
&edge_specs,
&edge_name,
fetch_full_edge,
&target_binding,
target_label.as_ref(),
predicate.as_deref(),
optional,
scan_batch_size,
&mut edge_cache,
&mut target_cache,
&mut out,
&mut emitted,
max_output_rows,
)
.await?;
}
SourceId::OptionalMiss => {
if optional {
out.push(null_filled_row(&row, &edge_name, &target_binding));
}
}
SourceId::Drop => {
debug!(
binding = %source,
"Expand: source binding is missing or not a record id; dropping row"
);
}
}
}
if !out.is_empty() {
yield ValueBatch { values: out };
}
}
};
Ok(monitor_stream(Box::pin(stream), self.name(), &self.metrics))
}
}
enum SourceId {
Found(RecordId),
OptionalMiss,
Drop,
}
fn extract_source_id(row: &Value, source: &str) -> SourceId {
let Value::Object(obj) = row else {
return SourceId::Drop;
};
match obj.get(source) {
Some(Value::Object(node)) => match node.get("id") {
Some(Value::RecordId(rid)) => SourceId::Found(rid.clone()),
_ => SourceId::Drop,
},
Some(Value::RecordId(rid)) => SourceId::Found(rid.clone()),
Some(Value::Null) => SourceId::OptionalMiss,
_ => SourceId::Drop,
}
}
fn candidate_row(
input: &Value,
edge_name: &str,
edge: Value,
target_binding: &str,
target: Value,
) -> Value {
let mut obj = match input {
Value::Object(o) => o.clone(),
_ => Object::default(),
};
obj.insert(edge_name, edge);
obj.insert(target_binding, target);
Value::Object(obj)
}
fn null_filled_row(input: &Value, edge_name: &str, target_binding: &str) -> Value {
candidate_row(input, edge_name, Value::Null, target_binding, Value::Null)
}
fn target_from_edge_obj(edge: &Value, direction: ExpandDir) -> Option<RecordId> {
let Value::Object(obj) = edge else {
return None;
};
match obj.get(direction.target_field()) {
Some(Value::RecordId(rid)) => Some(rid.clone()),
_ => None,
}
}
#[allow(clippy::too_many_arguments)]
async fn expand_row(
ctx: &ExecutionContext,
txn: &Arc<crate::kvs::Transaction>,
ns_id: crate::catalog::NamespaceId,
db_id: crate::catalog::DatabaseId,
version: Option<u64>,
input_row: &Value,
source_rid: &RecordId,
dir: Dir,
direction: ExpandDir,
edge_specs: &[EdgeTableSpec],
edge_name: &str,
fetch_full_edge: bool,
target_binding: &str,
target_label: Option<&TableName>,
predicate: Option<&dyn PhysicalExpr>,
optional: bool,
scan_batch_size: usize,
edge_cache: &mut FetchFieldStateCache,
target_cache: &mut FetchFieldStateCache,
out: &mut Vec<Value>,
emitted: &mut usize,
max_output_rows: usize,
) -> Result<(), ControlFlow> {
let ranges = compute_graph_ranges(ns_id, db_id, source_rid, dir, edge_specs, ctx).await?;
let mut edge_rids: Vec<RecordId> = Vec::new();
let mut decoded_targets: Vec<Option<RecordId>> = Vec::new();
for (beg, end) in ranges {
let mut cursor = txn
.open_keys_cursor(beg..end, ScanDirection::Forward, 0, version)
.await
.context("Failed to open graph cursor")?;
loop {
crate::exec::operators::check_cancelled(ctx)?;
let keys = cursor
.next_batch(crate::kvs::NORMAL_BATCH_SIZE)
.await
.context("Failed to scan graph edge")?;
if keys.is_empty() {
break;
}
for key in &keys {
let decoded = decode_graph_edge(key)?;
edge_rids.push(decoded.edge);
decoded_targets.push(decoded.target);
}
}
drop(cursor);
}
if edge_rids.is_empty() {
if optional {
out.push(null_filled_row(input_row, edge_name, target_binding));
}
return Ok(());
}
let gate_id_only_edges = !fetch_full_edge && edge_cache.check_perms(ctx)?;
let before = out.len();
let mut start = 0;
while start < edge_rids.len() {
let end = (start + scan_batch_size).min(edge_rids.len());
let edge_slice = &edge_rids[start..end];
let decoded_slice = &decoded_targets[start..end];
let (edge_values, edge_perm): (Vec<Option<Value>>, Vec<Option<Value>>) = if fetch_full_edge
{
let fetched = resolve_with_field_state(ctx, edge_cache, edge_slice).await?;
(fetched.clone(), fetched)
} else if gate_id_only_edges {
let gated = resolve_with_field_state(ctx, edge_cache, edge_slice).await?;
let values = edge_slice.iter().map(|rid| Some(Value::RecordId(rid.clone()))).collect();
(values, gated)
} else {
let values: Vec<Option<Value>> =
edge_slice.iter().map(|rid| Some(Value::RecordId(rid.clone()))).collect();
let perm = vec![None; edge_slice.len()];
(values, perm)
};
let mut target_rids: Vec<Option<RecordId>> = Vec::with_capacity(edge_slice.len());
let mut legacy_edge_fetch: Vec<RecordId> = Vec::new();
let mut legacy_positions: Vec<usize> = Vec::new();
for (i, edge_rid) in edge_slice.iter().enumerate() {
let edge_denied = (fetch_full_edge && edge_values[i].is_none())
|| (gate_id_only_edges && edge_perm[i].is_none());
if edge_denied {
target_rids.push(None);
} else if let Some(target) = decoded_slice[i].clone() {
target_rids.push(Some(target));
} else if let Some(edge_val) = edge_perm[i].as_ref() {
target_rids.push(target_from_edge_obj(edge_val, direction));
} else {
target_rids.push(None);
legacy_edge_fetch.push(edge_rid.clone());
legacy_positions.push(i);
}
}
if !legacy_edge_fetch.is_empty() {
let fetched = resolve_with_field_state(ctx, edge_cache, &legacy_edge_fetch).await?;
for (pos, edge_val) in legacy_positions.iter().zip(fetched) {
if let Some(edge_val) = edge_val {
target_rids[*pos] = target_from_edge_obj(&edge_val, direction);
}
}
}
let mut targets_to_fetch: Vec<RecordId> = Vec::new();
let mut candidate_positions: Vec<usize> = Vec::new();
for (i, target) in target_rids.iter().enumerate() {
let Some(target) = target else {
continue;
};
if let Some(label) = target_label
&& &target.table != label
{
continue;
}
targets_to_fetch.push(target.clone());
candidate_positions.push(i);
}
let target_values = resolve_with_field_state(ctx, target_cache, &targets_to_fetch).await?;
let base_eval = EvalContext::from_exec_ctx(ctx);
for (slot, target_value) in target_values.into_iter().enumerate() {
let Some(target_value) = target_value else {
continue;
};
let i = candidate_positions[slot];
let edge_value = match edge_values[i].clone() {
Some(v) => v,
None => continue,
};
let candidate =
candidate_row(input_row, edge_name, edge_value, target_binding, target_value);
if let Some(predicate) = predicate {
let result = predicate.evaluate(base_eval.with_value(&candidate)).await?;
if !result.is_truthy() {
continue;
}
}
*emitted += 1;
if *emitted > max_output_rows {
return Err(crate::exec::operators::gql_output_rows_exceeded(max_output_rows));
}
out.push(candidate);
}
start = end;
}
if optional && out.len() == before {
out.push(null_filled_row(input_row, edge_name, target_binding));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::exec::operators::CurrentValueSource;
fn obj_row(pairs: &[(&str, Value)]) -> Value {
let mut o = Object::default();
for (k, v) in pairs {
o.insert(*k, v.clone());
}
Value::Object(o)
}
fn rid(table: &str, key: &str) -> RecordId {
RecordId {
table: TableName::new(table.to_string()),
key: crate::val::RecordIdKey::from(key.to_string()),
}
}
fn sample_expand(optional: bool) -> Expand {
Expand::new(
Arc::new(CurrentValueSource::new()),
"a".to_string(),
ExpandDir::Out,
vec![TableName::new("knows".to_string())],
EdgeBinding::Full("k".to_string()),
"b".to_string(),
Some(TableName::new("person".to_string())),
None,
optional,
)
}
#[test]
fn name_reflects_optional() {
assert_eq!(sample_expand(false).name(), "Expand");
assert_eq!(sample_expand(true).name(), "OptionalExpand");
}
#[test]
fn attrs_render_direction_tables_and_label() {
let attrs = sample_expand(false).attrs();
assert!(attrs.iter().any(|(k, v)| k == "source" && v == "a"));
assert!(attrs.iter().any(|(k, v)| k == "direction" && v == "->"));
assert!(attrs.iter().any(|(k, v)| k == "tables" && v == "knows"));
assert!(attrs.iter().any(|(k, v)| k == "edge" && v == "k"));
assert!(attrs.iter().any(|(k, v)| k == "target" && v == "b"));
assert!(attrs.iter().any(|(k, v)| k == "target_label" && v == "person"));
}
#[test]
fn attrs_render_wildcard_tables() {
let mut e = sample_expand(false);
e.edge_tables.clear();
let attrs = e.attrs();
assert!(attrs.iter().any(|(k, v)| k == "tables" && v == "*"));
}
#[test]
fn required_context_is_at_least_database() {
assert_eq!(sample_expand(false).required_context(), ContextLevel::Database);
}
#[test]
fn edge_specs_unbounded_per_table() {
let e = sample_expand(false);
let specs = e.edge_specs();
assert_eq!(specs.len(), 1);
assert_eq!(specs[0].table.as_str(), "knows");
assert!(matches!(specs[0].range_start, std::ops::Bound::Unbounded));
assert!(matches!(specs[0].range_end, std::ops::Bound::Unbounded));
}
#[test]
fn edge_specs_empty_for_wildcard() {
let mut e = sample_expand(false);
e.edge_tables.clear();
assert!(e.edge_specs().is_empty());
}
#[test]
fn dir_mapping_and_target_field() {
assert_eq!(ExpandDir::Out.as_dir(), Dir::Out);
assert_eq!(ExpandDir::In.as_dir(), Dir::In);
assert_eq!(ExpandDir::Out.target_field(), "out");
assert_eq!(ExpandDir::In.target_field(), "in");
}
#[test]
fn edge_binding_name_and_fetch_flag() {
assert_eq!(EdgeBinding::Full("k".to_string()).name(), "k");
assert_eq!(EdgeBinding::IdOnly("k".to_string()).name(), "k");
assert!(EdgeBinding::Full("k".to_string()).fetch_full());
assert!(!EdgeBinding::IdOnly("k".to_string()).fetch_full());
}
#[test]
fn extract_source_id_from_full_object() {
let row = obj_row(&[("a", obj_row(&[("id", Value::RecordId(rid("person", "alice")))]))]);
match extract_source_id(&row, "a") {
SourceId::Found(r) => assert_eq!(r, rid("person", "alice")),
_ => panic!("expected Found"),
}
}
#[test]
fn extract_source_id_from_bare_record_id() {
let row = obj_row(&[("a", Value::RecordId(rid("person", "bob")))]);
match extract_source_id(&row, "a") {
SourceId::Found(r) => assert_eq!(r, rid("person", "bob")),
_ => panic!("expected Found"),
}
}
#[test]
fn extract_source_id_null_is_optional_miss() {
let row = obj_row(&[("a", Value::Null)]);
assert!(matches!(extract_source_id(&row, "a"), SourceId::OptionalMiss));
}
#[test]
fn extract_source_id_missing_or_bad_is_drop() {
let missing = obj_row(&[("z", Value::Bool(true))]);
assert!(matches!(extract_source_id(&missing, "a"), SourceId::Drop));
let bad = obj_row(&[("a", Value::Bool(true))]);
assert!(matches!(extract_source_id(&bad, "a"), SourceId::Drop));
let no_id = obj_row(&[("a", obj_row(&[("name", Value::Bool(true))]))]);
assert!(matches!(extract_source_id(&no_id, "a"), SourceId::Drop));
assert!(matches!(extract_source_id(&Value::Bool(true), "a"), SourceId::Drop));
}
#[test]
fn candidate_row_inserts_edge_and_target() {
let input = obj_row(&[("a", Value::RecordId(rid("person", "alice")))]);
let row = candidate_row(
&input,
"k",
Value::RecordId(rid("knows", "e1")),
"b",
Value::RecordId(rid("person", "bob")),
);
let Value::Object(o) = row else {
panic!("expected object");
};
assert_eq!(o.get("a"), Some(&Value::RecordId(rid("person", "alice"))));
assert_eq!(o.get("k"), Some(&Value::RecordId(rid("knows", "e1"))));
assert_eq!(o.get("b"), Some(&Value::RecordId(rid("person", "bob"))));
}
#[test]
fn null_filled_row_nulls_edge_and_target() {
let input = obj_row(&[("a", Value::RecordId(rid("person", "alice")))]);
let row = null_filled_row(&input, "k", "b");
let Value::Object(o) = row else {
panic!("expected object");
};
assert_eq!(o.get("a"), Some(&Value::RecordId(rid("person", "alice"))));
assert_eq!(o.get("k"), Some(&Value::Null));
assert_eq!(o.get("b"), Some(&Value::Null));
}
#[test]
fn target_from_edge_obj_reads_correct_field() {
let edge = obj_row(&[
("in", Value::RecordId(rid("person", "alice"))),
("out", Value::RecordId(rid("person", "bob"))),
]);
assert_eq!(target_from_edge_obj(&edge, ExpandDir::Out), Some(rid("person", "bob")));
assert_eq!(target_from_edge_obj(&edge, ExpandDir::In), Some(rid("person", "alice")));
let partial = obj_row(&[("in", Value::RecordId(rid("person", "alice")))]);
assert_eq!(target_from_edge_obj(&partial, ExpandDir::Out), None);
assert_eq!(target_from_edge_obj(&Value::Null, ExpandDir::Out), None);
}
}