fix: Reading of chunked parquet when columns parameter is specified (#2599)

* fix: Reading of chunked parquet when columns parameter is specified

* fix: Preserve schema metadata

Signed-off-by: Anton Kukushkin <[email protected]>

* add test case

Signed-off-by: Anton Kukushkin <[email protected]>

---------

Signed-off-by: Anton Kukushkin <[email protected]>
Co-authored-by: kukushking <[email protected]>

Author: rchromik

awswrangler/s3/_read_parquet.py

@@ -237,8 +237,13 @@ def _read_parquet_chunked(
             chunks = pq_file.iter_batches(
                 batch_size=batch_size, columns=columns, use_threads=use_threads_flag, use_pandas_metadata=False
             )
+
+            schema = pq_file.schema.to_arrow_schema()
+            if columns:
+                schema = pa.schema([schema.field(column) for column in columns], schema.metadata)
+
             table = _add_table_partitions(
-                table=pa.Table.from_batches(chunks, schema=pq_file.schema.to_arrow_schema()),
+                table=pa.Table.from_batches(chunks, schema=schema),
                 path=path,
                 path_root=path_root,
             )

tests/unit/test_s3_parquet.py

@@ -816,3 +816,18 @@ def test_read_parquet_schema_validation_with_index_column(path) -> None:
         validate_schema=True,
     )
     assert df0.shape == df1.shape
+
+
+@pytest.mark.parametrize("columns", [None, ["c0"], ["c0", "c1"]])
+@pytest.mark.parametrize("chunked", [False, True, 2])
+def test_chunked_columns(path, columns, chunked):
+    df = pd.DataFrame({"c0": [0, 1, 2, 3, 4], "c1": [2, 3, 4, 5, 6], "c3": [3, 4, 5, 6, 7]})
+    path_file = f"{path}0.parquet"
+
+    wr.s3.to_parquet(df, path_file)
+    df2 = wr.s3.read_parquet(path_file, columns=columns, chunked=chunked)
+
+    if chunked:
+        df2 = pd.concat(list(df2), ignore_index=True)
+
+    assert df[columns].shape if columns else df.shape == df2.shape