使用 Apache Arrow 加载和回流数据

Colab Google Colab 中的笔记本中跟着做

安装在 AuraDS 上的 GDS 企业版包含一个 Arrow Flight 服务器,默认情况下已配置并运行。Arrow Flight 服务器可加快数据密集型进程的速度,例如

  • 直接从内存中数据创建图。

  • 流式传输节点和关系属性。

  • 流式传输图的关系拓扑。

有两种方法可以使用 GDS 的 Arrow Flight 服务器

  1. 使用 GDS Python 客户端,其中包含一个 Arrow Flight 客户端。

  2. GDS 手册 中所述,实现一个自定义 Arrow Flight 客户端。

在以下示例中,我们使用 GDS 客户端,因为它是最方便的选择。所有加载和流式传输方法都可以在没有 Arrow 的情况下使用,但如果有 Arrow,效率更高。

设置

%pip install 'graphdatascience>=1.7'

from graphdatascience import GraphDataScience

# Replace with the actual connection URI and credentials
AURA_CONNECTION_URI = "neo4j+s://xxxxxxxx.databases.neo4j.io"
AURA_USERNAME = "neo4j"
AURA_PASSWORD = ""

# When initialized, the client tries to use Arrow if it is available on the server.
# This behaviour is controlled by the `arrow` parameter, which is set to `True` by default.
gds = GraphDataScience(AURA_CONNECTION_URI, auth=(AURA_USERNAME, AURA_PASSWORD), aura_ds=True)

# Necessary if Arrow is enabled (as is by default on Aura)
gds.set_database("neo4j")

您可以调用 gds.debug.arrow() 方法以验证 Arrow 是否已启用并正在运行

gds.debug.arrow()

加载数据

您可以使用 gds.graph.construct 客户端方法将数据直接加载到图中。

数据必须是 Pandas DataFrame,因此我们需要安装并导入 pandas 库。

%pip install pandas

import pandas as pd

然后,我们可以按照以下示例创建图。每个具有必需列的 DataFrame 的格式在 GDS 手册 中指定。

nodes = pd.DataFrame(
    {
        "nodeId": [0, 1, 2],
        "labels":  ["Article", "Article", "Article"],
        "pages": [3, 7, 12],
    }
)

relationships = pd.DataFrame(
    {
        "sourceNodeId": [0, 1],
        "targetNodeId": [1, 2],
        "relationshipType": ["CITES", "CITES"],
        "times": [2, 1]
    }
)

article_graph = gds.graph.construct(
    "article-graph",
    nodes,
    relationships
)

现在,我们可以检查图是否已创建

gds.graph.list()

流式传输节点和关系属性

创建图后,您可以 作为流 读取节点和关系属性。

# Read all the values for the node property `pages`
gds.graph.nodeProperties.stream(article_graph, "pages")
# Read all the values for the relationship property `times`
gds.graph.relationshipProperties.stream(article_graph, "times")

性能

要查看 Arrow 可用时的性能差异,我们可以测量将数据集加载到图中所需的时间。在本示例中,我们使用内置的 OGBN 数据集,因此我们需要安装 ogb 扩展。

%pip install 'graphdatascience[ogb]>=1.7'

# Load and immediately drop the dataset to download and cache the data
ogbn_arxiv = gds.graph.ogbn.load("ogbn-arxiv")
ogbn_arxiv.drop()

然后,我们可以计时加载过程。在 8 GB AuraDS 实例上,这应该需要不到 30 秒。

%%timeit -n 1 -r 1

# This call uses the cached dataset, so only the actual loading is timed
ogbn_arxiv = gds.graph.ogbn.load("ogbn-arxiv")

通过在 GraphDataScience 构造函数中添加 arrow=False 来禁用 Arrow,相同的加载过程将需要超过 1 分钟。因此,对于此数据集,Arrow 至少提供了 2 倍的速度提升。

清理

article_graph.drop()
ogbn_arxiv.drop()

gds.close()