Merge pull request #120 from awslabs/pandas-read-redshift-empty

Handling empty results for Pandas.read_sql_redshift()

Author: igorborgest

awswrangler/pandas.py

@@ -1408,6 +1408,8 @@ def read_parquet(self,
         :param wait_objects_timeout: Wait objects Timeout (seconds)
         :return: Pandas DataFrame
         """
+        if type(path) == list and not path:  # Empty list
+            return pd.DataFrame()
         procs_cpu_bound = procs_cpu_bound if procs_cpu_bound is not None else self._session.procs_cpu_bound if self._session.procs_cpu_bound is not None else 1
         logger.debug(f"procs_cpu_bound: {procs_cpu_bound}")
         dfs: List[pd.DataFrame] = []
@@ -1640,7 +1642,10 @@ def read_sql_redshift(self,
                                                       iam_role=iam_role,
                                                       connection=connection)
             logger.debug(f"paths: {paths}")
-            df: pd.DataFrame = self.read_parquet(path=paths, procs_cpu_bound=procs_cpu_bound)  # type: ignore
+            if paths:
+                df: pd.DataFrame = self.read_parquet(path=paths, procs_cpu_bound=procs_cpu_bound)  # type: ignore
+            else:
+                df = pd.DataFrame()
         except Exception as ex:
             connection.rollback()
             if paths is not None:

awswrangler/redshift.py

@@ -455,8 +455,9 @@ def to_parquet(self,
         logger.debug(f"paths: {paths}")
         connection.commit()
         cursor.close()
-        if manifest_str != "":
-            self._session.s3.wait_object_exists(path=f"{path}manifest", timeout=30.0)
-        for p in paths:
-            self._session.s3.wait_object_exists(path=p, timeout=30.0)
+        if paths:
+            if manifest_str != "":
+                self._session.s3.wait_object_exists(path=f"{path}manifest", timeout=30.0)
+            for p in paths:
+                self._session.s3.wait_object_exists(path=p, timeout=30.0)
         return paths

testing/test_awswrangler/test_pandas.py

@@ -2353,13 +2353,15 @@ def test_s3_overall_nan(bucket, database):
     df["c_int"] = df["c_int"].astype("Int64")
     print(df)
     wr.pandas.to_parquet(dataframe=df,
-                     database=database,
-                     path=path,
-                     mode="overwrite",
-                     preserve_index=False,
-                     procs_cpu_bound=1,
-                     inplace=False)
+                         database=database,
+                         path=path,
+                         mode="overwrite",
+                         preserve_index=False,
+                         procs_cpu_bound=1,
+                         inplace=False)
     sleep(15)
-    df2 = wr.pandas.read_sql_athena(database=database, sql="SELECT * FROM test_s3_overall_nan ORDER BY id", ctas_approach=True)
+    df2 = wr.pandas.read_sql_athena(database=database,
+                                    sql="SELECT * FROM test_s3_overall_nan ORDER BY id",
+                                    ctas_approach=True)
     wr.s3.delete_objects(path=path)
     assert df.equals(df2)

testing/test_awswrangler/test_redshift.py

@@ -752,3 +752,26 @@ def test_read_sql_redshift_pandas_glue_conn(session, bucket, redshift_parameters
                                            temp_s3_path=path2)
     assert len(df.index) == len(df2.index)
     assert len(df.columns) + 1 == len(df2.columns)
+
+
+def test_read_sql_redshift_pandas_empty(session, bucket, redshift_parameters):
+    path = f"s3://{bucket}/test_read_sql_redshift_pandas_empty/"
+    path2 = f"s3://{bucket}/test_read_sql_redshift_pandas_empty2/"
+    wr.s3.delete_objects(path=path)
+    wr.s3.delete_objects(path=path2)
+    df = pd.DataFrame({"id": [1, 2, 3, 4, 5]})
+    session.pandas.to_redshift(
+        dataframe=df,
+        path=path,
+        schema="public",
+        table="test",
+        connection="aws-data-wrangler-redshift",
+        iam_role=redshift_parameters.get("RedshiftRole"),
+        mode="overwrite",
+        preserve_index=True,
+    )
+    df2 = session.pandas.read_sql_redshift(sql="SELECT * FROM public.test WHERE id = 6",
+                                           iam_role=redshift_parameters.get("RedshiftRole"),
+                                           connection="aws-data-wrangler-redshift",
+                                           temp_s3_path=path2)
+    assert df2.equals(pd.DataFrame())