fix e2e test

wey-gu · Sep 6, 2023 · 8457c94 · 8457c94
1 parent 2977ec0
commit 8457c94
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 5 deletions.
diff --git a/examples/spark_engine.ipynb b/examples/spark_engine.ipynb
@@ -628,15 +628,14 @@
    "outputs": [],
    "source": [
     "writer = NebulaWriter(\n",
-    "    data=df_result, sink=\"nebulagraph_vertex\", config=config, engine=\"spark\"\n",
+    "    data=df_result, sink=\"nebulagraph_edge\", config=config, engine=\"spark\"\n",
     ")\n",
     "\n",
     "# map column louvain into property cluster_id\n",
     "properties = {\"similarity\": \"similarity\"}\n",
     "\n",
     "writer.set_options(\n",
     "    space=\"basketballplayer\",\n",
-    "    type=\"edge\",\n",
     "    edge_type=\"jaccard_similarity\",\n",
     "    src_id=\"srcId\",\n",
     "    dst_id=\"dstId\",\n",

diff --git a/ng_ai/nebula_writer.py b/ng_ai/nebula_writer.py
@@ -320,6 +320,9 @@ def set_options(self, **kwargs):
 
         # case switch based on sink
         if self.sink in SPARK_NEBULA_SINKS:
+            kwargs["type"] = (
+                "vertex" if self.sink == "nebulagraph_vertex" else "edge"
+            )
             self._set_options_with_nebula(**kwargs)
         elif self.sink in SPARK_FILE_SINKS:
             self._set_options_with_file(**kwargs)

diff --git a/tests/integration/notebooks/spark_engine_e2e.ipynb b/tests/integration/notebooks/spark_engine_e2e.ipynb
@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "8a6e2678",
    "metadata": {},
@@ -42,6 +43,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "170f8c38",
    "metadata": {},
@@ -70,6 +72,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "6051622d",
    "metadata": {},
@@ -116,6 +119,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "311b49c4",
    "metadata": {},
@@ -153,6 +157,7 @@
     "properties = {\"lpa\": \"cluster_id\"}\n",
     "\n",
     "writer.set_options(\n",
+    "    space=\"basketballplayer\",\n",
     "    tag=\"label_propagation\",\n",
     "    vid_field=\"_id\",\n",
     "    properties=properties,\n",
@@ -162,6 +167,47 @@
     "# write back to NebulaGraph\n",
     "writer.write()"
    ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "cbe44f42",
+   "metadata": {},
+   "source": [
+    "## Test Writer insert mode in sink: `nebulagraph_edge`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "79068641",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_result = df.algo.jaccard_similarity()\n",
+    "\n",
+    "writer = NebulaWriter(\n",
+    "    data=df_result, sink=\"nebulagraph_edge\", config=config, engine=\"spark\"\n",
+    ")\n",
+    "\n",
+    "# map column louvain into property cluster_id\n",
+    "properties = {\"similarity\": \"similarity\"}\n",
+    "\n",
+    "writer.set_options(\n",
+    "    space=\"basketballplayer\",\n",
+    "    edge_type=\"jaccard_similarity\",\n",
+    "    src_id=\"srcId\",\n",
+    "    dst_id=\"dstId\",\n",
+    "    src_id_policy=\"\",\n",
+    "    dst_id_policy=\"\",\n",
+    "    properties=properties,\n",
+    "    batch_size=256,\n",
+    "    write_mode=\"insert\",\n",
+    ")\n",
+    "\n",
+    "# write back to NebulaGraph\n",
+    "writer.write()"
+   ]
   }
  ],
  "metadata": {

diff --git a/tests/integration/spark_engine_cases/writer.py b/tests/integration/spark_engine_cases/writer.py
@@ -43,15 +43,14 @@
 df_result = df.algo.jaccard()
 
 writer = NebulaWriter(
-    data=df_result, sink="nebulagraph_vertex", config=config, engine="spark"
+    data=df_result, sink="nebulagraph_edge", config=config, engine="spark"
 )
 
 # map column louvain into property cluster_id
 properties = {"similarity": "similarity"}
 
 writer.set_options(
     space="basketballplayer",
-    type="edge",
     edge_type="jaccard_similarity",
     src_id="srcId",
     dst_id="dstId",

diff --git a/tests/integration/test_e2e_spark_engine.py b/tests/integration/test_e2e_spark_engine.py
@@ -107,7 +107,6 @@ def test_label_propagation_spark_engine_writer():
             "MATCH (v:player) RETURN v.label_propagation.cluster_id LIMIT 1"
         )
         print(result)
-    connection_pool.close()
 
     assert result.is_succeeded(), f"ERROR during query NebulaGraph: {result}"
     assert (