fix: column names and apply map (#2492)

* fix column names and apply map
---------

Co-authored-by: Julian Grueber <[email protected]>

Author: juliangrueber

awswrangler/neptune/_neptune.py

@@ -110,8 +110,8 @@ def execute_sparql(client: NeptuneClient, query: str) -> pd.DataFrame:
     data = client.read_sparql(query)
     df = None
     if "results" in data and "bindings" in data["results"]:
-        df = pd.DataFrame(data["results"]["bindings"])
-        df.applymap(lambda x: x["value"])
+        df = pd.DataFrame(data["results"]["bindings"], columns=data.get("head", {}).get("vars"))
+        df = df.applymap(lambda d: d["value"] if "value" in d else None)
     else:
         df = pd.DataFrame(data)
 

tests/_utils.py

@@ -578,3 +578,9 @@ def assert_pandas_equals(df1: Union[pd.DataFrame, pd.Series], df2: Union[pd.Data
         assert_series_equal(df1, df2)
     else:
         raise ValueError(f"Unsupported type {type(df1)}")
+
+
+def assert_columns_in_pandas_data_frame(df: pd.DataFrame, columns: List[str]) -> None:
+    """Check data frame for columns"""
+    for col in columns:
+        assert col in df.columns

tests/unit/test_neptune.py

@@ -12,7 +12,7 @@
 import awswrangler.pandas as pd
 from awswrangler.neptune._client import BulkLoadParserConfiguration
 
-from .._utils import extract_cloudformation_outputs
+from .._utils import assert_columns_in_pandas_data_frame, extract_cloudformation_outputs
 
 logging.getLogger("awswrangler").setLevel(logging.DEBUG)
 
@@ -160,9 +160,15 @@ def test_sparql_query(neptune_endpoint, neptune_port) -> Dict[str, Any]:
     df = wr.neptune.execute_sparql(client, "INSERT DATA { <test1> <test1> <test1>}")
     df = wr.neptune.execute_sparql(client, "SELECT ?s ?p ?o {?s ?p ?o} LIMIT 1")
     assert df.shape == (1, 3)
+    assert_columns_in_pandas_data_frame(df, ["s", "p", "o"])
 
     df = wr.neptune.execute_sparql(client, "SELECT ?s ?p ?o {?s ?p ?o} LIMIT 2")
     assert df.shape == (2, 3)
+    assert_columns_in_pandas_data_frame(df, ["s", "p", "o"])
+
+    df = wr.neptune.execute_sparql(client, "SELECT ?s ?p ?o {?s ?p ?o} LIMIT 0")
+    assert df.shape == (0, 3)
+    assert_columns_in_pandas_data_frame(df, ["s", "p", "o"])
 
 
 def test_write_vertex_property_nan(neptune_endpoint, neptune_port) -> Dict[str, Any]: