pyspark/graphar_pyspark/util.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""Bindings to org.apache.graphar.util."""

from __future__ import annotations

from typing import Optional

from pyspark.sql import DataFrame

from graphar_pyspark import GraphArSession, _check_session


class IndexGenerator:
    """IndexGenerator is an object to help generating the indices for vertex/edge DataFrames."""

    @staticmethod
    def construct_vertex_index_mapping(
        vertex_df: DataFrame,
        primary_key: str,
    ) -> DataFrame:
        """Generate a vertex index mapping from the primary key.

        The resulting DataFrame contains two columns: vertex index & primary key.

        :param vertex_df: input vertex DataFrame.
        :param primary_key: the primary key of vertex
        :returns: a DataFrame contains two columns: vertex index & primary key.
        """
        _check_session()
        return DataFrame(
            GraphArSession.graphar.util.IndexGenerator.constructVertexIndexMapping(
                vertex_df._jdf,
                primary_key,
            ),
            GraphArSession.ss,
        )

    @staticmethod
    def generate_vertex_index_column(vertex_df: DataFrame) -> DataFrame:
        """Add a column contains vertex index to DataFrame.

        :param vertex_df: the input vertex DataFrame.
        :returns: DataFrame that contains a new vertex index column.
        """
        _check_session()
        return DataFrame(
            GraphArSession.graphar.util.IndexGenerator.generateVertexIndexColumn(
                vertex_df._jdf,
            ),
            GraphArSession.ss,
        )

    @staticmethod
    def generate_vertex_index_column_and_index_mapping(
        vertex_df: DataFrame,
        primary_key: str = "",
    ) -> (DataFrame, DataFrame):
        """Add an index column and generate a new index mapping.

        :param vertex_df: the input vertex DataFrame.
        :param primary_key: the primary key of vertex.
        :returns: the new vertex DataFrame and mapping DataFrame.
        """
        _check_session()
        jvm_res = GraphArSession.graphar.util.IndexGenerator.generateVertexIndexColumnAndIndexMapping(
            vertex_df._jdf,
            primary_key,
        )

        return (
            DataFrame(jvm_res._1(), GraphArSession.ss),
            DataFrame(jvm_res._2(), GraphArSession.ss),
        )

    @staticmethod
    def generate_edge_index_column(edge_df: DataFrame) -> DataFrame:
        """Add a column contains edge index to input edge DataFrame.

        :param edge_df: DataFrame with edges.
        :returns: DataFrame with edges and index.
        """
        _check_session()
        return DataFrame(
            GraphArSession.graphar.util.IndexGenerator.generateEdgeIndexColumn(
                edge_df._jdf,
            ),
            GraphArSession.ss,
        )

    @staticmethod
    def generate_src_index_for_edges_from_mapping(
        edge_df: DataFrame,
        src_column_name: str,
        src_index_mapping: DataFrame,
    ) -> DataFrame:
        """Join the edge table with the vertex index mapping for source column.

        :param edge_df: edges DataFrame
        :param src_column_name: join-column
        :param src_index_mapping: mapping DataFrame
        :returns: DataFrame with index
        """
        _check_session()
        return DataFrame(
            GraphArSession.graphar.util.IndexGenerator.generateSrcIndexForEdgesFromMapping(
                edge_df._jdf,
                src_column_name,
                src_index_mapping._jdf,
            ),
            GraphArSession.ss,
        )

    @staticmethod
    def generate_dst_index_for_edges_from_mapping(
        edge_df: DataFrame,
        dst_column_name: str,
        dst_index_mapping: DataFrame,
    ) -> DataFrame:
        """Join the edge table with the vertex index mapping for destination column.

        :param edge_df: edges DataFrame
        :param dst_column_name: join-column
        :param dst_index_mapping: mapping DataFrame
        :returns: DataFrame with index
        """
        _check_session()
        return DataFrame(
            GraphArSession.graphar.util.IndexGenerator.generateDstIndexForEdgesFromMapping(
                edge_df._jdf,
                dst_column_name,
                dst_index_mapping._jdf,
            ),
            GraphArSession.ss,
        )

    @staticmethod
    def generate_src_and_dst_index_for_edges_from_mapping(
        edge_df: DataFrame,
        src_column_name: Optional[str],
        dst_column_name: Optional[str],
        src_index_mapping: DataFrame,
        dst_index_mapping: DataFrame,
    ) -> DataFrame:
        """Join the edge table with the vertex index mapping for source & destination columns.

        Assumes that the first and second columns are the src and dst columns if they are None.


        :param edge_df: edge DataFrame
        :param src_column_name: src column, optional (the first col from edge_df will be used if None)
        :param dst_column_name: dst column, optional (the second col from edge_df will be used if None)
        :param src_index_mapping: source mapping DataFrame
        :param dst_index_mapping: dest mapping DataFrame
        :returns: DataFrame with indices
        """
        _check_session()
        if src_column_name is None:
            src_column_name = edge_df.columns[0]

        if dst_column_name is None:
            dst_column_name = edge_df.columns[1]

        return DataFrame(
            GraphArSession.graphar.util.IndexGenerator.generateSrcAndDstIndexForEdgesFromMapping(
                edge_df._jdf,
                src_column_name,
                dst_column_name,
                src_index_mapping._jdf,
                dst_index_mapping._jdf,
            ),
            GraphArSession.ss,
        )

    @staticmethod
    def generate_scr_index_for_edges(
        edge_df: DataFrame,
        src_column_name: str,
    ) -> DataFrame:
        """Construct vertex index for source column.

        :param edge_df: edge DataFrame
        :param src_column_name: source column
        :returns: DataFrame with index
        """
        _check_session()
        return DataFrame(
            GraphArSession.graphar.util.IndexGenerator.generateSrcIndexForEdges(
                edge_df._jdf,
                src_column_name,
            ),
            GraphArSession.ss,
        )

    @staticmethod
    def generate_dst_index_for_edges(
        edge_df: DataFrame,
        dst_column_name: str,
    ) -> DataFrame:
        """Construct vertex index for destination column.

        :param edge_df: edge DataFrame
        :param src_column_name: destination column
        :returns: DataFrame with index
        """
        _check_session()
        return DataFrame(
            GraphArSession.graphar.util.IndexGenerator.generateDstIndexForEdges(
                edge_df._jdf,
                dst_column_name,
            ),
            GraphArSession.ss,
        )

    @staticmethod
    def generate_src_and_dst_index_unitedly_for_edges(
        edge_df: DataFrame,
        src_column_name: str,
        dst_column_name: str,
    ) -> DataFrame:
        """Union and construct vertex index for source & destination columns.

        :param edge_df: edge DataFrame
        :param src_column_name: source column name
        :param dst_column_name: destination column name
        :returns: DataFrame with index
        """
        _check_session()
        return DataFrame(
            GraphArSession.graphar.util.IndexGenerator.generateSrcAndDstIndexUnitedlyForEdges(
                edge_df._jdf,
                src_column_name,
                dst_column_name,
            ),
            GraphArSession.ss,
        )