Improve notebook descriptions and address PR comments

Author: DarthMax

examples/graph-analytics-serverless-spark.ipynb

@@ -40,9 +40,7 @@
    "source": [
     "## Prerequisites\n",
     "\n",
-    "This notebook requires having an AuraDB instance available and have the Aura Graph Analytics [feature](https://neo4j.com/docs/aura/graph-analytics/#aura-gds-serverless) enabled for your project.\n",
-    "\n",
-    "We also need to have the `graphdatascience` Python library installed, version `1.18` or later, as well as `pyspark`. For more information about setting up pyspark visit https://spark.apache.org/docs/latest/api/python/getting_started/"
+    "We also need to have the `graphdatascience` Python library installed, version `1.18` or later, as well as `pyspark`."
    ]
   },
   {
@@ -75,7 +73,7 @@
     "### Connecting to a Spark Session\n",
     "\n",
     "To interact with the Spark cluster we need to first instantiate a Spark session. In this example we will use a local Spark session, which will run Spark on the same machine.\n",
-    "Working with a remote Spark cluster will work similarly."
+    "Working with a remote Spark cluster will work similarly. For more information about setting up pyspark visit https://spark.apache.org/docs/latest/api/python/getting_started/"
    ]
   },
   {
@@ -164,8 +162,7 @@
    "source": [
     "## Adding a dataset\n",
     "\n",
-    "As the next step we will setup a dataset in Spark. In this example we will use the New York Bike trips dataset (https://www.kaggle.com/datasets/gabrielramos87/bike-trips).",
-    "The bike trips form a graph where nodes represent bike renting stations and relationships represent start and end points for a bike rental trip."
+    "As the next step we will setup a dataset in Spark. In this example we will use the New York Bike trips dataset (https://www.kaggle.com/datasets/gabrielramos87/bike-trips). The bike trips form a graph where nodes represent bike renting stations and relationships represent start and end points for a bike rental trip."
    ]
   },
   {
@@ -206,23 +203,37 @@
     "\n",
     "We first need to get access to the GDSArrowClient. This client allows us to directly communicate with the Arrow Flight server provided by the session.\n",
     "\n",
-    "Our input data already resembles edge triplets, where each of the rows represents an edge from a source station to a target station. This allows us to use the arrows servers graph import from triplets functionality, which requires the following protocol:\n",
+    "Our input data already resembles triplets, where each row represents an edge from a source station to a target station. This allows us to use the Arrow Server's \"graph import from triplets\" functionality, which requires the following protocol:\n",
     "\n",
     "1. Send an action `v2/graph.project.fromTriplets`\n",
     "   This will initialize the import process and allows us to specify the graph name, and settings like `undirected_relationship_types`. It returns a job id, that we need to reference the import job in the following steps.\n",
     "2. Send the data in batches to the Arrow server.\n",
     "3. Send another action called `v2/graph.project.fromTriplets.done` to tell the import process that no more data will be sent. This will trigger the final graph creation inside the GDS session.\n",
     "4. Wait for the import process to reach the `DONE` state.\n",
     "\n",
-    "While the overall process is straight forward, we need to somehow tell Spark to"
+    "The most complicated step here is to run the actual data upload on each spark worker. We will use the `mapInArrow` function to run custom code on each spark worker. Each worker will receive a number of arrow record batches that we can directly send to the GDS session's Arrow server. "
+   ]
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": [
+    "The user wants to add a 1-second delay (sleep) within the loop that waits for the import job to finish. This requires importing the `time` module and adding `time.sleep(1)` inside the `while` loop at the end of the cell.\n",
+    "\n"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
    "metadata": {},
+   "cell_type": "markdown",
+   "source": "<llm-snippet-file>graph-analytics-serverless-spark.ipynb</llm-snippet-file>\n"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
    "outputs": [],
+   "execution_count": null,
    "source": [
+    "import time\n",
     "import pandas as pd\n",
     "import pyarrow\n",
     "from pyspark.sql import functions\n",
@@ -244,30 +255,33 @@
     "\n",
     "# Select the source target pairs from our source data\n",
     "source_target_pairs = spark.sql(\"\"\"\n",
-    "  SELECT start_station_id AS sourceNode, end_station_id AS targetNode\n",
-    "  FROM bike_trips\n",
-    "\"\"\")\n",
+    "                                SELECT start_station_id AS sourceNode, end_station_id AS targetNode\n",
+    "                                FROM bike_trips\n",
+    "                                \"\"\")\n",
     "\n",
     "# 2. Use the `mapInArrow` function to upload the data to the GDS session. Returns a DataFrame with a single column containing the batch sizes.\n",
     "uploaded_batches = source_target_pairs.mapInArrow(upload_batch, \"batch_rows_imported long\")\n",
     "\n",
     "# Aggregate the batch sizes to receive the row count.\n",
-    "uploaded_batches.agg(functions.sum(\"batch_rows_imported\").alias(\"rows_imported\")).show()\n",
+    "aggregated_batch_sizes = uploaded_batches.agg(functions.sum(\"batch_rows_imported\").alias(\"rows_imported\"))\n",
+    "\n",
+    "# Show the result. This will trigger the computation and thus run the data upload.\n",
+    "aggregated_batch_sizes.show()\n",
     "\n",
     "# 3. Finish the import process\n",
     "arrow_client.triplet_load_done(job_id)\n",
     "\n",
     "# 4. Wait for the import to finish\n",
     "while not arrow_client.job_status(job_id).succeeded():\n",
-    "    pass\n",
+    "    time.sleep(1)\n",
     "\n",
     "G = gds.v2.graph.get(graph_name)\n",
     "G"
    ]
   },
   {
-   "cell_type": "markdown",
    "metadata": {},
+   "cell_type": "markdown",
    "source": [
     "## Running Algorithms\n",
     "\n",
@@ -322,7 +336,7 @@
     "# Optional: Repartition the data to make sure it is distributed equally\n",
     "result = received_batches.repartition(numPartitions=spark.sparkContext.defaultParallelism)\n",
     "\n",
-    "result.show()"
+    "result.toPandas()"
    ]
   },
   {

src/graphdatascience/arrow_client/v2/api_types.py

@@ -34,7 +34,7 @@ def sub_tasks(self) -> str | None:
         return None
 
     def aborted(self) -> bool:
-        return self.status == "Aborted"
+        return self.status.lower() == "aborted"
 
     def succeeded(self) -> bool:
         return self.status.lower() == "done"