Merge pull request #103 from awslabs/athena

Athena

Author: igorborgest

awswrangler/pandas.py

@@ -583,10 +583,15 @@ def _read_sql_athena_ctas(self,
         manifest_path: str = f"{s3_output}/tables/{query_id}-manifest.csv"
         paths: List[str] = self._session.athena.extract_manifest_paths(path=manifest_path)
         logger.debug(f"paths: {paths}")
-        return self.read_parquet(path=paths,
-                                 procs_cpu_bound=procs_cpu_bound,
-                                 wait_objects=True,
-                                 wait_objects_timeout=15.0)
+        if not paths:
+            df: pd.DataFrame = pd.DataFrame()
+        else:
+            df = self.read_parquet(path=paths,
+                                   procs_cpu_bound=procs_cpu_bound,
+                                   wait_objects=True,
+                                   wait_objects_timeout=15.0)
+        self._session.s3.delete_listed_objects(objects_paths=[manifest_path] + paths)
+        return df
 
     def _read_sql_athena_regular(self,
                                  sql: str,

testing/test_awswrangler/test_pandas.py

@@ -2020,3 +2020,17 @@ def test_read_csv_prefix_iterator(bucket, sample, row_num):
         total_count += count
     wr.s3.delete_listed_objects(objects_paths=paths)
     assert total_count == row_num * n
+
+
+@pytest.mark.parametrize("ctas_approach", [False, True])
+def test_read_sql_athena_empty(ctas_approach):
+    sql = """
+        WITH dataset AS (
+          SELECT 0 AS id
+        )
+        SELECT id
+        FROM dataset
+        WHERE id != 0
+    """
+    df = wr.pandas.read_sql_athena(sql=sql, ctas_approach=ctas_approach)
+    print(df)