graphdatascience 1.10a1__tar.gz → 1.11a2__tar.gz

{graphdatascience-1.10a1/graphdatascience.egg-info → graphdatascience-1.11a2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphdatascience
-Version: 1.10a1
+Version: 1.11a2
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j
@@ -31,7 +31,7 @@ License-File: LICENSE
 Requires-Dist: multimethod<2.0,>=1.0
 Requires-Dist: neo4j<6.0,>=4.4.2
 Requires-Dist: pandas<3.0,>=1.0
-Requires-Dist: pyarrow<15.0,>=10.0
+Requires-Dist: pyarrow<16.0,>=11.0
 Requires-Dist: textdistance<5.0,>=4.0
 Requires-Dist: tqdm<5.0,>=4.0
 Requires-Dist: typing-extensions<5.0,>=4.0

{graphdatascience-1.10a1 → graphdatascience-1.11a2}/graphdatascience/graph/base_graph_proc_runner.py

@@ -18,6 +18,7 @@ from .graph_entity_ops_runner import (
     GraphElementPropertyRunner,
     GraphLabelRunner,
     GraphNodePropertiesRunner,
+    GraphNodePropertyRunner,
     GraphPropertyRunner,
     GraphRelationshipPropertiesRunner,
     GraphRelationshipRunner,
@@ -379,9 +380,9 @@ class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
         )
 
     @property
-    def nodeProperty(self) -> GraphElementPropertyRunner:
+    def nodeProperty(self) -> GraphNodePropertyRunner:
         self._namespace += ".nodeProperty"
-        return GraphElementPropertyRunner(self._query_runner, self._namespace, self._server_version)
+        return GraphNodePropertyRunner(self._query_runner, self._namespace, self._server_version)
 
     @property
     def nodeProperties(self) -> GraphNodePropertiesRunner:
@@ -516,8 +517,7 @@ class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
         ).squeeze()
 
     @multimethod
-    def removeNodeProperties(self) -> None:
-        ...
+    def removeNodeProperties(self) -> None: ...
 
     @removeNodeProperties.register
     @graph_type_check

{graphdatascience-1.10a1 → graphdatascience-1.11a2}/graphdatascience/graph/graph_entity_ops_runner.py

@@ -77,6 +77,26 @@ class GraphElementPropertyRunner(GraphEntityOpsBaseRunner):
         return self._handle_properties(G, node_properties, node_labels, config)
 
 
+class GraphNodePropertyRunner(GraphEntityOpsBaseRunner):
+    @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
+    @filter_id_func_deprecation_warning()
+    def stream(
+        self,
+        G: Graph,
+        node_property: str,
+        node_labels: Strings = ["*"],
+        db_node_properties: List[str] = [],
+        **config: Any,
+    ) -> DataFrame:
+        self._namespace += ".stream"
+
+        result = self._handle_properties(G, node_property, node_labels, config)
+
+        return GraphNodePropertiesRunner._process_result(
+            self._query_runner, list(node_property), False, db_node_properties, result, config
+        )
+
+
 class GraphNodePropertiesRunner(GraphEntityOpsBaseRunner):
     @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
     @filter_id_func_deprecation_warning()
@@ -93,6 +113,19 @@ class GraphNodePropertiesRunner(GraphEntityOpsBaseRunner):
 
         result = self._handle_properties(G, node_properties, node_labels, config)
 
+        return GraphNodePropertiesRunner._process_result(
+            self._query_runner, node_properties, separate_property_columns, db_node_properties, result, config
+        )
+
+    @staticmethod
+    def _process_result(
+        query_runner: QueryRunner,
+        node_properties: List[str],
+        separate_property_columns: bool,
+        db_node_properties: List[str],
+        result: DataFrame,
+        config: Dict[str, Any],
+    ) -> DataFrame:
         # new format was requested, but the query was run via Cypher
         if separate_property_columns and "propertyValue" in result.keys():
             wide_result = result.pivot(index=["nodeId"], columns=["nodeProperty"], values="propertyValue")
@@ -106,7 +139,7 @@ class GraphNodePropertiesRunner(GraphEntityOpsBaseRunner):
         # old format was requested but the query was run via Arrow
         elif not separate_property_columns and "propertyValue" not in result.keys():
             id_vars = ["nodeId", "nodeLabels"] if config.get("listNodeLabels", False) else ["nodeId"]
-            result = result.melt(id_vars=id_vars).rename(columns={"variable": "nodeProperty", "value": "propertyValue"})
+            result = result.melt(id_vars=id_vars, var_name="nodeProperty", value_name="propertyValue")
 
         if db_node_properties:
             duplicate_properties = set(db_node_properties).intersection(set(node_properties))
@@ -116,16 +149,20 @@ class GraphNodePropertiesRunner(GraphEntityOpsBaseRunner):
                 )
 
             unique_node_ids = result["nodeId"].drop_duplicates().tolist()
-            db_properties_df = self._query_runner.run_cypher(
-                self._build_query(db_node_properties), {"ids": unique_node_ids}
+            db_properties_df = query_runner.run_cypher(
+                GraphNodePropertiesRunner._build_query(db_node_properties), {"ids": unique_node_ids}
             )
 
             if "propertyValue" not in result.keys():
                 result = result.join(db_properties_df.set_index("nodeId"), on="nodeId")
             else:
-                db_properties_df = db_properties_df.melt(id_vars=["nodeId"]).rename(
-                    columns={"variable": "nodeProperty", "value": "propertyValue"}
+                db_properties_df = db_properties_df.melt(
+                    id_vars=["nodeId"], var_name="nodeProperty", value_name="propertyValue"
                 )
+
+                if "nodeProperty" not in result.keys():
+                    result["nodeProperty"] = node_properties[0]
+
                 result = pd.concat([result, db_properties_df])
 
         return result
@@ -140,7 +177,7 @@ class GraphNodePropertiesRunner(GraphEntityOpsBaseRunner):
         return reduce(add_property, db_node_properties, query_prefix)
 
     @compatible_with("write", min_inclusive=ServerVersion(2, 2, 0))
-    def write(self, G: Graph, node_properties: List[str], node_labels: Strings = ["*"], **config: Any) -> "Series[Any]":
+    def write(self, G: Graph, node_properties: Strings, node_labels: Strings = ["*"], **config: Any) -> "Series[Any]":
         self._namespace += ".write"
         return self._handle_properties(G, node_properties, node_labels, config).squeeze()  # type: ignore
 

{graphdatascience-1.10a1 → graphdatascience-1.11a2}/graphdatascience/graph/graph_remote_proc_runner.py

@@ -5,5 +5,4 @@ from graphdatascience.graph.graph_remote_project_runner import GraphProjectRemot
 class GraphRemoteProcRunner(BaseGraphProcRunner):
     @property
     def project(self) -> GraphProjectRemoteRunner:
-        self._namespace += ".project.remoteDb"
         return GraphProjectRemoteRunner(self._query_runner, self._namespace, self._server_version)

graphdatascience-1.11a2/graphdatascience/graph/graph_remote_project_runner.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+from ..error.illegal_attr_checker import IllegalAttrChecker
+from ..query_runner.aura_db_query_runner import AuraDbQueryRunner
+from ..server_version.compatible_with import compatible_with
+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):
+    @compatible_with("project", min_inclusive=ServerVersion(2, 7, 0))
+    def __call__(
+        self,
+        graph_name: str,
+        query: str,
+        concurrency: int = 4,
+        undirected_relationship_types: Optional[List[str]] = None,
+        inverse_indexed_relationship_types: Optional[List[str]] = None,
+    ) -> GraphCreateResult:
+        if inverse_indexed_relationship_types is None:
+            inverse_indexed_relationship_types = []
+        if undirected_relationship_types is None:
+            undirected_relationship_types = []
+
+        params = CallParameters(
+            graph_name=graph_name,
+            query=query,
+            concurrency=concurrency,
+            undirected_relationship_types=undirected_relationship_types,
+            inverse_indexed_relationship_types=inverse_indexed_relationship_types,
+        )
+
+        result = self._query_runner.call_procedure(
+            endpoint=AuraDbQueryRunner.GDS_REMOTE_PROJECTION_PROC_NAME,
+            params=params,
+        ).squeeze()
+        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)

{graphdatascience-1.10a1 → graphdatascience-1.11a2}/graphdatascience/graph_data_science.py

@@ -23,11 +23,12 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
 
     def __init__(
         self,
+        /,
         endpoint: Union[str, Driver, QueryRunner],
         auth: Optional[Tuple[str, str]] = None,
         aura_ds: bool = False,
         database: Optional[str] = None,
-        arrow: bool = True,
+        arrow: Union[str, bool] = True,
         arrow_disable_server_verification: bool = True,
         arrow_tls_root_certs: Optional[bytes] = None,
         bookmarks: Optional[Any] = None,
@@ -43,19 +44,20 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
             A username, password pair for database authentication.
         aura_ds : bool, default False
             A flag that indicates that that the client is used to connect
-            to a Neo4j Aura instance.
+            to a Neo4j AuraDS instance.
         database: Optional[str], default None
             The Neo4j database to query against.
-        arrow : bool, default True
-            A flag that indicates that the client should use Apache Arrow
-            for data streaming if it is available on the server.
+        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_disable_server_verification : bool, default True
-            A flag that indicates that, if the flight client is connecting with
-            TLS, that it skips server verification. If this is enabled, all
-            other TLS settings are overridden.
+            A flag that overrides other TLS settings and disables server verification for TLS connections.
         arrow_tls_root_certs : Optional[bytes], default None
-            PEM-encoded certificates that are used for the connecting to the
-            Arrow Flight server.
+            PEM-encoded certificates that are used for the connection to the
+            GDS Arrow Flight server.
         bookmarks : Optional[Any], default None
             The Neo4j bookmarks to require a certain state before the next query gets executed.
         """
@@ -76,6 +78,7 @@ class GraphDataScience(DirectEndpoints, UncallableNamespace):
                 self._query_runner.encrypted(),
                 arrow_disable_server_verification,
                 arrow_tls_root_certs,
+                None if arrow is True else arrow,
             )
 
         super().__init__(self._query_runner, "gds", self._server_version)

{graphdatascience-1.10a1 → graphdatascience-1.11a2}/graphdatascience/query_runner/arrow_graph_constructor.py

@@ -1,17 +1,17 @@
+from __future__ import annotations
+
 import concurrent
-import json
 import math
 import warnings
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, NoReturn, Optional
 
 import numpy
-import pyarrow.flight as flight
 from pandas import DataFrame
 from pyarrow import Table
 from tqdm.auto import tqdm
 
-from .arrow_endpoint_version import ArrowEndpointVersion
+from .gds_arrow_client import GdsArrowClient
 from .graph_constructor import GraphConstructor
 
 
@@ -20,9 +20,8 @@ class ArrowGraphConstructor(GraphConstructor):
         self,
         database: str,
         graph_name: str,
-        flight_client: flight.FlightClient,
+        flight_client: GdsArrowClient,
         concurrency: int,
-        arrow_endpoint_version: ArrowEndpointVersion,
         undirected_relationship_types: Optional[List[str]],
         chunk_size: int = 10_000,
     ):
@@ -30,7 +29,6 @@ class ArrowGraphConstructor(GraphConstructor):
         self._concurrency = concurrency
         self._graph_name = graph_name
         self._client = flight_client
-        self._arrow_endpoint_version = arrow_endpoint_version
         self._undirected_relationship_types = (
             [] if undirected_relationship_types is None else undirected_relationship_types
         )
@@ -47,20 +45,20 @@ class ArrowGraphConstructor(GraphConstructor):
             if self._undirected_relationship_types:
                 config["undirected_relationship_types"] = self._undirected_relationship_types
 
-            self._send_action(
+            self._client.send_action(
                 "CREATE_GRAPH",
                 config,
             )
 
             self._send_dfs(node_dfs, "node")
 
-            self._send_action("NODE_LOAD_DONE", {"name": self._graph_name})
+            self._client.send_action("NODE_LOAD_DONE", {"name": self._graph_name})
 
             self._send_dfs(relationship_dfs, "relationship")
 
-            self._send_action("RELATIONSHIP_LOAD_DONE", {"name": self._graph_name})
+            self._client.send_action("RELATIONSHIP_LOAD_DONE", {"name": self._graph_name})
         except (Exception, KeyboardInterrupt) as e:
-            self._send_action("ABORT", {"name": self._graph_name})
+            self._client.send_action("ABORT", {"name": self._graph_name})
 
             raise e
 
@@ -83,31 +81,20 @@ class ArrowGraphConstructor(GraphConstructor):
 
         return partitioned_dfs
 
-    def _send_action(self, action_type: str, meta_data: Dict[str, Any]) -> None:
-        action_type = self._versioned_action_type(action_type)
-        result = self._client.do_action(flight.Action(action_type, json.dumps(meta_data).encode("utf-8")))
-
-        # Consume result fully to sanity check and avoid cancelled streams
-        collected_result = list(result)
-        assert len(collected_result) == 1
-
-        json.loads(collected_result[0].body.to_pybytes().decode())
-
-    def _send_df(self, df: DataFrame, entity_type: str, pbar: tqdm) -> None:
+    def _send_df(self, df: DataFrame, entity_type: str, pbar: tqdm[NoReturn]) -> None:
         table = Table.from_pandas(df)
         batches = table.to_batches(self._chunk_size)
         flight_descriptor = {"name": self._graph_name, "entity_type": entity_type}
-        flight_descriptor = self._versioned_flight_desriptor(flight_descriptor)
 
-        # Write schema
-        upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(flight_descriptor).encode("utf-8"))
-        writer, _ = self._client.do_put(upload_descriptor, table.schema)
+        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()
 
     def _send_dfs(self, dfs: List[DataFrame], entity_type: str) -> None:
         desc = "Uploading Nodes" if entity_type == "node" else "Uploading Relationships"
@@ -122,17 +109,3 @@ class ArrowGraphConstructor(GraphConstructor):
                 if not future.exception():
                     continue
                 raise future.exception()  # type: ignore
-
-    def _versioned_action_type(self, action_type: str) -> str:
-        return self._arrow_endpoint_version.prefix() + action_type
-
-    def _versioned_flight_desriptor(self, flight_descriptor: Dict[str, Any]) -> Dict[str, Any]:
-        return (
-            flight_descriptor
-            if self._arrow_endpoint_version == ArrowEndpointVersion.ALPHA
-            else {
-                "name": "PUT_MESSAGE",
-                "version": ArrowEndpointVersion.V1.version(),
-                "body": flight_descriptor,
-            }
-        )

{graphdatascience-1.10a1 → graphdatascience-1.11a2}/graphdatascience/query_runner/arrow_query_runner.py

@@ -1,21 +1,14 @@
 from __future__ import annotations
 
-import base64
-import json
-import time
 import warnings
 from typing import Any, Dict, List, Optional, Tuple
 
-import pyarrow.flight as flight
 from pandas import DataFrame
-from pyarrow import ChunkedArray, Table, chunked_array
-from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
-from pyarrow.types import is_dictionary  # type: ignore
 
 from ..call_parameters import CallParameters
 from ..server_version.server_version import ServerVersion
-from .arrow_endpoint_version import ArrowEndpointVersion
 from .arrow_graph_constructor import ArrowGraphConstructor
+from .gds_arrow_client import GdsArrowClient
 from .graph_constructor import GraphConstructor
 from .query_runner import QueryRunner
 from graphdatascience.server_version.compatible_with import (
@@ -31,58 +24,31 @@ class ArrowQueryRunner(QueryRunner):
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
+        connection_string_override: Optional[str] = None,
     ) -> QueryRunner:
-        arrow_info = (
-            fallback_query_runner.call_procedure(endpoint="gds.debug.arrow", custom_error=False).squeeze().to_dict()
-        )
-        server_version = fallback_query_runner.server_version()
-        listen_address: str = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])
-        arrow_endpoint_version = ArrowEndpointVersion.from_arrow_info(arrow_info.get("versions", []))
-
-        if arrow_info["running"]:
-            return ArrowQueryRunner(
-                listen_address,
-                fallback_query_runner,
-                server_version,
-                auth,
-                encrypted,
-                disable_server_verification,
-                tls_root_certs,
-                arrow_endpoint_version,
-            )
-        else:
+        if not GdsArrowClient.is_arrow_enabled(fallback_query_runner):
             return fallback_query_runner
 
+        gds_arrow_client = GdsArrowClient.create(
+            fallback_query_runner,
+            auth,
+            encrypted,
+            disable_server_verification,
+            tls_root_certs,
+            connection_string_override,
+        )
+
+        return ArrowQueryRunner(gds_arrow_client, fallback_query_runner, fallback_query_runner.server_version())
+
     def __init__(
         self,
-        uri: str,
+        gds_arrow_client: GdsArrowClient,
         fallback_query_runner: QueryRunner,
         server_version: ServerVersion,
-        auth: Optional[Tuple[str, str]] = None,
-        encrypted: bool = False,
-        disable_server_verification: bool = False,
-        tls_root_certs: Optional[bytes] = None,
-        arrow_endpoint_version: ArrowEndpointVersion = ArrowEndpointVersion.ALPHA,
     ):
         self._fallback_query_runner = fallback_query_runner
+        self._gds_arrow_client = gds_arrow_client
         self._server_version = server_version
-        self._arrow_endpoint_version = arrow_endpoint_version
-
-        host, port_string = uri.split(":")
-
-        location = (
-            flight.Location.for_grpc_tls(host, int(port_string))
-            if encrypted
-            else flight.Location.for_grpc_tcp(host, int(port_string))
-        )
-
-        client_options: Dict[str, Any] = {"disable_server_verification": disable_server_verification}
-        if auth:
-            client_options["middleware"] = [AuthFactory(auth)]
-        if tls_root_certs:
-            client_options["tls_root_certs"] = tls_root_certs
-
-        self._flight_client = flight.FlightClient(location, **client_options)
 
     def warn_about_deprecation(self, old_endpoint: str, new_endpoint: str) -> None:
         warnings.warn(
@@ -135,7 +101,7 @@ class ArrowQueryRunner(QueryRunner):
                         old_endpoint="gds.graph.streamNodeProperty", new_endpoint="gds.graph.nodeProperty.stream"
                     )
 
-            return self._run_arrow_property_get(graph_name, endpoint, config)
+            return self._gds_arrow_client.get_property(self.database(), graph_name, endpoint, config)
         elif (
             old_endpoint := ("gds.graph.streamNodeProperties" == endpoint)
         ) or "gds.graph.nodeProperties.stream" == endpoint:
@@ -154,7 +120,8 @@ class ArrowQueryRunner(QueryRunner):
                     self.warn_about_deprecation(
                         old_endpoint="gds.graph.streamNodeProperties", new_endpoint="gds.graph.nodeProperties.stream"
                     )
-            return self._run_arrow_property_get(
+            return self._gds_arrow_client.get_property(
+                self.database(),
                 graph_name,
                 endpoint,
                 config,
@@ -175,7 +142,8 @@ class ArrowQueryRunner(QueryRunner):
                         old_endpoint="gds.graph.streamRelationshipProperty",
                         new_endpoint="gds.graph.relationshipProperty.stream",
                     )
-            return self._run_arrow_property_get(
+            return self._gds_arrow_client.get_property(
+                self.database(),
                 graph_name,
                 endpoint,
                 {"relationship_property": property_name, "relationship_types": relationship_types},
@@ -197,7 +165,8 @@ class ArrowQueryRunner(QueryRunner):
                         new_endpoint="gds.graph.relationshipProperties.stream",
                     )
 
-            return self._run_arrow_property_get(
+            return self._gds_arrow_client.get_property(
+                self.database(),
                 graph_name,
                 endpoint,
                 {"relationship_properties": property_names, "relationship_types": relationship_types},
@@ -224,7 +193,9 @@ class ArrowQueryRunner(QueryRunner):
                             new_endpoint="gds.graph.relationships.stream",
                         )
 
-            return self._run_arrow_property_get(graph_name, endpoint, {"relationship_types": relationship_types})
+            return self._gds_arrow_client.get_property(
+                self.database(), graph_name, endpoint, {"relationship_types": relationship_types}
+            )
 
         return self._fallback_query_runner.call_procedure(endpoint, params, yields, database, logging, custom_error)
 
@@ -254,52 +225,11 @@ class ArrowQueryRunner(QueryRunner):
 
     def close(self) -> None:
         self._fallback_query_runner.close()
-        # PyArrow 7 did not expose a close method yet
-        if hasattr(self._flight_client, "close"):
-            self._flight_client.close()
+        self._gds_arrow_client.close()
 
     def fallback_query_runner(self) -> QueryRunner:
         return self._fallback_query_runner
 
-    def _run_arrow_property_get(self, graph_name: str, procedure_name: str, configuration: Dict[str, Any]) -> DataFrame:
-        if not self.database():
-            raise ValueError(
-                "For this call you must have explicitly specified a valid Neo4j database to execute on, "
-                "using `GraphDataScience.set_database`."
-            )
-
-        payload = {
-            "database_name": self.database(),
-            "graph_name": graph_name,
-            "procedure_name": procedure_name,
-            "configuration": configuration,
-        }
-
-        if self._arrow_endpoint_version == ArrowEndpointVersion.V1:
-            payload = {
-                "name": "GET_MESSAGE",
-                "version": ArrowEndpointVersion.V1.version(),
-                "body": payload,
-            }
-
-        ticket = flight.Ticket(json.dumps(payload).encode("utf-8"))
-        get = self._flight_client.do_get(ticket)
-        arrow_table = get.read_all()
-
-        if configuration.get("list_node_labels", False):
-            # GDS 2.5 had an inconsistent naming of the node labels column
-            new_colum_names = ["nodeLabels" if i == "labels" else i for i in arrow_table.column_names]
-            arrow_table = arrow_table.rename_columns(new_colum_names)
-
-        # Pandas 2.2.0 deprecated an API used by ArrowTable.to_pandas() (< pyarrow 15.0)
-        warnings.filterwarnings(
-            "ignore",
-            category=DeprecationWarning,
-            message=r"Passing a BlockManager to DataFrame is deprecated",
-        )
-
-        return self._sanitize_arrow_table(arrow_table).to_pandas()  # type: ignore
-
     def create_graph_constructor(
         self, graph_name: str, concurrency: int, undirected_relationship_types: Optional[List[str]]
     ) -> GraphConstructor:
@@ -313,75 +243,7 @@ class ArrowQueryRunner(QueryRunner):
         return ArrowGraphConstructor(
             database,
             graph_name,
-            self._flight_client,
+            self._gds_arrow_client,
             concurrency,
-            self._arrow_endpoint_version,
             undirected_relationship_types,
         )
-
-    def _sanitize_arrow_table(self, arrow_table: Table) -> Table:
-        dict_encoded_fields = [
-            (idx, field) for idx, field in enumerate(arrow_table.schema) if is_dictionary(field.type)
-        ]
-        for idx, field in dict_encoded_fields:
-            try:
-                field.type.to_pandas_dtype()
-            except NotImplementedError:
-                # we need to decode the dictionary column before transforming to pandas
-                if isinstance(arrow_table[field.name], ChunkedArray):
-                    decoded_col = chunked_array([chunk.dictionary_decode() for chunk in arrow_table[field.name].chunks])
-                else:
-                    decoded_col = arrow_table[field.name].dictionary_decode()
-                arrow_table = arrow_table.set_column(idx, field.name, decoded_col)
-        return arrow_table
-
-
-class AuthFactory(ClientMiddlewareFactory):  # type: ignore
-    def __init__(self, auth: Tuple[str, str], *args: Any, **kwargs: Any) -> None:
-        super().__init__(*args, **kwargs)
-        self._auth = auth
-        self._token: Optional[str] = None
-        self._token_timestamp = 0
-
-    def start_call(self, info: Any) -> "AuthMiddleware":
-        return AuthMiddleware(self)
-
-    def token(self) -> Optional[str]:
-        # check whether the token is older than 10 minutes. If so, reset it.
-        if self._token and int(time.time()) - self._token_timestamp > 600:
-            self._token = None
-
-        return self._token
-
-    def set_token(self, token: str) -> None:
-        self._token = token
-        self._token_timestamp = int(time.time())
-
-    @property
-    def auth(self) -> Tuple[str, str]:
-        return self._auth
-
-
-class AuthMiddleware(ClientMiddleware):  # type: ignore
-    def __init__(self, factory: AuthFactory, *args: Any, **kwargs: Any) -> None:
-        super().__init__(*args, **kwargs)
-        self._factory = factory
-
-    def received_headers(self, headers: Dict[str, Any]) -> None:
-        auth_header: str = headers.get("Authorization", None)
-        if not auth_header:
-            return
-        [auth_type, token] = auth_header.split(" ", 1)
-        if auth_type == "Bearer":
-            self._factory.set_token(token)
-
-    def sending_headers(self) -> Dict[str, str]:
-        token = self._factory.token()
-        if not token:
-            username, password = self._factory.auth
-            auth_token = f"{username}:{password}"
-            auth_token = "Basic " + base64.b64encode(auth_token.encode("utf-8")).decode("ASCII")
-            # There seems to be a bug, `authorization` must be lower key
-            return {"authorization": auth_token}
-        else:
-            return {"authorization": "Bearer " + token}