graphdatascience 1.11a4__tar.gz → 1.12__tar.gz

{graphdatascience-1.11a4 → graphdatascience-1.12}/MANIFEST.in

@@ -1,6 +1,7 @@
 include requirements/base/base.txt
 include requirements/base/ogb.txt
 include requirements/base/networkx.txt
+include requirements/base/rust-ext.txt
 include LICENSE
 prune graphdatascience/tests
 prune graphdatascience/resources/cora/serialize_cora.py

{graphdatascience-1.11a4/graphdatascience.egg-info → graphdatascience-1.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphdatascience
-Version: 1.11a4
+Version: 1.12
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j
@@ -41,6 +41,8 @@ Provides-Extra: ogb
 Requires-Dist: ogb<2.0,>=1.0; extra == "ogb"
 Provides-Extra: networkx
 Requires-Dist: networkx<4.0,>=2.0; extra == "networkx"
+Provides-Extra: rust-ext
+Requires-Dist: neo4j-rust-ext<6.0,>=4.4.2; extra == "rust-ext"
 
 # Neo4j Graph Data Science Client
 

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/algo/algo_proc_runner.py

@@ -3,11 +3,11 @@ from typing import Any, Dict, Tuple
 
 from pandas import DataFrame, Series
 
+from ..call_parameters import CallParameters
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..graph.graph_object import Graph
 from ..graph.graph_type_check import graph_type_check
 from ..model.graphsage_model import GraphSageModel
-from graphdatascience.call_parameters import CallParameters
 
 
 class AlgoProcRunner(IllegalAttrChecker, ABC):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/base_graph_proc_runner.py

@@ -9,11 +9,13 @@ from multimethod import multimethod
 from neo4j import __version__ as neo4j_driver_version
 from pandas import DataFrame, Series, read_parquet
 
+from ..call_parameters import CallParameters
 from ..error.client_only_endpoint import client_only_endpoint
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..error.uncallable_namespace import UncallableNamespace
 from ..server_version.compatible_with import compatible_with
 from ..server_version.server_version import ServerVersion
+from .graph_create_result import GraphCreateResult
 from .graph_entity_ops_runner import (
     GraphLabelRunner,
     GraphNodePropertiesRunner,
@@ -33,8 +35,6 @@ from .graph_type_check import (
     graph_type_check_optional,
 )
 from .ogb_loader import OGBLLoader, OGBNLoader
-from graphdatascience.call_parameters import CallParameters
-from graphdatascience.graph.graph_create_result import GraphCreateResult
 
 Strings = Union[str, List[str]]
 
@@ -120,7 +120,7 @@ class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
         )
         constructor.run(nodes, relationships)
 
-        return Graph(graph_name, self._query_runner, self._server_version)
+        return Graph(graph_name, self._query_runner)
 
     @client_only_endpoint("gds.graph")
     def load_cora(self, graph_name: str = "cora", undirected: bool = False) -> Graph:
@@ -250,7 +250,7 @@ class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
             params=params,
         ).squeeze()
 
-        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)
+        return GraphCreateResult(Graph(graph_name, self._query_runner), result)
 
     @from_graph_type_check
     @compatible_with("filter", min_inclusive=ServerVersion(2, 5, 0))
@@ -276,7 +276,7 @@ class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
             params=params,
         ).squeeze()
 
-        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)
+        return GraphCreateResult(Graph(graph_name, self._query_runner), result)
 
     @property
     def export(self) -> GraphExportRunner:
@@ -357,7 +357,7 @@ class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
                 f"No projected graph named '{graph_name}' exists in current database '{self._query_runner.database()}'"
             )
 
-        return Graph(graph_name, self._query_runner, self._server_version)
+        return Graph(graph_name, self._query_runner)
 
     @graph_type_check
     def _handle_properties(

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/graph_alpha_proc_runner.py

@@ -79,4 +79,4 @@ class GraphAlphaProcRunner(UncallableNamespace, IllegalAttrChecker):
         )
         constructor.run(nodes, relationships)
 
-        return Graph(graph_name, self._query_runner, self._server_version)
+        return Graph(graph_name, self._query_runner)

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/graph_beta_proc_runner.py

@@ -1,13 +1,13 @@
 from typing import Any, List, Union
 
+from ..call_parameters import CallParameters
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..error.uncallable_namespace import UncallableNamespace
+from .graph_create_result import GraphCreateResult
+from .graph_entity_ops_runner import GraphRelationshipsBetaRunner
 from .graph_export_runner import GraphExportCsvEndpoints
 from .graph_object import Graph
 from .graph_project_runner import GraphProjectBetaRunner
-from graphdatascience.call_parameters import CallParameters
-from graphdatascience.graph.graph_create_result import GraphCreateResult
-from graphdatascience.graph.graph_entity_ops_runner import GraphRelationshipsBetaRunner
 
 Strings = Union[str, List[str]]
 
@@ -40,4 +40,4 @@ class GraphBetaProcRunner(UncallableNamespace, IllegalAttrChecker):
 
         result = self._query_runner.call_procedure(endpoint=self._namespace, params=params).squeeze()
 
-        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)
+        return GraphCreateResult(Graph(graph_name, self._query_runner), result)

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/graph_cypher_runner.py

@@ -57,7 +57,7 @@ class GraphCypherRunner(CallerBase):
                 f"Invalid query, the query must end with the `RETURN {self._namespace}(...)` call: {query}"
             )
 
-        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), Series(data=result))
+        return GraphCreateResult(Graph(graph_name, self._query_runner), Series(data=result))
 
     __separators = re.compile(r"[,(.]")
 

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/graph_entity_ops_runner.py

@@ -5,6 +5,10 @@ from warnings import filterwarnings
 import pandas as pd
 from pandas import DataFrame, Series
 
+from ..call_parameters import CallParameters
+from ..error.cypher_warning_handler import (
+    filter_id_func_deprecation_warning,
+)
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..error.uncallable_namespace import UncallableNamespace
 from ..query_runner.query_runner import QueryRunner
@@ -13,10 +17,6 @@ from ..server_version.server_version import ServerVersion
 from ..utils.util_proc_runner import UtilProcRunner
 from .graph_object import Graph
 from .graph_type_check import graph_type_check
-from graphdatascience.call_parameters import CallParameters
-from graphdatascience.error.cypher_warning_handler import (
-    filter_id_func_deprecation_warning,
-)
 
 Strings = Union[str, List[str]]
 

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/graph_object.py

@@ -5,9 +5,8 @@ from typing import Any, List, Optional, Type, Union
 
 from pandas import Series
 
+from ..call_parameters import CallParameters
 from ..query_runner.query_runner import QueryRunner
-from ..server_version.server_version import ServerVersion
-from graphdatascience.call_parameters import CallParameters
 
 
 class Graph:
@@ -17,11 +16,10 @@ class Graph:
     It contains summary information about the graph.
     """
 
-    def __init__(self, name: str, query_runner: QueryRunner, server_version: ServerVersion):
+    def __init__(self, name: str, query_runner: QueryRunner):
         self._name = name
         self._query_runner = query_runner
         self._db = query_runner.database()
-        self._server_version = server_version
 
     def __enter__(self: Graph) -> Graph:
         return self

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/graph_proc_runner.py

@@ -1,6 +1,6 @@
+from .base_graph_proc_runner import BaseGraphProcRunner
+from .graph_cypher_runner import GraphCypherRunner
 from .graph_project_runner import GraphProjectRunner
-from graphdatascience.graph.base_graph_proc_runner import BaseGraphProcRunner
-from graphdatascience.graph.graph_cypher_runner import GraphCypherRunner
 
 
 class GraphProcRunner(BaseGraphProcRunner):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/graph_project_runner.py

@@ -4,11 +4,11 @@ from typing import Any
 
 from pandas import Series
 
+from ..call_parameters import CallParameters
 from ..error.illegal_attr_checker import IllegalAttrChecker
+from .graph_create_result import GraphCreateResult
 from .graph_object import Graph
 from .graph_type_check import from_graph_type_check
-from graphdatascience.call_parameters import CallParameters
-from graphdatascience.graph.graph_create_result import GraphCreateResult
 
 
 class GraphProjectRunner(IllegalAttrChecker):
@@ -25,7 +25,7 @@ class GraphProjectRunner(IllegalAttrChecker):
             logging=True,
         ).squeeze()
 
-        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)
+        return GraphCreateResult(Graph(graph_name, self._query_runner), result)
 
     def estimate(self, node_projection: Any, relationship_projection: Any, **config: Any) -> "Series[Any]":
         self._namespace += ".estimate"
@@ -69,4 +69,4 @@ class GraphProjectBetaRunner(IllegalAttrChecker):
             logging=True,
         ).squeeze()
 
-        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)
+        return GraphCreateResult(Graph(graph_name, self._query_runner), result)

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/graph_remote_project_runner.py

@@ -2,13 +2,13 @@ from __future__ import annotations
 
 from typing import List, Optional
 
+from ..call_parameters import CallParameters
 from ..error.illegal_attr_checker import IllegalAttrChecker
-from ..query_runner.aura_db_query_runner import AuraDbQueryRunner
+from ..query_runner.session_query_runner import SessionQueryRunner
 from ..server_version.compatible_with import compatible_with
+from ..server_version.server_version import ServerVersion
+from .graph_create_result import GraphCreateResult
 from .graph_object import Graph
-from graphdatascience.call_parameters import CallParameters
-from graphdatascience.graph.graph_create_result import GraphCreateResult
-from graphdatascience.server_version.server_version import ServerVersion
 
 
 class GraphProjectRemoteRunner(IllegalAttrChecker):
@@ -41,7 +41,7 @@ class GraphProjectRemoteRunner(IllegalAttrChecker):
         )
 
         result = self._query_runner.call_procedure(
-            endpoint=AuraDbQueryRunner.GDS_REMOTE_PROJECTION_PROC_NAME,
+            endpoint=SessionQueryRunner.GDS_REMOTE_PROJECTION_PROC_NAME,
             params=params,
         ).squeeze()
-        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)
+        return GraphCreateResult(Graph(graph_name, self._query_runner), result)

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/graph_sample_runner.py

@@ -2,13 +2,13 @@ from typing import Any
 
 from pandas import Series
 
+from ..call_parameters import CallParameters
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..server_version.compatible_with import compatible_with
 from ..server_version.server_version import ServerVersion
+from .graph_create_result import GraphCreateResult
 from .graph_object import Graph
 from .graph_type_check import from_graph_type_check
-from graphdatascience.call_parameters import CallParameters
-from graphdatascience.graph.graph_create_result import GraphCreateResult
 
 
 class GraphAlphaSampleRunner(IllegalAttrChecker):
@@ -41,7 +41,7 @@ class RWRRunner(IllegalAttrChecker):
 
         result = self._query_runner.call_procedure(endpoint=self._namespace, params=params, logging=True).squeeze()
 
-        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)
+        return GraphCreateResult(Graph(graph_name, self._query_runner), result)
 
 
 class CNARWRunner(IllegalAttrChecker):
@@ -55,7 +55,7 @@ class CNARWRunner(IllegalAttrChecker):
         )
         result = self._query_runner.call_procedure(endpoint=self._namespace, params=params, logging=True).squeeze()
 
-        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)
+        return GraphCreateResult(Graph(graph_name, self._query_runner), result)
 
     def estimate(self, from_G: Graph, **config: Any) -> "Series[Any]":
         self._namespace += ".estimate"

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/nx_loader.py

@@ -25,7 +25,7 @@ class NXLoader(UncallableNamespace, IllegalAttrChecker):
         constructor = self._query_runner.create_graph_constructor(graph_name, concurrency, undirected_rel_types)
         constructor.run(nodes, rels)
 
-        return Graph(graph_name, self._query_runner, self._server_version)
+        return Graph(graph_name, self._query_runner)
 
     @staticmethod
     def _attr_to_labels_key(labels_attr: Any, node_id: Any, no_node_labels: Optional[bool]) -> Tuple[str, ...]:

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph/ogb_loader.py

@@ -96,7 +96,7 @@ class OGBLoader(UncallableNamespace, IllegalAttrChecker, ABC):
         constructor = self._query_runner.create_graph_constructor(graph_name, concurrency, [])
         constructor.run(nodes, rels)
 
-        return Graph(graph_name, self._query_runner, self._server_version)
+        return Graph(graph_name, self._query_runner)
 
 
 class OGBNLoader(OGBLoader):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/graph_data_science.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from typing import Any, Dict, Optional, Tuple, Type, Union
 
 from neo4j import Driver
@@ -8,12 +9,14 @@ from pandas import DataFrame
 from .call_builder import IndirectCallBuilder
 from .endpoints import AlphaEndpoints, BetaEndpoints, DirectEndpoints
 from .error.uncallable_namespace import UncallableNamespace
+from .graph.graph_proc_runner import GraphProcRunner
+from .query_runner.arrow_info import ArrowInfo
 from .query_runner.arrow_query_runner import ArrowQueryRunner
 from .query_runner.neo4j_query_runner import Neo4jQueryRunner
 from .query_runner.query_runner import QueryRunner
 from .server_version.server_version import ServerVersion
-from graphdatascience.graph.graph_proc_runner import GraphProcRunner
-from graphdatascience.utils.util_proc_runner import UtilProcRunner
+from .utils.util_proc_runner import UtilProcRunner
+from .version import __min_server_version__
 
 
 class GraphDataScience(DirectEndpoints, UncallableNamespace):
@@ -33,6 +36,7 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
         arrow_disable_server_verification: bool = True,
         arrow_tls_root_certs: Optional[bytes] = None,
         bookmarks: Optional[Any] = None,
+        show_progress: bool = True,
     ):
         """
         Construct a new GraphDataScience object.
@@ -49,11 +53,12 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
         database: Optional[str], default None
             The Neo4j database to query against.
         arrow : Union[str, bool], default True
-            Arrow connection information. This is either a bool or a string.
-            If it is a string, it will be interpreted as a connection URL to a GDS Arrow Server.
-            If it is a bool,
-                True will make the client discover the connection URI to the GDS Arrow server via the Neo4j endpoint,
-                while False will make the client use Bolt for all operations.
+            Arrow connection information. This is either a string or a bool.
+
+            - If it is a string, it will be interpreted as a connection URL to a GDS Arrow Server.
+            - If it is a bool:
+                - True will make the client discover the connection URI to the GDS Arrow server via the Neo4j endpoint.
+                - False will make the client use Bolt for all operations.
         arrow_disable_server_verification : bool, default True
             A flag that overrides other TLS settings and disables server verification for TLS connections.
         arrow_tls_root_certs : Optional[bytes], default None
@@ -61,6 +66,8 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
             GDS Arrow Flight server.
         bookmarks : Optional[Any], default None
             The Neo4j bookmarks to require a certain state before the next query gets executed.
+        show_progress : bool, default True
+            A flag to indicate whether to show progress bars for running procedures.
         """
         if aura_ds:
             GraphDataScience._validate_endpoint(endpoint)
@@ -68,13 +75,24 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
         if isinstance(endpoint, QueryRunner):
             self._query_runner = endpoint
         else:
-            self._query_runner = Neo4jQueryRunner.create(endpoint, auth, aura_ds, database, bookmarks)
+            self._query_runner = Neo4jQueryRunner.create(endpoint, auth, aura_ds, database, bookmarks, show_progress)
 
         self._server_version = self._query_runner.server_version()
 
-        if arrow and self._server_version >= ServerVersion(2, 1, 0):
+        if self._server_version < ServerVersion.from_string(__min_server_version__):
+            warnings.warn(
+                DeprecationWarning(
+                    f"Client does not support the given server version `{self._server_version}`."
+                    + " We recommend to either update the GDS server version or use a compatible version of the `graphdatascience` package."
+                    + " Please refer to the compatibility matrix at https://neo4j.com/docs/graph-data-science-client/current/installation/#python-client-system-requirements."
+                )
+            )
+
+        arrow_info = ArrowInfo.create(self._query_runner)
+        if arrow and arrow_info.enabled and self._server_version >= ServerVersion(2, 1, 0):
             self._query_runner = ArrowQueryRunner.create(
                 self._query_runner,
+                arrow_info,
                 auth,
                 self._query_runner.encrypted(),
                 arrow_disable_server_verification,
@@ -82,6 +100,7 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
                 None if arrow is True else arrow,
             )
 
+        self._query_runner.set_show_progress(show_progress)
         super().__init__(self._query_runner, namespace="gds", server_version=self._server_version)
 
     @property
@@ -125,6 +144,17 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
         """
         self._query_runner.set_bookmarks(bookmarks)
 
+    def set_show_progress(self, show_progress: bool) -> None:
+        """
+        Set whether to show progress for running procedures.
+
+        Parameters
+        ----------
+        show_progress: bool
+            Whether to show progress for procedures.
+        """
+        self._query_runner.set_show_progress(show_progress)
+
     def database(self) -> Optional[str]:
         """
         Get the database which queries are run against.
@@ -227,3 +257,10 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
         Close the GraphDataScience object and release any resources held by it.
         """
         self._query_runner.close()
+
+    def __enter__(self) -> GraphDataScience:
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        self.close()
+        return False

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/model/graphsage_model.py

@@ -2,10 +2,10 @@ from typing import Any
 
 from pandas import Series
 
+from ..call_parameters import CallParameters
 from ..graph.graph_object import Graph
 from ..graph.graph_type_check import graph_type_check
 from .model import Model
-from graphdatascience.call_parameters import CallParameters
 
 
 class GraphSageModel(Model):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/model/model.py

@@ -5,12 +5,12 @@ from typing import Any, Dict
 
 from pandas import DataFrame, Series
 
+from ..call_parameters import CallParameters
 from ..graph.graph_object import Graph
 from ..graph.graph_type_check import graph_type_check
 from ..query_runner.query_runner import QueryRunner
 from ..server_version.compatible_with import compatible_with
 from ..server_version.server_version import ServerVersion
-from graphdatascience.call_parameters import CallParameters
 
 
 class Model(ABC):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/model/model_alpha_proc_runner.py

@@ -2,10 +2,10 @@ from typing import Any, Tuple
 
 from pandas import Series
 
+from ..call_parameters import CallParameters
 from .model import Model
 from .model_proc_runner import ModelProcRunner
 from .model_resolver import ModelResolver
-from graphdatascience.call_parameters import CallParameters
 
 
 class ModelAlphaProcRunner(ModelResolver):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/model/model_beta_proc_runner.py

@@ -2,10 +2,10 @@ from typing import Any, Optional
 
 from pandas import DataFrame, Series
 
+from ..call_parameters import CallParameters
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..error.uncallable_namespace import UncallableNamespace
 from .model import Model
-from graphdatascience.call_parameters import CallParameters
 
 
 class ModelBetaProcRunner(UncallableNamespace, IllegalAttrChecker):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/model/model_proc_runner.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 from pandas import DataFrame, Series
 
+from ..call_parameters import CallParameters
 from ..error.client_only_endpoint import client_only_endpoint
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..error.uncallable_namespace import UncallableNamespace
@@ -13,7 +14,6 @@ from ..server_version.compatible_with import compatible_with
 from ..server_version.server_version import ServerVersion
 from .model import Model
 from .model_resolver import ModelResolver
-from graphdatascience.call_parameters import CallParameters
 
 
 class DistMultCreator(UncallableNamespace, IllegalAttrChecker):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/model/node_classification_model.py

@@ -2,10 +2,10 @@ from typing import Any, Dict, List
 
 from pandas import Series
 
+from ..call_parameters import CallParameters
 from ..graph.graph_object import Graph
 from ..graph.graph_type_check import graph_type_check
 from .pipeline_model import PipelineModel
-from graphdatascience.call_parameters import CallParameters
 
 
 class NCModel(PipelineModel):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/pipeline/classification_training_pipeline.py

@@ -3,9 +3,9 @@ from typing import Any
 
 from pandas import Series
 
+from ..call_parameters import CallParameters
 from ..server_version.server_version import ServerVersion
 from .training_pipeline import MODEL_TYPE, TrainingPipeline
-from graphdatascience.call_parameters import CallParameters
 
 
 class ClassificationTrainingPipeline(TrainingPipeline[MODEL_TYPE], ABC):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/pipeline/lp_pipeline_create_runner.py

@@ -2,10 +2,10 @@ from typing import Any, Tuple
 
 from pandas import Series
 
+from ..call_parameters import CallParameters
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..error.uncallable_namespace import UncallableNamespace
 from .lp_training_pipeline import LPTrainingPipeline
-from graphdatascience.call_parameters import CallParameters
 
 
 class LPPipelineCreateRunner(UncallableNamespace, IllegalAttrChecker):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/pipeline/lp_training_pipeline.py

@@ -2,10 +2,10 @@ from typing import Any
 
 from pandas import DataFrame, Series
 
+from ..call_parameters import CallParameters
 from ..model.link_prediction_model import LPModel
 from ..query_runner.query_runner import QueryRunner
 from .classification_training_pipeline import ClassificationTrainingPipeline
-from graphdatascience.call_parameters import CallParameters
 
 
 class LPTrainingPipeline(ClassificationTrainingPipeline[LPModel]):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/pipeline/nc_training_pipeline.py

@@ -3,10 +3,10 @@ from typing import Any, List, Union
 
 from pandas import Series
 
+from ..call_parameters import CallParameters
 from ..model.node_classification_model import NCModel
 from ..pipeline.classification_training_pipeline import ClassificationTrainingPipeline
 from ..query_runner.query_runner import QueryRunner
-from graphdatascience.call_parameters import CallParameters
 
 
 class NCTrainingPipeline(ClassificationTrainingPipeline[NCModel], ABC):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/pipeline/nr_training_pipeline.py

@@ -2,10 +2,10 @@ from typing import Any, List, Union
 
 from pandas import Series
 
+from ..call_parameters import CallParameters
 from ..model.node_regression_model import NRModel
 from ..query_runner.query_runner import QueryRunner
 from .training_pipeline import TrainingPipeline
-from graphdatascience.call_parameters import CallParameters
 
 
 class NRTrainingPipeline(TrainingPipeline[NRModel]):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/pipeline/pipeline_proc_runner.py

@@ -2,16 +2,16 @@ from typing import Any, Optional
 
 from pandas import DataFrame, Series
 
+from ..call_parameters import CallParameters
 from ..error.client_only_endpoint import client_only_endpoint
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..error.uncallable_namespace import UncallableNamespace
 from ..model.pipeline_model import PipelineModel
+from ..server_version.server_version import ServerVersion
 from .lp_training_pipeline import LPTrainingPipeline
 from .nc_training_pipeline import NCTrainingPipeline
 from .nr_training_pipeline import NRTrainingPipeline
 from .training_pipeline import TrainingPipeline
-from graphdatascience.call_parameters import CallParameters
-from graphdatascience.server_version.server_version import ServerVersion
 
 
 class PipelineProcRunner(UncallableNamespace, IllegalAttrChecker):

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/pipeline/training_pipeline.py

@@ -3,12 +3,12 @@ from typing import Any, Dict, Generic, Tuple, TypeVar
 
 from pandas import DataFrame, Series
 
+from ..call_parameters import CallParameters
 from ..graph.graph_object import Graph
 from ..graph.graph_type_check import graph_type_check
 from ..model.pipeline_model import PipelineModel
 from ..query_runner.query_runner import QueryRunner
 from ..server_version.server_version import ServerVersion
-from graphdatascience.call_parameters import CallParameters
 
 MODEL_TYPE = TypeVar("MODEL_TYPE", bound=PipelineModel, covariant=True)
 
@@ -238,9 +238,7 @@ class TrainingPipeline(ABC, Generic[MODEL_TYPE]):
             params=CallParameters(name=self.name()),
             yields=["exists"],
             custom_error=False,
-        )[
-            "exists"
-        ].squeeze()  # type: ignore
+        )["exists"].squeeze()  # type: ignore
 
     def drop(self, failIfMissing: bool = False) -> "Series[Any]":
         """

{graphdatascience-1.11a4 → graphdatascience-1.12}/graphdatascience/query_runner/arrow_graph_constructor.py

@@ -88,13 +88,16 @@ class ArrowGraphConstructor(GraphConstructor):
 
         writer, _ = self._client.start_put(flight_descriptor, table.schema)
 
-        with writer:
-            # Write table in chunks
-            for partition in batches:
-                writer.write_batch(partition)
-                pbar.update(partition.num_rows)
-        # Force a refresh to avoid the progress bar getting stuck at 0%
-        pbar.refresh()
+        try:
+            with writer:
+                # Write table in chunks
+                for partition in batches:
+                    writer.write_batch(partition)
+                    pbar.update(partition.num_rows)
+            # Force a refresh to avoid the progress bar getting stuck at 0%
+            pbar.refresh()
+        except Exception as e:
+            GdsArrowClient.handle_flight_error(e)
 
     def _send_dfs(self, dfs: List[DataFrame], entity_type: str) -> None:
         desc = "Uploading Nodes" if entity_type == "node" else "Uploading Relationships"