query_planning/plan/
graph_algorithm.rs

1//! Native graph algorithm IR.
2//!
3//! This module defines stable graph algorithm names, graph construction
4//! options, edge-list normalization, parameter defaults, and result schemas.
5//! The IR can be built without the optional cuGraph feature; execution support
6//! is selected later by the adapter/runtime.
7use std::sync::Arc;
8
9use arrow::datatypes::DataType as ArrowDataType;
10
11use super::{Expr, Field, NodeId, PlanNode, ProjectionNode, QueryPlan, Schema, SchemaRef};
12
13mod centrality;
14mod community;
15mod definition;
16mod layout;
17mod residency;
18mod similarity;
19mod structure;
20mod traversal;
21
22pub use centrality::{
23    DEFAULT_BETWEENNESS_EXACT_VERTEX_THRESHOLD, GraphBetweennessCentralityParams, GraphBetweennessSeeds,
24    GraphEdgeBetweennessCentralityParams, GraphEigenvectorCentralityParams, GraphHitsParams, GraphKatzCentralityParams,
25    GraphPageRankParams, GraphPageRankPersonalizationInput, GraphPersonalizedPageRankParams,
26};
27pub use community::{
28    GraphEcgParams, GraphLeidenParams, GraphLouvainParams, GraphSccParams, GraphSpectralClusteringParams,
29    GraphWccParams,
30};
31pub use definition::{
32    EmptyGraphEdgeInputAction, EmptyGraphInputPolicy, EmptyGraphSideInputPolicy, GRAPH_OPERATION_REGISTRY,
33    GraphAlgorithm, GraphConstructionOptions, GraphConstructionPolicy, GraphEdgeColumnBindings,
34    GraphOperationDefinition, GraphOutputSchemaKind, GraphSourceVertexRequirement, GraphVertexIdTypeSupport,
35    GraphWeightPolicy,
36};
37pub use layout::{
38    GraphForceAtlas2InitialPositionsInput, GraphForceAtlas2Params, GraphForceAtlas2VertexAttributesInput,
39    GraphLayoutRequirement,
40};
41pub use residency::{
42    GraphResidencyDescriptor, GraphResidencyOrientation, GraphResidencyPlanAnalysis, GraphSourceFingerprint,
43    GraphSourceFingerprintBypassReason, ResidentGraphKey, edge_plan_has_simple_directed_edge_topology,
44    graph_residency_descriptor_for_edge_plan, graph_residency_descriptor_for_local_parquet_source,
45};
46pub use similarity::{GraphSimilarityParams, GraphVertexPairInput};
47pub use structure::{
48    EdgeDedupStrategy, EdgeNormalizeSpec, EdgeWeightCombiner, GraphCoreNumberParams, GraphDegreesAllParams,
49    GraphInDegreesAllParams, GraphKCoreDegreeType, GraphKCoreParams, GraphMstParams, GraphOutDegreesAllParams,
50    GraphTriangleCountAllParams,
51};
52pub use traversal::{
53    GraphBfsEndpointPolicy, GraphBfsParams, GraphBfsPathOutput, GraphBfsPredicateInputs, GraphEdgePredicateInput,
54    GraphPredicateInput, GraphSsspParams, GraphSsspSource,
55};
56
57/// Algorithm-specific parameters.
58#[derive(Debug, Clone, PartialEq)]
59pub enum GraphAlgorithmParams {
60    /// PageRank damping, convergence tolerance, and iteration cap.
61    PageRank(GraphPageRankParams),
62    /// PageRank solver options with a relation-backed personalization vector.
63    PersonalizedPageRank(GraphPersonalizedPageRankParams),
64    /// Breadth-first traversal source, depth bound, predicates, and optional path output.
65    Bfs(Box<GraphBfsParams>),
66    /// Louvain modularity optimization controls.
67    Louvain(GraphLouvainParams),
68    /// Power-iteration tolerance and cap for vertex eigenvector scores.
69    EigenvectorCentrality(GraphEigenvectorCentralityParams),
70    /// Attenuation, constant contribution, and convergence controls for Katz scores.
71    KatzCentrality(GraphKatzCentralityParams),
72    /// Convergence controls and optional normalization for hub and authority scores.
73    Hits(GraphHitsParams),
74    /// Vertex betweenness seed selection, normalization, and endpoint inclusion.
75    BetweennessCentrality(GraphBetweennessCentralityParams),
76    /// Edge betweenness seed selection and normalization.
77    EdgeBetweennessCentrality(GraphEdgeBetweennessCentralityParams),
78    /// Single-source weighted shortest-path source and optional distance cutoff.
79    Sssp(GraphSsspParams),
80    /// Request both incoming and outgoing degree counts for every vertex.
81    DegreesAll(GraphDegreesAllParams),
82    /// Request incoming degree counts for every vertex.
83    InDegreesAll(GraphInDegreesAllParams),
84    /// Request outgoing degree counts for every vertex.
85    OutDegreesAll(GraphOutDegreesAllParams),
86    /// Label vertices connected when edge direction is ignored.
87    WeaklyConnectedComponents(GraphWccParams),
88    /// Label vertices mutually reachable along directed edges.
89    StronglyConnectedComponents(GraphSccParams),
90    /// Count triangles incident to every vertex.
91    TriangleCountAll(GraphTriangleCountAllParams),
92    /// Select the degree convention used to compute vertex core numbers.
93    CoreNumber(GraphCoreNumberParams),
94    /// Select the minimum core degree and degree convention for the returned edge list.
95    KCore(GraphKCoreParams),
96    /// Request a minimum-weight spanning edge list without additional solver options.
97    MinimumSpanningTree(GraphMstParams),
98    /// Leiden community resolution, iteration bound, and random seed.
99    Leiden(GraphLeidenParams),
100    /// Ensemble clustering weights, ensemble size, modularity controls, and random seed.
101    Ecg(GraphEcgParams),
102    /// Spectral clustering count, eigensolver tolerance, and random seed.
103    SpectralModularityMaximization(GraphSpectralClusteringParams),
104    /// ForceAtlas2 iteration controls and optional per-vertex attributes and initial positions.
105    ForceAtlas2(GraphForceAtlas2Params),
106    /// Explicit vertex-pair relation shared by Jaccard, Sorensen, overlap, and cosine metrics.
107    Similarity(GraphSimilarityParams),
108}
109
110impl GraphAlgorithmParams {
111    /// Whether this payload is valid for an algorithm variant.
112    ///
113    /// Similarity deliberately has one shared payload for four metric
114    /// identities. The metric stays in [`GraphAlgorithm`] so the pair side
115    /// input and its validation cannot drift across copied parameter variants.
116    pub const fn supports_algorithm(&self, algorithm: GraphAlgorithm) -> bool {
117        match self {
118            Self::PageRank(_) => matches!(algorithm, GraphAlgorithm::PageRank),
119            Self::PersonalizedPageRank(_) => matches!(algorithm, GraphAlgorithm::PersonalizedPageRank),
120            Self::Bfs(_) => matches!(algorithm, GraphAlgorithm::Bfs),
121            Self::Louvain(_) => matches!(algorithm, GraphAlgorithm::Louvain),
122            Self::EigenvectorCentrality(_) => matches!(algorithm, GraphAlgorithm::EigenvectorCentrality),
123            Self::KatzCentrality(_) => matches!(algorithm, GraphAlgorithm::KatzCentrality),
124            Self::Hits(_) => matches!(algorithm, GraphAlgorithm::Hits),
125            Self::BetweennessCentrality(_) => matches!(algorithm, GraphAlgorithm::BetweennessCentrality),
126            Self::EdgeBetweennessCentrality(_) => matches!(algorithm, GraphAlgorithm::EdgeBetweennessCentrality),
127            Self::Sssp(_) => matches!(algorithm, GraphAlgorithm::Sssp),
128            Self::DegreesAll(_) => matches!(algorithm, GraphAlgorithm::DegreesAll),
129            Self::InDegreesAll(_) => matches!(algorithm, GraphAlgorithm::InDegreesAll),
130            Self::OutDegreesAll(_) => matches!(algorithm, GraphAlgorithm::OutDegreesAll),
131            Self::WeaklyConnectedComponents(_) => matches!(algorithm, GraphAlgorithm::WeaklyConnectedComponents),
132            Self::StronglyConnectedComponents(_) => matches!(algorithm, GraphAlgorithm::StronglyConnectedComponents),
133            Self::TriangleCountAll(_) => matches!(algorithm, GraphAlgorithm::TriangleCountAll),
134            Self::CoreNumber(_) => matches!(algorithm, GraphAlgorithm::CoreNumber),
135            Self::KCore(_) => matches!(algorithm, GraphAlgorithm::KCore),
136            Self::MinimumSpanningTree(_) => matches!(algorithm, GraphAlgorithm::MinimumSpanningTree),
137            Self::Leiden(_) => matches!(algorithm, GraphAlgorithm::Leiden),
138            Self::Ecg(_) => matches!(algorithm, GraphAlgorithm::Ecg),
139            Self::SpectralModularityMaximization(_) => {
140                matches!(algorithm, GraphAlgorithm::SpectralModularityMaximization)
141            },
142            Self::ForceAtlas2(_) => matches!(algorithm, GraphAlgorithm::ForceAtlas2),
143            Self::Similarity(_) => algorithm.is_similarity(),
144        }
145    }
146
147    /// Visits every relation-backed side-input ordinal in stable execution order.
148    pub fn side_input_ordinals(&self) -> Box<dyn Iterator<Item = usize> + '_> {
149        match self {
150            Self::Bfs(params) => Box::new(params.side_input_ordinals()),
151            Self::PersonalizedPageRank(params) => Box::new(params.side_input_ordinals()),
152            Self::Similarity(params) => Box::new(params.side_input_ordinals()),
153            Self::Sssp(params) => Box::new(params.side_input_ordinals()),
154            Self::BetweennessCentrality(params) => Box::new(params.side_input_ordinals()),
155            Self::EdgeBetweennessCentrality(params) => Box::new(params.side_input_ordinals()),
156            Self::ForceAtlas2(params) => Box::new(params.side_input_ordinals()),
157            _ => Box::new(std::iter::empty()),
158        }
159    }
160}
161
162/// Native graph algorithm specification.
163#[derive(Debug, Clone, PartialEq)]
164pub struct GraphAlgorithmSpec {
165    /// Algorithm to execute.
166    pub algorithm: GraphAlgorithm,
167    /// Edge-list column bindings.
168    pub edge_columns: GraphEdgeColumnBindings,
169    /// Shared graph construction options.
170    pub graph_options: GraphConstructionOptions,
171    /// Whether this graph algorithm consumes an explicit edge-normalize node
172    /// inserted from Python-compatible construction semantics.
173    pub edge_normalize_inserted: bool,
174    /// Algorithm-specific parameter payload.
175    pub params: GraphAlgorithmParams,
176    /// Stable output schema expected from execution.
177    pub output_schema: SchemaRef,
178}
179
180impl GraphAlgorithmSpec {
181    /// Create a graph algorithm spec with an explicit output schema.
182    ///
183    /// # Panics
184    ///
185    /// Panics when `params` does not support `algorithm`.
186    pub fn new(
187        algorithm: GraphAlgorithm,
188        edge_columns: GraphEdgeColumnBindings,
189        graph_options: GraphConstructionOptions,
190        params: GraphAlgorithmParams,
191        output_schema: SchemaRef,
192    ) -> Self {
193        assert!(params.supports_algorithm(algorithm), "graph algorithm and parameter variant must match");
194        Self { algorithm, edge_columns, graph_options, edge_normalize_inserted: false, params, output_schema }
195    }
196
197    /// Sets whether this spec consumes an explicit edge-normalize input.
198    pub fn with_edge_normalize_inserted(mut self, inserted: bool) -> Self {
199        self.edge_normalize_inserted = inserted;
200        self
201    }
202
203    /// Create a graph algorithm spec using this crate's stable result schema.
204    ///
205    /// # Panics
206    ///
207    /// Panics when `params` does not support `algorithm`.
208    pub fn with_stable_schema(
209        algorithm: GraphAlgorithm,
210        edge_columns: GraphEdgeColumnBindings,
211        graph_options: GraphConstructionOptions,
212        params: GraphAlgorithmParams,
213    ) -> Self {
214        let output_schema = algorithm.definition().result_schema_for_params(&params);
215        Self::new(algorithm, edge_columns, graph_options, params, output_schema)
216    }
217
218    /// Unique key emitted by this exact algorithm parameterization, when proven.
219    pub fn unique_output_key_columns(&self) -> &'static [&'static str] {
220        if matches!(&self.params, GraphAlgorithmParams::Bfs(params) if params.path_output.is_some()) {
221            return &["path_index"];
222        }
223        if self.algorithm.definition().output_vertex_is_unique() { &["vertex"] } else { &[] }
224    }
225}
226
227/// Stable graph identity available to transient coalescing and resident-cache
228/// lookup, when known at planning time.
229///
230/// Keep this function in sync with graph execution dispatch. It intentionally
231/// returns `None` for paths whose graph construction cannot form a stable
232/// identity, such as edge-id graph construction and BFS include-edges
233/// filtering.
234pub fn planned_resident_graph_key(
235    descriptor: &GraphResidencyDescriptor,
236    spec: &GraphAlgorithmSpec,
237) -> Option<ResidentGraphKey> {
238    if spec.edge_columns.edge_id.is_some() {
239        return None;
240    }
241    match (&spec.algorithm, &spec.params) {
242        (GraphAlgorithm::Bfs, GraphAlgorithmParams::Bfs(params)) if params.predicates.include_edges.is_none() => {
243            Some(descriptor.resident_key(spec, GraphResidencyOrientation::Forward))
244        },
245        (GraphAlgorithm::PageRank, GraphAlgorithmParams::PageRank(_)) => {
246            Some(descriptor.resident_key(spec, GraphResidencyOrientation::Transposed))
247        },
248        (
249            GraphAlgorithm::Jaccard | GraphAlgorithm::Sorensen | GraphAlgorithm::Overlap | GraphAlgorithm::Cosine,
250            GraphAlgorithmParams::Similarity(_),
251        ) => Some(descriptor.resident_key(spec, GraphResidencyOrientation::Forward)),
252        _ => None,
253    }
254}
255
256/// Native graph algorithm operation over one edge-list input.
257#[derive(Debug, Clone, PartialEq)]
258pub struct GraphAlgorithmNode {
259    /// The edge-list input first, then the side inputs; parameter side-input
260    /// bindings name their relation by its position here.
261    pub inputs: Vec<NodeId>,
262    /// Graph algorithm specification.
263    pub spec: GraphAlgorithmSpec,
264}
265
266impl GraphAlgorithmNode {
267    /// The edge-list input node.
268    pub fn edge_input(&self) -> NodeId {
269        self.inputs[0]
270    }
271
272    /// Output schema declared by the graph algorithm spec.
273    pub fn output_schema(&self) -> &SchemaRef {
274        &self.spec.output_schema
275    }
276}
277
278/// Native base result schema for graph algorithms using Int64 vertex ids.
279///
280/// BFS returns raw `vertex`, `distance`, and `predecessor` columns. The SQL
281/// adapter owns the four-column `vertices` result with nullable distances and
282/// `reachable`; native callers retain the underlying distance sentinel.
283/// Use [`graph_result_schema_for_params`] for bounded BFS path output and
284/// [`graph_result_schema_for_vertex_type`] for the input graph's vertex-id dtype.
285pub fn stable_graph_result_schema(algorithm: GraphAlgorithm) -> SchemaRef {
286    let fields = match algorithm {
287        GraphAlgorithm::PageRank
288        | GraphAlgorithm::PersonalizedPageRank
289        | GraphAlgorithm::EigenvectorCentrality
290        | GraphAlgorithm::KatzCentrality
291        | GraphAlgorithm::BetweennessCentrality => {
292            // Centrality-style algorithms expose one score per vertex.
293            vec![Field::new("vertex", ArrowDataType::Int64, false), Field::new("value", ArrowDataType::Float64, false)]
294        },
295        GraphAlgorithm::EdgeBetweennessCentrality => {
296            vec![
297                Field::new("src", ArrowDataType::Int64, false),
298                Field::new("dst", ArrowDataType::Int64, false),
299                Field::new("value", ArrowDataType::Float64, false),
300            ]
301        },
302        GraphAlgorithm::Bfs => vec![
303            Field::new("vertex", ArrowDataType::Int64, false),
304            Field::new("distance", ArrowDataType::Int64, false),
305            Field::new("predecessor", ArrowDataType::Int64, true),
306        ],
307        GraphAlgorithm::Louvain => {
308            vec![
309                Field::new("vertex", ArrowDataType::Int64, false),
310                Field::new("community", ArrowDataType::Int64, false),
311            ]
312        },
313        GraphAlgorithm::Hits => {
314            // HITS exposes both hub and authority scores per vertex;
315            // upstream cuGraph names them `hub` / `authority`. The
316            // adapter renames to `hub_score` / `authority_score` in
317            // the stable SQL schema so the meaning is obvious to query
318            // authors.
319            vec![
320                Field::new("vertex", ArrowDataType::Int64, false),
321                Field::new("hub_score", ArrowDataType::Float64, false),
322                Field::new("authority_score", ArrowDataType::Float64, false),
323            ]
324        },
325        GraphAlgorithm::Sssp => vec![
326            Field::new("vertex", ArrowDataType::Int64, false),
327            // SSSP distances are weighted (Float64); BFS uses Int64 hop counts.
328            Field::new("distance", ArrowDataType::Float64, false),
329            // -1 sentinel means no predecessor (source vertex or unreachable);
330            // nullable=true matches the BFS predecessor convention.
331            Field::new("predecessor", ArrowDataType::Int64, true),
332        ],
333        GraphAlgorithm::DegreesAll => vec![
334            Field::new("vertex", ArrowDataType::Int64, false),
335            Field::new("in_degree", ArrowDataType::Int64, false),
336            Field::new("out_degree", ArrowDataType::Int64, false),
337        ],
338        GraphAlgorithm::InDegreesAll => vec![
339            Field::new("vertex", ArrowDataType::Int64, false),
340            Field::new("in_degree", ArrowDataType::Int64, false),
341        ],
342        GraphAlgorithm::OutDegreesAll => vec![
343            Field::new("vertex", ArrowDataType::Int64, false),
344            Field::new("out_degree", ArrowDataType::Int64, false),
345        ],
346        // WCC and SCC both return a vertex-component assignment. Component
347        // ids are not semantically ordered; parity tests use
348        // partition-equivalence comparison.
349        GraphAlgorithm::WeaklyConnectedComponents | GraphAlgorithm::StronglyConnectedComponents => {
350            vec![
351                Field::new("vertex", ArrowDataType::Int64, false),
352                Field::new("component", ArrowDataType::Int64, false),
353            ]
354        },
355        GraphAlgorithm::TriangleCountAll => vec![
356            Field::new("vertex", ArrowDataType::Int64, false),
357            Field::new("triangle_count", ArrowDataType::Int64, false),
358        ],
359        GraphAlgorithm::CoreNumber => vec![
360            Field::new("vertex", ArrowDataType::Int64, false),
361            Field::new("core_number", ArrowDataType::Int64, false),
362        ],
363        // KCore and MST return an edge list; column names follow the
364        // graph-SQL `src`/`dst` convention. The current output contract omits weight.
365        GraphAlgorithm::KCore | GraphAlgorithm::MinimumSpanningTree => {
366            vec![Field::new("src", ArrowDataType::Int64, false), Field::new("dst", ArrowDataType::Int64, false)]
367        },
368        // Leiden, ECG, and spectral modularity maximization return a
369        // vertex-community assignment. Community ids are arbitrary
370        // identifiers; parity tests use partition-equivalence.
371        GraphAlgorithm::Leiden | GraphAlgorithm::Ecg | GraphAlgorithm::SpectralModularityMaximization => {
372            vec![
373                Field::new("vertex", ArrowDataType::Int64, false),
374                Field::new("community", ArrowDataType::Int64, false),
375            ]
376        },
377        GraphAlgorithm::ForceAtlas2 => vec![
378            Field::new("vertex", ArrowDataType::Int64, false),
379            Field::new("x", ArrowDataType::Float32, false),
380            Field::new("y", ArrowDataType::Float32, false),
381        ],
382        GraphAlgorithm::Jaccard | GraphAlgorithm::Sorensen | GraphAlgorithm::Overlap | GraphAlgorithm::Cosine => vec![
383            Field::new("first", ArrowDataType::Int64, false),
384            Field::new("second", ArrowDataType::Int64, false),
385            Field::new("similarity", ArrowDataType::Float64, false),
386        ],
387    };
388    Arc::new(Schema::new(fields))
389}
390
391/// Native result schema after algorithm-specific output options.
392///
393/// Most graph algorithms use [`stable_graph_result_schema`] without change. BFS
394/// returns the raw three-column vertex result unless bounded path output is requested.
395pub fn graph_result_schema_for_params(algorithm: GraphAlgorithm, params: &GraphAlgorithmParams) -> SchemaRef {
396    algorithm.definition().result_schema_for_params(params)
397}
398
399/// Native result schema in the input graph's numeric vertex-id domain.
400pub fn graph_result_schema_for_vertex_type(
401    algorithm: GraphAlgorithm,
402    params: &GraphAlgorithmParams,
403    vertex_type: &ArrowDataType,
404) -> SchemaRef {
405    assert!(
406        matches!(vertex_type, ArrowDataType::Int32 | ArrowDataType::Int64),
407        "native graph vertex type must be Int32 or Int64"
408    );
409    let vertex_columns = algorithm.definition().vertex_output_columns();
410    let fields = graph_result_schema_for_params(algorithm, params)
411        .fields()
412        .iter()
413        .map(|field| {
414            if vertex_columns.contains(&field.name().as_str()) {
415                field.as_ref().clone().with_data_type(vertex_type.clone())
416            } else {
417                field.as_ref().clone()
418            }
419        })
420        .collect::<Vec<_>>();
421    Arc::new(Schema::new(fields))
422}
query_planning/plan/graph_algorithm.rs

query_planning/plan/
graph_algorithm.rs
