algeon-datafusion 0.1.2

use arrow::datatypes::DataType;
use query_planning::plan::{GraphAlgorithm as NativeGraphAlgorithm, GraphVertexIdTypeSupport};

use super::types::CugraphGraphOptions;

/// SQL function name for [`CugraphAlgorithm::PageRank`].
pub const CUGRAPH_PAGERANK_FUNCTION: &str = "cugraph_pagerank";
/// SQL function name for [`CugraphAlgorithm::PersonalizedPageRank`].
pub const CUGRAPH_PERSONALIZED_PAGERANK_FUNCTION: &str = "cugraph_personalized_pagerank";
/// SQL function name for [`CugraphAlgorithm::Bfs`].
pub const CUGRAPH_BFS_FUNCTION: &str = "cugraph_bfs";
/// SQL function name for [`CugraphAlgorithm::Louvain`].
pub const CUGRAPH_LOUVAIN_FUNCTION: &str = "cugraph_louvain";
/// SQL function name for [`CugraphAlgorithm::EigenvectorCentrality`].
pub const CUGRAPH_EIGENVECTOR_CENTRALITY_FUNCTION: &str = "cugraph_eigenvector_centrality";
/// SQL function name for [`CugraphAlgorithm::KatzCentrality`].
pub const CUGRAPH_KATZ_CENTRALITY_FUNCTION: &str = "cugraph_katz_centrality";
/// SQL function name for [`CugraphAlgorithm::Hits`].
pub const CUGRAPH_HITS_FUNCTION: &str = "cugraph_hits";
/// SQL function name for [`CugraphAlgorithm::BetweennessCentrality`].
pub const CUGRAPH_BETWEENNESS_CENTRALITY_FUNCTION: &str = "cugraph_betweenness_centrality";
/// SQL function name for [`CugraphAlgorithm::EdgeBetweennessCentrality`].
pub const CUGRAPH_EDGE_BETWEENNESS_CENTRALITY_FUNCTION: &str = "cugraph_edge_betweenness_centrality";
/// SQL function name for [`CugraphAlgorithm::Sssp`].
pub const CUGRAPH_SSSP_FUNCTION: &str = "cugraph_sssp";
/// SQL function name for [`CugraphAlgorithm::DegreesAll`].
pub const CUGRAPH_DEGREES_ALL_FUNCTION: &str = "cugraph_degrees_all";
/// SQL function name for [`CugraphAlgorithm::InDegreesAll`].
pub const CUGRAPH_IN_DEGREES_ALL_FUNCTION: &str = "cugraph_in_degrees_all";
/// SQL function name for [`CugraphAlgorithm::OutDegreesAll`].
pub const CUGRAPH_OUT_DEGREES_ALL_FUNCTION: &str = "cugraph_out_degrees_all";
/// SQL function name for [`CugraphAlgorithm::WeaklyConnectedComponents`].
pub const CUGRAPH_WEAKLY_CONNECTED_COMPONENTS_FUNCTION: &str = "cugraph_weakly_connected_components";
/// SQL function name for [`CugraphAlgorithm::StronglyConnectedComponents`].
pub const CUGRAPH_STRONGLY_CONNECTED_COMPONENTS_FUNCTION: &str = "cugraph_strongly_connected_components";
/// SQL function name for [`CugraphAlgorithm::TriangleCountAll`].
pub const CUGRAPH_TRIANGLE_COUNT_ALL_FUNCTION: &str = "cugraph_triangle_count_all";
/// SQL function name for [`CugraphAlgorithm::CoreNumber`].
pub const CUGRAPH_CORE_NUMBER_FUNCTION: &str = "cugraph_core_number";
/// SQL function name for [`CugraphAlgorithm::KCore`].
pub const CUGRAPH_K_CORE_FUNCTION: &str = "cugraph_k_core";
/// SQL function name for [`CugraphAlgorithm::MinimumSpanningTree`].
pub const CUGRAPH_MINIMUM_SPANNING_TREE_FUNCTION: &str = "cugraph_minimum_spanning_tree";
/// SQL function name for [`CugraphAlgorithm::Leiden`].
pub const CUGRAPH_LEIDEN_FUNCTION: &str = "cugraph_leiden";
/// SQL function name for [`CugraphAlgorithm::Ecg`].
pub const CUGRAPH_ECG_FUNCTION: &str = "cugraph_ecg";
/// SQL function name for [`CugraphAlgorithm::SpectralModularityMaximization`].
pub const CUGRAPH_SPECTRAL_MODULARITY_MAXIMIZATION_FUNCTION: &str = "cugraph_spectral_modularity_maximization";
/// SQL function name for [`CugraphAlgorithm::ForceAtlas2`].
pub const CUGRAPH_FORCE_ATLAS2_FUNCTION: &str = "cugraph_force_atlas2";
/// SQL function name for [`CugraphAlgorithm::Jaccard`].
pub const CUGRAPH_JACCARD_FUNCTION: &str = "cugraph_jaccard";
/// SQL function name for [`CugraphAlgorithm::Sorensen`].
pub const CUGRAPH_SORENSEN_FUNCTION: &str = "cugraph_sorensen";
/// SQL function name for [`CugraphAlgorithm::Overlap`].
pub const CUGRAPH_OVERLAP_FUNCTION: &str = "cugraph_overlap";
/// SQL function name for [`CugraphAlgorithm::Cosine`].
pub const CUGRAPH_COSINE_FUNCTION: &str = "cugraph_cosine";

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
/// Algorithms registered as relation-aware cuGraph SQL functions.
pub enum CugraphAlgorithm {
    /// Rank vertices by damped incoming-link influence.
    PageRank,
    /// Rank vertices using a caller-supplied personalization relation.
    PersonalizedPageRank,
    /// Traverse unweighted paths from one or more sources.
    Bfs,
    /// Detect communities by hierarchical modularity optimization.
    Louvain,
    /// Compute vertex influence from neighboring centrality scores.
    EigenvectorCentrality,
    /// Compute attenuated neighbor influence with a constant bias.
    KatzCentrality,
    /// Compute hub and authority scores.
    Hits,
    /// Score vertices by their contribution to shortest paths.
    BetweennessCentrality,
    /// Score edges by their contribution to shortest paths.
    EdgeBetweennessCentrality,
    /// Compute weighted shortest distances from one source.
    Sssp,
    /// Return incoming and outgoing vertex degrees.
    DegreesAll,
    /// Return incoming vertex degrees.
    InDegreesAll,
    /// Return outgoing vertex degrees.
    OutDegreesAll,
    /// Label components using connectivity without edge direction.
    WeaklyConnectedComponents,
    /// Label components whose vertices are mutually reachable by directed paths.
    StronglyConnectedComponents,
    /// Count triangles involving each vertex.
    TriangleCountAll,
    /// Return each vertex's largest containing k-core.
    CoreNumber,
    /// Extract the subgraph meeting the requested core-number bound.
    KCore,
    /// Return minimum-weight spanning edges.
    MinimumSpanningTree,
    /// Detect communities with modularity refinement.
    Leiden,
    /// Detect communities using an ensemble of Louvain partitions.
    Ecg,
    /// Cluster vertices from a modularity spectral embedding.
    SpectralModularityMaximization,
    /// Compute two-dimensional vertex layout coordinates.
    ForceAtlas2,
    /// Compare neighbor sets by intersection over union.
    Jaccard,
    /// Compare neighbor sets by twice their intersection over combined sizes.
    Sorensen,
    /// Compare neighbor sets by intersection over the smaller set.
    Overlap,
    /// Compare neighbor vectors using cosine similarity.
    Cosine,
}

impl CugraphAlgorithm {
    /// All algorithms exposed by the SQL surface, in registry order.
    pub fn all() -> impl ExactSizeIterator<Item = Self> + Clone {
        CUGRAPH_OPERATION_REGISTRY.iter().map(|entry| entry.algorithm)
    }

    /// Returns the registered SQL function name for this algorithm.
    pub fn function_name(self) -> &'static str {
        self.definition().sql_function_name
    }

    /// Returns the stable native algorithm code used by diagnostics and reports.
    pub fn code(self) -> &'static str {
        self.native_algorithm().as_str()
    }

    /// Look up the SQL-side registry definition for this algorithm.
    pub fn definition(self) -> &'static CugraphOperationDefinition {
        CUGRAPH_OPERATION_REGISTRY
            .iter()
            .find(|entry| entry.algorithm == self)
            .expect("cuGraph SQL operation registry is missing an algorithm variant")
    }

    pub(crate) fn native_algorithm(self) -> NativeGraphAlgorithm {
        self.definition().native_algorithm
    }

    /// Vertex id type of the native graph this algorithm runs on. Vertex
    /// domains the native graph cannot take are mapped to ids of this type by
    /// the per-call codebook.
    pub(crate) fn native_vertex_id_type(self) -> DataType {
        match self.native_algorithm().vertex_id_type_support() {
            GraphVertexIdTypeSupport::Int32OnlyLegacyEdgeType => DataType::Int32,
            GraphVertexIdTypeSupport::Int32OrInt64 => DataType::Int64,
        }
    }

    pub(crate) const fn is_similarity(self) -> bool {
        matches!(self, Self::Jaccard | Self::Sorensen | Self::Overlap | Self::Cosine)
    }

    pub(crate) fn effective_default_graph_options(self, defaults: &CugraphGraphOptions) -> CugraphGraphOptions {
        CugraphGraphOptions::from_native(
            &self.native_algorithm().definition().default_graph_options(&defaults.to_native()),
        )
    }
}

/// Registry-level metadata for a SQL-visible cuGraph operation.
///
/// One entry per [`CugraphAlgorithm`] variant. Mirrors
/// [`query_planning::plan::GRAPH_OPERATION_REGISTRY`]: each SQL UDTF resolves to
/// the same native algorithm advertised by the engine, and the SQL function
/// name follows the `cugraph_<native_name>` convention.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct CugraphOperationDefinition {
    /// SQL algorithm represented by this registry entry.
    pub algorithm: CugraphAlgorithm,
    native_algorithm: NativeGraphAlgorithm,
    /// Function identifier recognized by the cuGraph SQL binder.
    pub sql_function_name: &'static str,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) enum CugraphLogicalOutputMultiplicity {
    OnePerVertex,
    EdgeEndpointOccurrences,
    SideInputOccurrences,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct CugraphLogicalOutputColumn {
    pub(crate) name: &'static str,
    pub(crate) nullable: bool,
    pub(crate) multiplicity: CugraphLogicalOutputMultiplicity,
}

const ONE_VERTEX_LOGICAL_OUTPUT: &[CugraphLogicalOutputColumn] = &[CugraphLogicalOutputColumn {
    name: "vertex",
    nullable: false,
    multiplicity: CugraphLogicalOutputMultiplicity::OnePerVertex,
}];

const BFS_LOGICAL_OUTPUTS: &[CugraphLogicalOutputColumn] = &[
    CugraphLogicalOutputColumn {
        name: "vertex",
        nullable: false,
        multiplicity: CugraphLogicalOutputMultiplicity::OnePerVertex,
    },
    CugraphLogicalOutputColumn {
        name: "predecessor",
        nullable: true,
        multiplicity: CugraphLogicalOutputMultiplicity::EdgeEndpointOccurrences,
    },
];

const EDGE_LOGICAL_OUTPUTS: &[CugraphLogicalOutputColumn] = &[
    CugraphLogicalOutputColumn {
        name: "src",
        nullable: false,
        multiplicity: CugraphLogicalOutputMultiplicity::EdgeEndpointOccurrences,
    },
    CugraphLogicalOutputColumn {
        name: "dst",
        nullable: false,
        multiplicity: CugraphLogicalOutputMultiplicity::EdgeEndpointOccurrences,
    },
];

const PAIR_LOGICAL_OUTPUTS: &[CugraphLogicalOutputColumn] = &[
    CugraphLogicalOutputColumn {
        name: "first",
        nullable: false,
        multiplicity: CugraphLogicalOutputMultiplicity::SideInputOccurrences,
    },
    CugraphLogicalOutputColumn {
        name: "second",
        nullable: false,
        multiplicity: CugraphLogicalOutputMultiplicity::SideInputOccurrences,
    },
];

impl CugraphOperationDefinition {
    /// Result columns that carry vertex ids and are decoded back to the input
    /// vertex type when the per-call codebook mapped the call.
    pub(crate) const fn vertex_output_columns(&self) -> &'static [CugraphLogicalOutputColumn] {
        match self.algorithm {
            CugraphAlgorithm::Bfs | CugraphAlgorithm::Sssp => BFS_LOGICAL_OUTPUTS,
            CugraphAlgorithm::PageRank
            | CugraphAlgorithm::PersonalizedPageRank
            | CugraphAlgorithm::Louvain
            | CugraphAlgorithm::EigenvectorCentrality
            | CugraphAlgorithm::KatzCentrality
            | CugraphAlgorithm::Hits
            | CugraphAlgorithm::BetweennessCentrality
            | CugraphAlgorithm::DegreesAll
            | CugraphAlgorithm::InDegreesAll
            | CugraphAlgorithm::OutDegreesAll
            | CugraphAlgorithm::WeaklyConnectedComponents
            | CugraphAlgorithm::StronglyConnectedComponents
            | CugraphAlgorithm::TriangleCountAll
            | CugraphAlgorithm::CoreNumber
            | CugraphAlgorithm::Leiden
            | CugraphAlgorithm::Ecg
            | CugraphAlgorithm::SpectralModularityMaximization
            | CugraphAlgorithm::ForceAtlas2 => ONE_VERTEX_LOGICAL_OUTPUT,
            CugraphAlgorithm::KCore
            | CugraphAlgorithm::MinimumSpanningTree
            | CugraphAlgorithm::EdgeBetweennessCentrality => EDGE_LOGICAL_OUTPUTS,
            CugraphAlgorithm::Jaccard
            | CugraphAlgorithm::Sorensen
            | CugraphAlgorithm::Overlap
            | CugraphAlgorithm::Cosine => PAIR_LOGICAL_OUTPUTS,
        }
    }
}

/// Source-of-truth list of SQL-visible cuGraph algorithms.
///
/// Order is preserved by [`CugraphAlgorithm::all`] and is the order in which
/// cuGraph UDTFs are available to the relation planner.
pub const CUGRAPH_OPERATION_REGISTRY: &[CugraphOperationDefinition] = &[
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::PageRank,
        native_algorithm: NativeGraphAlgorithm::PageRank,
        sql_function_name: CUGRAPH_PAGERANK_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::Bfs,
        native_algorithm: NativeGraphAlgorithm::Bfs,
        sql_function_name: CUGRAPH_BFS_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::Louvain,
        native_algorithm: NativeGraphAlgorithm::Louvain,
        sql_function_name: CUGRAPH_LOUVAIN_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::EigenvectorCentrality,
        native_algorithm: NativeGraphAlgorithm::EigenvectorCentrality,
        sql_function_name: CUGRAPH_EIGENVECTOR_CENTRALITY_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::KatzCentrality,
        native_algorithm: NativeGraphAlgorithm::KatzCentrality,
        sql_function_name: CUGRAPH_KATZ_CENTRALITY_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::Hits,
        native_algorithm: NativeGraphAlgorithm::Hits,
        sql_function_name: CUGRAPH_HITS_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::BetweennessCentrality,
        native_algorithm: NativeGraphAlgorithm::BetweennessCentrality,
        sql_function_name: CUGRAPH_BETWEENNESS_CENTRALITY_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::EdgeBetweennessCentrality,
        native_algorithm: NativeGraphAlgorithm::EdgeBetweennessCentrality,
        sql_function_name: CUGRAPH_EDGE_BETWEENNESS_CENTRALITY_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::Sssp,
        native_algorithm: NativeGraphAlgorithm::Sssp,
        sql_function_name: CUGRAPH_SSSP_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::DegreesAll,
        native_algorithm: NativeGraphAlgorithm::DegreesAll,
        sql_function_name: CUGRAPH_DEGREES_ALL_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::InDegreesAll,
        native_algorithm: NativeGraphAlgorithm::InDegreesAll,
        sql_function_name: CUGRAPH_IN_DEGREES_ALL_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::OutDegreesAll,
        native_algorithm: NativeGraphAlgorithm::OutDegreesAll,
        sql_function_name: CUGRAPH_OUT_DEGREES_ALL_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::WeaklyConnectedComponents,
        native_algorithm: NativeGraphAlgorithm::WeaklyConnectedComponents,
        sql_function_name: CUGRAPH_WEAKLY_CONNECTED_COMPONENTS_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::StronglyConnectedComponents,
        native_algorithm: NativeGraphAlgorithm::StronglyConnectedComponents,
        sql_function_name: CUGRAPH_STRONGLY_CONNECTED_COMPONENTS_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::TriangleCountAll,
        native_algorithm: NativeGraphAlgorithm::TriangleCountAll,
        sql_function_name: CUGRAPH_TRIANGLE_COUNT_ALL_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::CoreNumber,
        native_algorithm: NativeGraphAlgorithm::CoreNumber,
        sql_function_name: CUGRAPH_CORE_NUMBER_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::KCore,
        native_algorithm: NativeGraphAlgorithm::KCore,
        sql_function_name: CUGRAPH_K_CORE_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::MinimumSpanningTree,
        native_algorithm: NativeGraphAlgorithm::MinimumSpanningTree,
        sql_function_name: CUGRAPH_MINIMUM_SPANNING_TREE_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::PersonalizedPageRank,
        native_algorithm: NativeGraphAlgorithm::PersonalizedPageRank,
        sql_function_name: CUGRAPH_PERSONALIZED_PAGERANK_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::Leiden,
        native_algorithm: NativeGraphAlgorithm::Leiden,
        sql_function_name: CUGRAPH_LEIDEN_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::Ecg,
        native_algorithm: NativeGraphAlgorithm::Ecg,
        sql_function_name: CUGRAPH_ECG_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::SpectralModularityMaximization,
        native_algorithm: NativeGraphAlgorithm::SpectralModularityMaximization,
        sql_function_name: CUGRAPH_SPECTRAL_MODULARITY_MAXIMIZATION_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::ForceAtlas2,
        native_algorithm: NativeGraphAlgorithm::ForceAtlas2,
        sql_function_name: CUGRAPH_FORCE_ATLAS2_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::Jaccard,
        native_algorithm: NativeGraphAlgorithm::Jaccard,
        sql_function_name: CUGRAPH_JACCARD_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::Sorensen,
        native_algorithm: NativeGraphAlgorithm::Sorensen,
        sql_function_name: CUGRAPH_SORENSEN_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::Overlap,
        native_algorithm: NativeGraphAlgorithm::Overlap,
        sql_function_name: CUGRAPH_OVERLAP_FUNCTION,
    },
    CugraphOperationDefinition {
        algorithm: CugraphAlgorithm::Cosine,
        native_algorithm: NativeGraphAlgorithm::Cosine,
        sql_function_name: CUGRAPH_COSINE_FUNCTION,
    },
];