You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@iceberg.apache.org by "HonahX (via GitHub)" <gi...@apache.org> on 2024/03/06 07:55:13 UTC
Re: [PR] Bin-pack Writes Operation into multiple parquet files, and parallelize writing `WriteTask`s [iceberg-python]
HonahX commented on code in PR #444:
URL: https://github.com/apache/iceberg-python/pull/444#discussion_r1513978231
##########
pyiceberg/io/pyarrow.py:
##########
@@ -1717,54 +1717,65 @@ def fill_parquet_file_metadata(
def write_file(table: Table, tasks: Iterator[WriteTask], file_schema: Optional[Schema] = None) -> Iterator[DataFile]:
- task = next(tasks)
-
- try:
- _ = next(tasks)
- # If there are more tasks, raise an exception
- raise NotImplementedError("Only unpartitioned writes are supported: https://github.com/apache/iceberg-python/issues/208")
- except StopIteration:
- pass
-
- parquet_writer_kwargs = _get_parquet_writer_kwargs(table.properties)
-
- file_path = f'{table.location()}/data/{task.generate_data_file_filename("parquet")}'
file_schema = file_schema or table.schema()
arrow_file_schema = schema_to_pyarrow(file_schema)
-
- fo = table.io.new_output(file_path)
+ parquet_writer_kwargs = _get_parquet_writer_kwargs(table.properties)
row_group_size = PropertyUtil.property_as_int(
properties=table.properties,
property_name=TableProperties.PARQUET_ROW_GROUP_SIZE_BYTES,
default=TableProperties.PARQUET_ROW_GROUP_SIZE_BYTES_DEFAULT,
)
- with fo.create(overwrite=True) as fos:
- with pq.ParquetWriter(fos, schema=arrow_file_schema, **parquet_writer_kwargs) as writer:
- writer.write_table(task.df, row_group_size=row_group_size)
-
- data_file = DataFile(
- content=DataFileContent.DATA,
- file_path=file_path,
- file_format=FileFormat.PARQUET,
- partition=Record(),
- file_size_in_bytes=len(fo),
- # After this has been fixed:
- # https://github.com/apache/iceberg-python/issues/271
- # sort_order_id=task.sort_order_id,
- sort_order_id=None,
- # Just copy these from the table for now
- spec_id=table.spec().spec_id,
- equality_ids=None,
- key_metadata=None,
- )
- fill_parquet_file_metadata(
- data_file=data_file,
- parquet_metadata=writer.writer.metadata,
- stats_columns=compute_statistics_plan(file_schema, table.properties),
- parquet_column_mapping=parquet_path_to_id_mapping(file_schema),
+ def write_parquet(task: WriteTask) -> DataFile:
+ file_path = f'{table.location()}/data/{task.generate_data_file_filename("parquet")}'
+ fo = table.io.new_output(file_path)
+ with fo.create(overwrite=True) as fos:
+ with pq.ParquetWriter(fos, schema=arrow_file_schema, **parquet_writer_kwargs) as writer:
+ writer.write(pa.Table.from_batches(task.record_batches), row_group_size=row_group_size)
+
+ data_file = DataFile(
+ content=DataFileContent.DATA,
+ file_path=file_path,
+ file_format=FileFormat.PARQUET,
+ partition=Record(),
+ file_size_in_bytes=len(fo),
+ # After this has been fixed:
+ # https://github.com/apache/iceberg-python/issues/271
+ # sort_order_id=task.sort_order_id,
+ sort_order_id=None,
+ # Just copy these from the table for now
+ spec_id=table.spec().spec_id,
+ equality_ids=None,
+ key_metadata=None,
+ )
+ fill_parquet_file_metadata(
+ data_file=data_file,
+ parquet_metadata=writer.writer.metadata,
+ stats_columns=compute_statistics_plan(file_schema, table.properties),
+ parquet_column_mapping=parquet_path_to_id_mapping(file_schema),
+ )
+ return data_file
+
+ executor = ExecutorFactory.get_or_create()
+ data_files = executor.map(write_parquet, tasks)
+
+ return iter(data_files)
+
+
+def bin_pack_arrow_table(tbl: pa.Table, target_file_size: int) -> Iterator[List[pa.RecordBatch]]:
+ from pyiceberg.utils.bin_packing import PackingIterator
+
+ avg_row_size_bytes = tbl.nbytes / tbl.num_rows
+ target_rows_per_file = target_file_size // avg_row_size_bytes
+ batches = tbl.to_batches(max_chunksize=target_rows_per_file)
+ bin_packed_record_batches = PackingIterator(
+ items=batches,
+ target_weight=target_file_size,
+ lookback=2,
Review Comment:
Shall we add a comment to explain the choice of `lookback=2`?
##########
pyiceberg/table/__init__.py:
##########
@@ -2359,17 +2362,25 @@ def _dataframe_to_data_files(
Returns:
An iterable that supplies datafiles that represent the table.
"""
- from pyiceberg.io.pyarrow import write_file
+ from pyiceberg.io.pyarrow import bin_pack_arrow_table, write_file
if len(table.spec().fields) > 0:
raise ValueError("Cannot write to partitioned tables")
counter = itertools.count(0)
write_uuid = write_uuid or uuid.uuid4()
+ target_file_size = PropertyUtil.property_as_int(
+ properties=table.properties,
+ property_name=TableProperties.WRITE_TARGET_FILE_SIZE_BYTES,
+ default=TableProperties.WRITE_TARGET_FILE_SIZE_BYTES_DEFAULT,
+ )
+ assert isinstance(target_file_size, int)
Review Comment:
We try to avoid `assert` outside the test. If the primary purpose for this check it to pass the linter, I think could just use `# type: ignore` comment. WDYT?
##########
tests/integration/test_writes.py:
##########
@@ -424,6 +424,30 @@ def test_data_files(spark: SparkSession, session_catalog: Catalog, arrow_table_w
assert [row.deleted_data_files_count for row in rows] == [0, 0, 1, 0, 0]
+@pytest.mark.integration
+def test_write_multiple_data_files(spark: SparkSession, session_catalog: Catalog, arrow_table_with_null: pa.Table) -> None:
Review Comment:
The test name implies it will write multiple data files, but the test body only writes 1 data file per overwrite. Shall we add a third overwrite which writes multiple data files to the table?
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org