use std::fmt;
use arrow::datatypes::DataType as ArrowDataType;
use super::{
EdgeDedupStrategy, EdgeNormalizeSpec, EdgeWeightCombiner, Expr, Field, GraphAlgorithmSpec, NodeId, PlanNode,
ProjectionNode, QueryPlan,
};
use crate::{
error::{QueryError, QueryErrorCode, Result},
source::{LocalParquetFile, LocalParquetSource, ParquetSourceLocation, SourceColumnRef, SourceNode},
util::{StableFnv1a64, escape_kv_delimited},
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum GraphResidencyOrientation {
Forward,
Transposed,
}
impl GraphResidencyOrientation {
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::Forward => "forward",
Self::Transposed => "transposed",
}
}
}
impl fmt::Display for GraphResidencyOrientation {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct GraphSourceFingerprint {
fingerprint: String,
diagnostic_fingerprint: String,
}
impl fmt::Debug for GraphSourceFingerprint {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("GraphSourceFingerprint").field("fingerprint", &self.diagnostic_fingerprint).finish()
}
}
impl GraphSourceFingerprint {
fn new(fingerprint: String) -> Self {
let mut hasher = StableFnv1a64::new();
hasher.update_field(&fingerprint);
let diagnostic_fingerprint = format!("graph-source-v1-fnv64-{:016x}", hasher.finish());
Self { fingerprint, diagnostic_fingerprint }
}
pub fn as_str(&self) -> &str {
&self.diagnostic_fingerprint
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum GraphSourceFingerprintBypassReason {
NonLocalParquetSource,
UnsupportedTransformedInput,
MissingEdgeColumn,
EdgeVertexDataTypeMismatch,
UnstableSourceIdentity,
}
impl GraphSourceFingerprintBypassReason {
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::NonLocalParquetSource => "non_local_parquet_source",
Self::UnsupportedTransformedInput => "unsupported_transformed_input",
Self::MissingEdgeColumn => "missing_edge_column",
Self::EdgeVertexDataTypeMismatch => "edge_vertex_data_type_mismatch",
Self::UnstableSourceIdentity => "unstable_source_identity",
}
}
}
impl fmt::Display for GraphSourceFingerprintBypassReason {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GraphResidencyPlanAnalysis {
Keyable(GraphResidencyDescriptor),
Bypass(GraphSourceFingerprintBypassReason),
}
impl GraphResidencyPlanAnalysis {
pub fn descriptor(&self) -> Option<&GraphResidencyDescriptor> {
match self {
Self::Keyable(descriptor) => Some(descriptor),
Self::Bypass(_) => None,
}
}
pub fn bypass_reason(&self) -> Option<GraphSourceFingerprintBypassReason> {
match self {
Self::Keyable(_) => None,
Self::Bypass(reason) => Some(*reason),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct GraphResidencyDescriptor {
pub source_fingerprint: GraphSourceFingerprint,
graph_vertex_upper_bound: Option<u64>,
graph_raw_edge_row_count: Option<u64>,
pub vertex_data_type: ArrowDataType,
pub weight_data_type: Option<ArrowDataType>,
pub edge_id_data_type: Option<ArrowDataType>,
}
impl GraphResidencyDescriptor {
pub const fn graph_vertex_upper_bound(&self) -> Option<u64> {
self.graph_vertex_upper_bound
}
pub const fn graph_raw_edge_row_count(&self) -> Option<u64> {
self.graph_raw_edge_row_count
}
pub fn for_content_identity(
content_identity: impl AsRef<str>,
vertex_data_type: ArrowDataType,
weight_data_type: Option<ArrowDataType>,
edge_id_data_type: Option<ArrowDataType>,
) -> Result<Self> {
let content_identity = content_identity.as_ref().trim();
if content_identity.is_empty() {
return Err(QueryError::new(
QueryErrorCode::InvalidPlan,
"query_engine.plan.graph_residency_content_identity",
"graph residency requires a non-empty immutable content identity",
)
.with_fact("reason", "missing_content_identity"));
}
Ok(Self {
source_fingerprint: GraphSourceFingerprint::new(format!(
"mode=external_content_identity_graph_edges_v1;content_identity={content_identity}"
)),
graph_vertex_upper_bound: None,
graph_raw_edge_row_count: None,
vertex_data_type,
weight_data_type,
edge_id_data_type,
})
}
pub fn resident_key(&self, spec: &GraphAlgorithmSpec, orientation: GraphResidencyOrientation) -> ResidentGraphKey {
ResidentGraphKey {
source_fingerprint: self.source_fingerprint.clone(),
source_column: spec.edge_columns.source.clone(),
destination_column: spec.edge_columns.destination.clone(),
weight_column: spec.algorithm.consumes_weight().then(|| spec.edge_columns.weight.clone()).flatten(),
edge_id_column: spec.edge_columns.edge_id.clone(),
directed: spec.graph_options.directed,
renumber: spec.graph_options.renumber,
edge_preparation: GraphEdgePreparationKey::for_spec(spec),
orientation,
vertex_data_type: self.vertex_data_type.clone(),
weight_data_type: self.weight_data_type.clone(),
edge_id_data_type: self.edge_id_data_type.clone(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ResidentGraphKey {
pub(crate) source_fingerprint: GraphSourceFingerprint,
pub(crate) source_column: String,
pub(crate) destination_column: String,
pub weight_column: Option<String>,
pub(crate) edge_id_column: Option<String>,
pub(crate) directed: bool,
pub(crate) renumber: bool,
pub(crate) edge_preparation: GraphEdgePreparationKey,
pub orientation: GraphResidencyOrientation,
pub(crate) vertex_data_type: ArrowDataType,
pub(crate) weight_data_type: Option<ArrowDataType>,
pub(crate) edge_id_data_type: Option<ArrowDataType>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) enum GraphEdgePreparationKey {
None,
PythonCompatible {
add_reverse: bool,
dedup: EdgeDedupStrategy,
weight_combiner: EdgeWeightCombiner,
source_column: String,
destination_column: String,
weight_column: Option<String>,
edge_id_column: Option<String>,
},
}
impl GraphEdgePreparationKey {
fn for_spec(spec: &GraphAlgorithmSpec) -> Self {
if !spec.edge_normalize_inserted {
return Self::None;
}
let normalize = EdgeNormalizeSpec::python_cugraph(spec.edge_columns.clone(), &spec.graph_options);
Self::PythonCompatible {
add_reverse: normalize.add_reverse,
dedup: normalize.dedup,
weight_combiner: normalize.weight_combiner,
source_column: normalize.edge_columns.source,
destination_column: normalize.edge_columns.destination,
weight_column: normalize.edge_columns.weight,
edge_id_column: normalize.edge_columns.edge_id,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct EffectiveGraphEdgeBindings {
source: String,
destination: String,
weight: Option<String>,
edge_id: Option<String>,
}
impl EffectiveGraphEdgeBindings {
fn for_spec(spec: &GraphAlgorithmSpec) -> Self {
let weight_is_semantic = spec.algorithm.consumes_weight() || spec.edge_normalize_inserted;
Self {
source: spec.edge_columns.source.clone(),
destination: spec.edge_columns.destination.clone(),
weight: weight_is_semantic.then(|| spec.edge_columns.weight.clone()).flatten(),
edge_id: spec.edge_columns.edge_id.clone(),
}
}
fn columns(&self) -> impl Iterator<Item = (&'static str, &str)> {
std::iter::once(("source", self.source.as_str()))
.chain(std::iter::once(("destination", self.destination.as_str())))
.chain(self.weight.as_deref().map(|column| ("weight", column)))
.chain(self.edge_id.as_deref().map(|column| ("edge_id", column)))
}
fn try_map_columns(
&self,
mut map: impl FnMut(&str) -> Option<String>,
) -> std::result::Result<Self, GraphSourceFingerprintBypassReason> {
let weight = match self.weight.as_deref() {
Some(column) => Some(map(column).ok_or(GraphSourceFingerprintBypassReason::MissingEdgeColumn)?),
None => None,
};
let edge_id = match self.edge_id.as_deref() {
Some(column) => Some(map(column).ok_or(GraphSourceFingerprintBypassReason::MissingEdgeColumn)?),
None => None,
};
Ok(Self {
source: map(&self.source).ok_or(GraphSourceFingerprintBypassReason::MissingEdgeColumn)?,
destination: map(&self.destination).ok_or(GraphSourceFingerprintBypassReason::MissingEdgeColumn)?,
weight,
edge_id,
})
}
}
pub fn graph_residency_descriptor_for_edge_plan(
plan: &QueryPlan,
spec: &GraphAlgorithmSpec,
) -> GraphResidencyPlanAnalysis {
let bindings = EffectiveGraphEdgeBindings::for_spec(spec);
match graph_residency_descriptor_for_edge_node(plan, plan.root_id(), &bindings) {
Ok(descriptor) => GraphResidencyPlanAnalysis::Keyable(descriptor),
Err(reason) => GraphResidencyPlanAnalysis::Bypass(reason),
}
}
pub fn edge_plan_has_simple_directed_edge_topology(
plan: &QueryPlan,
source_column: &str,
destination_column: &str,
) -> bool {
edge_node_has_simple_directed_edge_topology(plan, plan.root_id(), source_column, destination_column)
}
fn edge_node_has_simple_directed_edge_topology(
plan: &QueryPlan,
node_id: NodeId,
source_column: &str,
destination_column: &str,
) -> bool {
let Some(node) = plan.node(node_id) else {
return false;
};
match &node.kind {
PlanNode::Source(source) => match source.as_ref() {
SourceNode::LocalParquet(source) => {
source.has_simple_directed_edge_topology(source_column, destination_column)
},
SourceNode::GpuResident(_) | SourceNode::LiteralBatch(_) => false,
},
PlanNode::Projection(projection) => {
let (Some(source_column), Some(destination_column)) = (
projection_source_column(projection, source_column),
projection_source_column(projection, destination_column),
) else {
return false;
};
edge_node_has_simple_directed_edge_topology(plan, projection.input, &source_column, &destination_column)
},
PlanNode::Filter(_)
| PlanNode::Aggregate(_)
| PlanNode::Window(_)
| PlanNode::Sort(_)
| PlanNode::Limit(_)
| PlanNode::Join(_)
| PlanNode::Union(_)
| PlanNode::EdgeNormalize(_)
| PlanNode::GraphAlgorithm(_)
| PlanNode::VectorAlgorithm(_)
| PlanNode::MlAlgorithm(_) => false,
}
}
pub fn graph_residency_descriptor_for_local_parquet_source(
source: &LocalParquetSource,
spec: &GraphAlgorithmSpec,
) -> GraphResidencyPlanAnalysis {
let bindings = EffectiveGraphEdgeBindings::for_spec(spec);
match graph_residency_descriptor_for_local_parquet_source_with_bindings(source, &bindings) {
Ok(descriptor) => GraphResidencyPlanAnalysis::Keyable(descriptor),
Err(reason) => GraphResidencyPlanAnalysis::Bypass(reason),
}
}
fn graph_residency_descriptor_for_edge_node(
plan: &QueryPlan,
node_id: NodeId,
bindings: &EffectiveGraphEdgeBindings,
) -> std::result::Result<GraphResidencyDescriptor, GraphSourceFingerprintBypassReason> {
let node = plan.node(node_id).ok_or(GraphSourceFingerprintBypassReason::UnsupportedTransformedInput)?;
match &node.kind {
PlanNode::Source(source) => match source.as_ref() {
SourceNode::LocalParquet(source) => {
graph_residency_descriptor_for_local_parquet_source_with_bindings(source, bindings)
},
SourceNode::GpuResident(_) | SourceNode::LiteralBatch(_) => {
Err(GraphSourceFingerprintBypassReason::NonLocalParquetSource)
},
},
PlanNode::Projection(projection) => {
let mapped = bindings.try_map_columns(|column| projection_source_column(projection, column))?;
graph_residency_descriptor_for_edge_node(plan, projection.input, &mapped)
},
PlanNode::Filter(_)
| PlanNode::Aggregate(_)
| PlanNode::Window(_)
| PlanNode::Sort(_)
| PlanNode::Limit(_)
| PlanNode::Join(_)
| PlanNode::Union(_)
| PlanNode::EdgeNormalize(_)
| PlanNode::GraphAlgorithm(_)
| PlanNode::VectorAlgorithm(_)
| PlanNode::MlAlgorithm(_) => Err(GraphSourceFingerprintBypassReason::UnsupportedTransformedInput),
}
}
fn projection_source_column(projection: &ProjectionNode, output_column: &str) -> Option<String> {
let output = projection.outputs.iter().find(|output| output.field.name() == output_column)?;
match &output.expr {
Expr::Column(column) => Some(column.name().to_string()),
Expr::Alias { expr, alias } if alias == output_column => {
let Expr::Column(column) = expr.as_ref() else {
return None;
};
Some(column.name().to_string())
},
_ => None,
}
}
fn graph_residency_descriptor_for_local_parquet_source_with_bindings(
source: &LocalParquetSource,
bindings: &EffectiveGraphEdgeBindings,
) -> std::result::Result<GraphResidencyDescriptor, GraphSourceFingerprintBypassReason> {
if source.scan_limit.is_some() {
return Err(GraphSourceFingerprintBypassReason::UnsupportedTransformedInput);
}
if source.files.iter().any(|file| !local_parquet_file_identity_is_graph_stable(file)) {
return Err(GraphSourceFingerprintBypassReason::UnstableSourceIdentity);
}
let source_field = graph_source_field(source, &bindings.source)?;
let destination_field = graph_source_field(source, &bindings.destination)?;
if source_field.data_type() != destination_field.data_type() {
return Err(GraphSourceFingerprintBypassReason::EdgeVertexDataTypeMismatch);
}
let weight_field = bindings.weight.as_deref().map(|column| graph_source_field(source, column)).transpose()?;
let edge_id_field = bindings.edge_id.as_deref().map(|column| graph_source_field(source, column)).transpose()?;
let mut fingerprint = String::new();
push_graph_fingerprint_token(&mut fingerprint, "mode", "local_parquet_graph_edges_v1");
push_graph_fingerprint_token(
&mut fingerprint,
"reader_predicate",
source
.reader_predicate
.as_ref()
.map_or_else(|| "none".to_string(), |predicate| format!("display={predicate};debug={predicate:?}")),
);
push_graph_fingerprint_token(
&mut fingerprint,
"scan_limit",
source.scan_limit.map_or_else(|| "none".to_string(), |limit| limit.to_string()),
);
push_graph_fingerprint_token(&mut fingerprint, "file_count", source.files.len().to_string());
for (index, file) in source.files.iter().enumerate() {
push_graph_fingerprint_token(&mut fingerprint, format!("file.{index}.stable_id"), &file.stable_id);
push_graph_fingerprint_token(
&mut fingerprint,
format!("file.{index}.content_identity"),
file.resolved_source.content_identity.as_deref().unwrap_or("none"),
);
push_graph_fingerprint_token(
&mut fingerprint,
format!("file.{index}.size"),
file.file_size_bytes.map_or_else(|| "unknown".to_string(), |bytes| bytes.to_string()),
);
push_graph_optional_usize_vec(&mut fingerprint, format!("file.{index}.row_groups"), file.row_groups.as_deref());
let partition_values = file.partition_values.iter().map(ToString::to_string).collect::<Vec<_>>();
push_graph_string_vec(&mut fingerprint, format!("file.{index}.partition_values"), &partition_values);
}
for (role, column) in bindings.columns() {
push_graph_source_column(&mut fingerprint, role, source, column)?;
}
Ok(GraphResidencyDescriptor {
source_fingerprint: GraphSourceFingerprint::new(fingerprint),
graph_vertex_upper_bound: source.graph_vertex_upper_bound(),
graph_raw_edge_row_count: source.graph_raw_edge_row_count(),
vertex_data_type: source_field.data_type().clone(),
weight_data_type: weight_field.map(|field| field.data_type().clone()),
edge_id_data_type: edge_id_field.map(|field| field.data_type().clone()),
})
}
fn graph_source_field<'a>(
source: &'a LocalParquetSource,
column: &str,
) -> std::result::Result<&'a Field, GraphSourceFingerprintBypassReason> {
source.output_schema().field_with_name(column).map_err(|_| GraphSourceFingerprintBypassReason::MissingEdgeColumn)
}
fn local_parquet_file_identity_is_graph_stable(file: &LocalParquetFile) -> bool {
match &file.location {
ParquetSourceLocation::LocalPath(_) => {
file.file_size_bytes.is_some() && file.stable_id.starts_with("local:") && file.stable_id.contains(":mtime=")
},
ParquetSourceLocation::ObjectStoreUri { .. } => {
file.resolved_source.content_identity.as_deref().is_some_and(graph_content_identity_is_strong)
},
}
}
fn graph_content_identity_is_strong(identity: &str) -> bool {
let identity = identity.trim();
!identity.is_empty()
&& !identity.starts_with("weak_")
&& !identity.get(..2).is_some_and(|prefix| prefix.eq_ignore_ascii_case("w/"))
}
fn push_graph_source_column(
output: &mut String,
role: &str,
source: &LocalParquetSource,
column: &str,
) -> std::result::Result<(), GraphSourceFingerprintBypassReason> {
let field = graph_source_field(source, column)?;
let source_ref = source.source_column_ref(column).ok_or(GraphSourceFingerprintBypassReason::MissingEdgeColumn)?;
push_graph_fingerprint_token(output, format!("edge.{role}.column"), column);
push_graph_fingerprint_token(output, format!("edge.{role}.source_ref"), graph_source_ref_label(source_ref));
push_graph_fingerprint_token(output, format!("edge.{role}.data_type"), format!("{:?}", field.data_type()));
push_graph_fingerprint_token(output, format!("edge.{role}.nullable"), field.is_nullable().to_string());
Ok(())
}
fn graph_source_ref_label(source_ref: SourceColumnRef) -> String {
match source_ref {
SourceColumnRef::File(index) => format!("file:{index}"),
SourceColumnRef::Partition(index) => format!("partition:{index}"),
SourceColumnRef::MissingNull(index) => format!("missing_null:{index}"),
}
}
fn push_graph_optional_usize_vec(output: &mut String, key: impl AsRef<str>, values: Option<&[usize]>) {
match values {
Some(values) => {
push_graph_fingerprint_token(output, format!("{}.len", key.as_ref()), values.len().to_string());
for (index, value) in values.iter().enumerate() {
push_graph_fingerprint_token(output, format!("{}.{}", key.as_ref(), index), value.to_string());
}
},
None => push_graph_fingerprint_token(output, key, "all"),
}
}
fn push_graph_string_vec(output: &mut String, key: impl AsRef<str>, values: &[String]) {
push_graph_fingerprint_token(output, format!("{}.len", key.as_ref()), values.len().to_string());
for (index, value) in values.iter().enumerate() {
push_graph_fingerprint_token(output, format!("{}.{}", key.as_ref(), index), value);
}
}
fn push_graph_fingerprint_token(output: &mut String, key: impl AsRef<str>, value: impl ToString) {
output.push_str(&escape_kv_delimited(key.as_ref()));
output.push('=');
output.push_str(&escape_kv_delimited(value.to_string().as_str()));
output.push(';');
}