read_parquet_table fails when chunked is true for partitioned table

[bentkibler]

If I try to read a partitioned table in batch chunks using s3.read_parquet_table(chunked=True) on a table that has partitions defined, I get the following exception:

<ipython-input-3-6b4a443538f2> in test_table(dbname, tablename, use_threads, chunked, partition_filter, columns)
      3     tot_rows = 0
      4     tot_files = 0
----> 5     dfs = wr.s3.read_parquet_table(boto3_session=session, 
      6                                    database=dbname,
      7                                    table=tablename,

~/.conda/envs/awswrangler/lib/python3.9/site-packages/awswrangler/_config.py in wrapper(*args_raw, **kwargs)
    415                 del args[name]
    416                 args = {**args, **keywords}
--> 417         return function(**args)
    418 
    419     wrapper.__doc__ = _inject_config_doc(doc=function.__doc__, available_configs=available_configs)

~/.conda/envs/awswrangler/lib/python3.9/site-packages/awswrangler/s3/_read_parquet.py in read_parquet_table(table, database, filename_suffix, filename_ignore_suffix, catalog_id, partition_filter, columns, validate_schema, categories, safe, map_types, chunked, use_threads, boto3_session, s3_additional_kwargs)
    774     except KeyError as ex:
    775         raise exceptions.InvalidTable(f"Missing s3 location for {database}.{table}.") from ex
--> 776     return _data_types.cast_pandas_with_athena_types(
    777         df=read_parquet(
    778             path=path,

~/.conda/envs/awswrangler/lib/python3.9/site-packages/awswrangler/_data_types.py in cast_pandas_with_athena_types(df, dtype)
    597     for col, athena_type in dtype.items():
    598         if (
--> 599             (col in df.columns)
    600             and (athena_type.startswith("array") is False)
    601             and (athena_type.startswith("struct") is False)

AttributeError: 'generator' object has no attribute 'columns'

The exception is thrown in line 599 of cast_pandas_with_athena_types because it's expecting an actual dataframe, not a generator of dataframes. It seems the calling methodread_parquet_table() is where a fix is needed, such that it calls the cast_ method on each dataframe returned by the generator from read_parquet().

My environment is Python 3.9, awswrangler 2.6.0, on AWS Linux2

[maxispeicher]

I've created a fix for reading a parquet table in chunked mode. If you find the time, you can try if it works for you 🙂

pip uninstall awswrangler -y
pip install git+https://github.com/maxispeicher/aws-data-wrangler.git@read_pq_table_chunked

[bentkibler]

I tried that fix in my environment, and it works as expected now. Thanks for the quick turnaround!

[jaidisido]

Covered in 2.7.0 release