Fixed ISSUE-1406 (#1407)

* Fixed ISSUE-1406 where edge data frames were having properties set to 'set' cardinality which is not supported.  Version prior to 1.1.1.0 were not throwing an error

* Fixed formatting issue and updated SPARQL examples in tutorial with more accurate examples

* Fixed validation issue with the client needing to ensure something is a string

* Fixed issue with Flake8 validation

* Minor - Remove quarantine in Neptune tests

Co-authored-by: Dave Bechberger <[email protected]>
Co-authored-by: Lucas Hanson <[email protected]>
Co-authored-by: Abdel Jaidi <[email protected]>

Author: 3 people

awswrangler/neptune/client.py

@@ -249,7 +249,7 @@ def _execute_sparql(self, query: str, headers: Any) -> Any:
 
         s = SPARQLWrapper("")
         s.setQuery(query)
-        query_type = s.queryType.upper()
+        query_type = str(s.queryType).upper()
         if query_type in ["SELECT", "CONSTRUCT", "ASK", "DESCRIBE"]:
             data = {"query": query}
         else:

awswrangler/neptune/neptune.py

@@ -7,6 +7,7 @@
 from typing import Any, Callable, TypeVar
 
 import pandas as pd
+from gremlin_python.process.anonymous_traversal import traversal
 from gremlin_python.process.graph_traversal import GraphTraversalSource, __
 from gremlin_python.process.translator import Translator
 from gremlin_python.process.traversal import Cardinality, T
@@ -175,7 +176,7 @@ def to_property_graph(
     ... )
     """
     # check if ~id and ~label column exist and if not throw error
-    g = Graph().traversal()
+    g = traversal().withGraph(Graph())
     is_edge_df = False
     is_update_df = True
     if "~id" in df.columns:
@@ -307,18 +308,23 @@ def _get_column_name(column: str) -> str:
     return column
 
 
-def _set_properties(g: GraphTraversalSource, use_header_cardinality: bool, row: Any) -> GraphTraversalSource:
+def _set_properties(
+    g: GraphTraversalSource, use_header_cardinality: bool, row: Any, ignore_cardinality: bool = False
+) -> GraphTraversalSource:
     for (column, value) in row.items():
         if column not in ["~id", "~label", "~to", "~from"]:
-            # If the column header is specifying the cardinality then use it
-            if use_header_cardinality:
-                if column.lower().find("(single)") > 0 and pd.notna(value):
-                    g = g.property(Cardinality.single, _get_column_name(column), value)
-                else:
-                    g = _expand_properties(g, _get_column_name(column), value)
+            if ignore_cardinality and pd.notna(value):
+                g = g.property(_get_column_name(column), value)
             else:
-                # If not using header cardinality then use the default of set
-                g = _expand_properties(g, column, value)
+                # If the column header is specifying the cardinality then use it
+                if use_header_cardinality:
+                    if column.lower().find("(single)") > 0 and pd.notna(value):
+                        g = g.property(Cardinality.single, _get_column_name(column), value)
+                    else:
+                        g = _expand_properties(g, _get_column_name(column), value)
+                else:
+                    # If not using header cardinality then use the default of set
+                    g = _expand_properties(g, column, value)
     return g
 
 
@@ -361,7 +367,7 @@ def _build_gremlin_insert_edges(
             .coalesce(__.unfold(), _build_gremlin_insert_vertices(__, {"~id": row["~to"], "~label": "Vertex"}))
         )
     )
-    g = _set_properties(g, use_header_cardinality, row)
+    g = _set_properties(g, use_header_cardinality, row, ignore_cardinality=True)
 
     return g
 
@@ -370,6 +376,9 @@ def _run_gremlin_insert(client: NeptuneClient, g: GraphTraversalSource) -> bool:
     translator = Translator("g")
     s = translator.translate(g.bytecode)
     s = s.replace("Cardinality.", "")  # hack to fix parser error for set cardinality
+    s = s.replace(
+        ".values('shape')", ""
+    )  # hack to fix parser error for adding unknown values('shape') steps to translation.
     _logger.debug(s)
     res = client.write_gremlin(s)
     return res

tests/test_neptune.py

@@ -170,7 +170,19 @@ def test_sparql_query(neptune_endpoint, neptune_port) -> Dict[str, Any]:
     assert df.shape == (2, 3)
 
 
-@pytest.mark.skip(reason="In quarantine due to breaking change in Neptune engine.")
+def test_write_vertex_property_nan(neptune_endpoint, neptune_port) -> Dict[str, Any]:
+    client = wr.neptune.connect(neptune_endpoint, neptune_port, iam_enabled=False)
+    id = uuid.uuid4()
+    wr.neptune.execute_gremlin(client, f"g.addV().property(T.id, '{str(id)}')")
+
+    data = [_create_dummy_edge(), _create_dummy_edge()]
+    del data[1]["str"]
+    data[1]["int"] = np.nan
+    df = pd.DataFrame(data)
+    res = wr.neptune.to_property_graph(client, df)
+    assert res
+
+
 def test_gremlin_write_different_cols(neptune_endpoint, neptune_port) -> Dict[str, Any]:
     client = wr.neptune.connect(neptune_endpoint, neptune_port, iam_enabled=False)
     id = uuid.uuid4()
@@ -293,7 +305,6 @@ def test_gremlin_write_vertices(neptune_endpoint, neptune_port) -> Dict[str, Any
     assert len(saved_row["str"]) == 2
 
 
-@pytest.mark.skip(reason="In quarantine due to breaking change in Neptune engine.")
 def test_gremlin_write_edges(neptune_endpoint, neptune_port) -> Dict[str, Any]:
     client = wr.neptune.connect(neptune_endpoint, neptune_port, iam_enabled=False)
 

tutorials/033 - Amazon Neptune.ipynb

@@ -27,29 +27,29 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "outputs": [],
-   "source": [
-    "!pip install awswrangler[sparql]"
-   ],
    "metadata": {
     "collapsed": false,
     "pycharm": {
      "name": "#%%\n"
     }
-   }
+   },
+   "outputs": [],
+   "source": [
+    "!pip install awswrangler[sparql]"
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false
+   },
    "source": [
     "## Initialize\n",
     "\n",
     "The first step to using AWS Data Wrangler with Amazon Neptune is to import the library and create a client connection.\n",
     "\n",
     "<div style=\"background-color:#eeeeee; padding:10px; text-align:left; border-radius:10px; margin-top:10px; margin-bottom:10px; \"><b>Note</b>: Connecting to Amazon Neptune requires that the application you are running has access to the Private VPC where Neptune is located.  Without this access you will not be able to connect using AWS Data Wrangler.</div>"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
+   ]
   },
   {
    "cell_type": "code",
@@ -124,7 +124,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "query = \"SELECT ?s ?o WHERE { ?s ?p ?o .} LIMIT 5\"\n",
+    "query = \"\"\"\n",
+    "        PREFIX foaf: <https://xmlns.com/foaf/0.1/> \n",
+    "        PREFIX ex: <https://www.example.com/> \n",
+    "        SELECT ?firstName WHERE { ex:JaneDoe foaf:knows ?person . ?person foaf:firstName ?firstName }\"\"\"\n",
     "df = wr.neptune.execute_sparql(client, query)\n",
     "display(df.head(5))"
    ]
@@ -316,15 +319,17 @@
    "source": [
     "def _create_dummy_triple():\n",
     "    data = dict()\n",
-    "    data[\"s\"] = \"foo\"\n",
+    "    data[\"s\"] = \"http://example.com/resources/foo\"\n",
     "    data[\"p\"] = uuid.uuid4()\n",
     "    data[\"o\"] = random.randint(0, 1000)\n",
     "    return data\n",
     "\n",
     "data = [_create_dummy_triple(), _create_dummy_triple(), _create_dummy_triple()]\n",
     "df = pd.DataFrame(data)\n",
     "res = wr.neptune.to_rdf_graph(client, df)\n",
-    "query = \"SELECT ?o WHERE { <foo> <\" + str(data[0]['p']) + \"> ?o .}\"\n",
+    "query = \"\"\"\n",
+    "    PREFIX foo: <http://example.com/resources/>\n",
+    "    SELECT ?o WHERE { <foo:foo> <\" + str(data[0]['p']) + \"> ?o .}\"\"\"\n",
     "df = wr.neptune.execute_sparql(client, query)\n",
     "display(df)"
    ]
@@ -352,7 +357,9 @@
     "data = [_create_dummy_quad(), _create_dummy_quad(), _create_dummy_quad()]\n",
     "df = pd.DataFrame(data)\n",
     "res = wr.neptune.to_rdf_graph(client, df)\n",
-    "query = \"SELECT ?o WHERE { <foo> <\" + str(data[0]['p']) + \"> ?o .}\"\n",
+    "query = \"\"\"\n",
+    "    PREFIX foo: <http://example.com/resources/>\n",
+    "    SELECT ?o WHERE { <foo:foo> <\" + str(data[0]['p']) + \"> ?o .}\"\"\"\n",
     "df = wr.neptune.execute_sparql(client, query)\n",
     "display(df)"
    ]
@@ -542,4 +549,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}