corleyma commented on code in PR #786:
URL: https://github.com/apache/iceberg-python/pull/786#discussion_r1637425676
##########
pyiceberg/io/pyarrow.py:
##########
@@ -1795,15 +1873,19 @@ def write_file(io: FileIO, table_metadata:
TableMetadata, tasks: Iterator[WriteT
def write_parquet(task: WriteTask) -> DataFile:
table_schema = task.schema
- arrow_table = pa.Table.from_batches(task.record_batches)
+
# if schema needs to be transformed, use the transformed schema and
adjust the arrow table accordingly
# otherwise use the original schema
if (sanitized_schema := sanitize_column_names(table_schema)) !=
table_schema:
file_schema = sanitized_schema
- arrow_table = to_requested_schema(requested_schema=file_schema,
file_schema=table_schema, table=arrow_table)
+ batches = [
+ to_requested_schema(requested_schema=file_schema,
file_schema=table_schema, batch=batch)
+ for batch in task.record_batches
+ ]
else:
file_schema = table_schema
-
+ batches = task.record_batches
+ arrow_table = pa.Table.from_batches(batches)
Review Comment:
i.e., maybe we could do something like the following?:
```python
# Check if schema needs to be transformed
sanitized_schema = sanitize_column_names(table_schema)
needs_sanitization = sanitized_schema != table_schema
file_path =
f'{table_metadata.location}/data/{task.generate_data_file_path("parquet")}'
fo = io.new_output(file_path)
with fo.create(overwrite=True) as fos:
with pq.ParquetWriter(fos, schema=sanitized_schema.as_arrow(),
**parquet_writer_kwargs) as writer:
for batch in task.record_batches:
if needs_sanitization:
sanitized_batch = to_requested_schema(
requested_schema=sanitized_schema,
file_schema=table_schema, batch=batch
)
else:
sanitized_batch = batch
writer.write_table(pa.Table.from_batches([sanitized_batch]),
row_group_size=row_group_size)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
