Fix athena.read_sql_query for empty table and chunk size not returning an empty frame generator (#1685)

Author: LeonLuttenberger

awswrangler/athena/_read.py

@@ -103,12 +103,18 @@ def _fetch_parquet_result(
     if not paths:
         if not temp_table_fqn:
             raise exceptions.EmptyDataFrame("Query would return untyped, empty dataframe.")
+
         database, temp_table_name = map(lambda x: x.replace('"', ""), temp_table_fqn.split("."))
         dtype_dict = catalog.get_table_types(database=database, table=temp_table_name, boto3_session=boto3_session)
         df = pd.DataFrame(columns=list(dtype_dict.keys()))
         df = cast_pandas_with_athena_types(df=df, dtype=dtype_dict)
         df = _apply_query_metadata(df=df, query_metadata=query_metadata)
+
+        if chunked:
+            return (df,)
+
         return df
+
     ret = s3.read_parquet(
         path=paths,
         use_threads=use_threads,

tests/test_athena.py

@@ -595,6 +595,19 @@ def test_read_sql_query_wo_results(path, glue_database, glue_table):
     ensure_athena_query_metadata(df=df, ctas_approach=False, encrypted=False)
 
 
+@pytest.mark.parametrize("ctas_approach", [False, True])
+def test_read_sql_query_wo_results_chunked(path, glue_database, glue_table, ctas_approach):
+    wr.catalog.create_parquet_table(database=glue_database, table=glue_table, path=path, columns_types={"c0": "int"})
+    sql = f"SELECT * FROM {glue_database}.{glue_table}"
+
+    counter = 0
+    for df in wr.athena.read_sql_query(sql, database=glue_database, ctas_approach=ctas_approach, chunksize=100):
+        assert df.empty
+        counter += 1
+
+    assert counter == 1
+
+
 @pytest.mark.xfail()
 def test_read_sql_query_wo_results_ctas(path, glue_database, glue_table):
     wr.catalog.create_parquet_table(database=glue_database, table=glue_table, path=path, columns_types={"c0": "int"})