graphar_pyspark.graph

  1# Licensed to the Apache Software Foundation (ASF) under one
  2# or more contributor license agreements.  See the NOTICE file
  3# distributed with this work for additional information
  4# regarding copyright ownership.  The ASF licenses this file
  5# to you under the Apache License, Version 2.0 (the
  6# "License"); you may not use this file except in compliance
  7# with the License.  You may obtain a copy of the License at
  8#
  9#   http://www.apache.org/licenses/LICENSE-2.0
 10#
 11# Unless required by applicable law or agreed to in writing,
 12# software distributed under the License is distributed on an
 13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 14# KIND, either express or implied.  See the License for the
 15# specific language governing permissions and limitations
 16# under the License.
 17
 18"""Bidnings to org.apache.graphar.graph."""
 19
 20from __future__ import annotations
 21
 22from collections.abc import Mapping
 23from dataclasses import dataclass
 24from typing import Optional, Union
 25
 26from py4j.java_gateway import JavaObject
 27from pyspark.sql import DataFrame
 28
 29from graphar_pyspark import GraphArSession, _check_session
 30from graphar_pyspark.enums import FileType
 31from graphar_pyspark.errors import InvalidGraphFormatError
 32from graphar_pyspark.info import GraphInfo
 33
 34
 35@dataclass(frozen=True)
 36class EdgeLabels:
 37    """A triplet that describe edge. Contains source, edge and dest labels. Immutable."""
 38
 39    src_label: str
 40    edge_label: str
 41    dst_label: str
 42
 43
 44@dataclass(frozen=True)
 45class GraphReaderResult:
 46    """A simple immutable class, that represent results of reading a graph with GraphReader."""
 47
 48    vertex_dataframes: Mapping[str, DataFrame]
 49    edge_dataframes: Mapping[EdgeLabels, Mapping[str, DataFrame]]
 50
 51    @staticmethod
 52    def from_scala(
 53        jvm_result: tuple[
 54            dict[str, JavaObject],
 55            dict[tuple[str, str, str], dict[str, JavaObject]],
 56        ],
 57    ) -> "GraphReaderResult":
 58        """Create an instance of the Class from JVM method output.
 59
 60        :param jvm_result: structure, returned from JVM.
 61        :returns: instance of Python Class.
 62        """
 63        first_dict = {}
 64        first_scala_map = jvm_result._1()
 65        first_scala_map_iter = first_scala_map.keySet().iterator()
 66
 67        while first_scala_map_iter.hasNext():
 68            k = first_scala_map_iter.next()
 69            first_dict[k] = DataFrame(first_scala_map.get(k).get(), GraphArSession.ss)
 70
 71        second_dict = {}
 72        second_scala_map = jvm_result._2()
 73        second_scala_map_iter = second_scala_map.keySet().iterator()
 74
 75        while second_scala_map_iter.hasNext():
 76            k = second_scala_map_iter.next()
 77            nested_scala_map = second_scala_map.get(k).get()
 78            nested_scala_map_iter = nested_scala_map.keySet().iterator()
 79            inner_dict = {}
 80
 81            while nested_scala_map_iter.hasNext():
 82                kk = nested_scala_map_iter.next()
 83                inner_dict[kk] = DataFrame(
 84                    nested_scala_map.get(kk).get(),
 85                    GraphArSession.ss,
 86                )
 87
 88            second_dict[EdgeLabels(k._1(), k._2(), k._3())] = inner_dict
 89
 90        return GraphReaderResult(
 91            vertex_dataframes=first_dict,
 92            edge_dataframes=second_dict,
 93        )
 94
 95
 96class GraphReader:
 97    """The helper object for reading graph through the definitions of graph info."""
 98
 99    @staticmethod
100    def read(
101        graph_info: Union[GraphInfo, str],
102    ) -> GraphReaderResult:
103        """Read the graph as vertex and edge DataFrames with the graph info yaml file or GraphInfo object.
104
105        :param graph_info: The path of the graph info yaml or GraphInfo instance.
106        :returns: GraphReaderResults, that contains vertex and edge dataframes.
107        """
108        _check_session()
109        if isinstance(graph_info, str):
110            graph_info = GraphInfo.load_graph_info(graph_info)
111
112        jvm_result = GraphArSession.graphar.graph.GraphReader.readWithGraphInfo(
113            graph_info.to_scala(),
114            GraphArSession.jss,
115        )
116        return GraphReaderResult.from_scala(jvm_result)
117
118
119class GraphWriter:
120    """The helper class for writing graph."""
121
122    def __init__(self, jvm_obj: JavaObject) -> None:
123        """One should not use this constructor directly, please use `from_scala` or `from_python`."""
124        _check_session()
125        self._jvm_graph_writer_obj = jvm_obj
126
127    def to_scala(self) -> JavaObject:
128        """Transform object to JVM representation.
129
130        :returns: JavaObject
131        """
132        return self._jvm_graph_writer_obj
133
134    @staticmethod
135    def from_scala(jvm_obj: JavaObject) -> "GraphWriter":
136        """Create an instance of the Class from the corresponding JVM object.
137
138        :param jvm_obj: scala object in JVM.
139        :returns: instance of Python Class.
140        """
141        return GraphWriter(jvm_obj)
142
143    @staticmethod
144    def from_python() -> "GraphWriter":
145        """Create an instance of the Class from Python arguments."""
146        return GraphWriter(GraphArSession.graphar.graph.GraphWriter())
147
148    def put_vertex_data(self, label: str, df: DataFrame, primary_key: str) -> None:
149        """Put the vertex DataFrame into writer.
150
151        :param label: label of vertex.
152        :param df: DataFrame of the vertex type.
153        :param primary_key: primary key of the vertex type, default is empty, which take the first property column as primary key.
154        """
155        self._jvm_graph_writer_obj.PutVertexData(label, df._jdf, primary_key)
156
157    def put_edge_data(self, relation: tuple[str, str, str], df: DataFrame) -> None:
158        """Put the egde datafrme into writer.
159
160        :param relation: 3-Tuple (source label, edge label, target label) to indicate edge type.
161        :param df: data frame of edge type.
162        """
163        relation_jvm = GraphArSession.jvm.scala.Tuple3(
164            relation[0], relation[1], relation[2],
165        )
166        self._jvm_graph_writer_obj.PutEdgeData(relation_jvm, df._jdf)
167
168    def write_with_graph_info(self, graph_info: Union[GraphInfo, str]) -> None:
169        """Write the graph data in GraphAr format with graph info.
170
171        Note: original method is `write` but there is not directly overloading in Python.
172
173        :param graph_info: the graph info object for the graph or the path to graph info object.
174        """
175        if isinstance(graph_info, str):
176            self._jvm_graph_writer_obj.write(graph_info, GraphArSession.jss)
177        else:
178            self._jvm_graph_writer_obj.write(graph_info.to_scala(), GraphArSession.jss)
179
180    def write(
181        self,
182        path: str,
183        name: str = "graph",
184        vertex_chunk_size: Optional[int] = None,
185        edge_chunk_size: Optional[int] = None,
186        file_type: Optional[FileType] = None,
187        version: Optional[str] = None,
188    ) -> None:
189        """Write graph data in GraphAr format.
190
191        Note: for default parameters check org.apache.graphar.GeneralParams;
192        For this method None for any of arguments means that the default value will be used.
193
194        :param path: the directory to write.
195        :param name: the name of graph, default is 'grpah'
196        :param vertex_chunk_size: the chunk size for vertices, default is 2^18
197        :param edge_chunk_size: the chunk size for edges, default is 2^22
198        :param file_type: the file type for data payload file, support [parquet, orc, csv, json], default is parquet.
199        :param version: version of GraphAr format, default is v1.
200        """
201        if vertex_chunk_size is None:
202            vertex_chunk_size = (
203                GraphArSession.graphar.GeneralParams.defaultVertexChunkSize
204            )
205
206        if edge_chunk_size is None:
207            edge_chunk_size = GraphArSession.graphar.GeneralParams.defaultEdgeChunkSize
208
209        file_type = (
210            GraphArSession.graphar.GeneralParams.defaultFileType
211            if file_type is None
212            else file_type.value
213        )
214
215        if version is None:
216            version = GraphArSession.graphar.GeneralParams.defaultVersion
217
218        self._jvm_graph_writer_obj.write(
219            path,
220            GraphArSession.jss,
221            name,
222            vertex_chunk_size,
223            edge_chunk_size,
224            file_type,
225            version,
226        )
227
228
229class GraphTransformer:
230    """The helper object for transforming graphs through the definitions of their infos."""
231
232    @staticmethod
233    def transform(
234        source_graph_info: Union[str, GraphInfo],
235        dest_graph_info: Union[str, GraphInfo],
236    ) -> None:
237        """Transform the graphs following the meta data provided or defined in info files.
238
239        Note: both arguments should be strings or GrapInfo instances! Mixed arguments type is not supported.
240
241        :param source_graph_info: The path of the graph info yaml file for the source graph OR the info object for the source graph.
242        :param dest_graph_info: The path of the graph info yaml file for the destination graph OR the info object for the destination graph.
243        :raise InvalidGraphFormatException: if you pass mixed format of source and dest graph info.
244        """
245        _check_session()
246        if isinstance(source_graph_info, str) and isinstance(dest_graph_info, str):
247            GraphArSession.graphar.graph.GraphTransformer.transform(
248                source_graph_info,
249                dest_graph_info,
250                GraphArSession.jss,
251            )
252        elif isinstance(source_graph_info, GraphInfo) and isinstance(
253            dest_graph_info,
254            GraphInfo,
255        ):
256            GraphArSession.graphar.graph.GraphTransformer.transform(
257                source_graph_info.to_scala(),
258                dest_graph_info.to_scala(),
259                GraphArSession.jss,
260            )
261        else:
262            msg = "Both src and dst graph info objects should be of the same type. "
263            msg += f"But {type(source_graph_info)} and {type(dest_graph_info)} were provided!"
264            raise InvalidGraphFormatError(msg)