Copilot commented on code in PR #8781:
URL: https://github.com/apache/paimon/pull/8781#discussion_r3636261035
##########
paimon-python/pypaimon/daft/daft_datasink.py:
##########
@@ -280,3 +322,445 @@ def finalize(self, write_results:
list[WriteResult[list[Any]]]) -> MicroPartitio
"file_name": pa.array([f.file_name for f in all_files],
type=pa.string()),
}
)
+
+ def _primary_key_write_error(
+ self, non_empty_workers: int
+ ) -> str | None:
+ if not self._table.is_primary_key_table:
+ return None
+
+ from pypaimon.table.bucket_mode import BucketMode
+
+ bucket_mode = self._table.bucket_mode()
+ if bucket_mode == BucketMode.POSTPONE_MODE:
+ return None
+ if bucket_mode in (
+ BucketMode.HASH_FIXED,
+ BucketMode.HASH_DYNAMIC,
+ ) and non_empty_workers <= 1:
+ return None
+ if bucket_mode in (BucketMode.HASH_FIXED, BucketMode.HASH_DYNAMIC):
+ reason = "require a single non-empty Daft write task"
+ elif bucket_mode == BucketMode.CROSS_PARTITION:
+ reason = (
+ "require a persistent global primary-key index, which "
+ "PyPaimon does not yet support"
+ )
+ else:
+ reason = f"are unsafe for {bucket_mode.name}"
+ return (
+ f"Direct PaimonDataSink primary-key writes {reason}. Use "
+ "pypaimon.daft.write_paimon or PaimonTable.append/overwrite "
+ "for coordinated input."
+ )
+
+
+class PaimonCommitDataSink(PaimonDataSink):
+ """Commit files already produced by partition/bucket group UDFs."""
+
+ def __init__(
+ self,
+ table: FileStoreTable,
+ mode: str = "append",
+ commit_messages_column: str = COMMIT_MESSAGES_COLUMN,
+ ) -> None:
+ self._commit_messages_column = commit_messages_column
+ super().__init__(table, mode)
+
+ def __getstate__(self) -> dict[str, Any]:
+ state = super().__getstate__()
+ state["_commit_messages_column"] = self._commit_messages_column
+ return state
+
+ def __setstate__(self, state: dict[str, Any]) -> None:
+ self._commit_messages_column = state.pop(
+ "_commit_messages_column", COMMIT_MESSAGES_COLUMN
+ )
+ super().__setstate__(state)
+
+ @property
+ def commit_user(self) -> str | None:
+ return self._commit_user
+
+ def name(self) -> str:
+ return "Paimon Commit"
+
+ def _primary_key_write_error(
+ self, non_empty_workers: int
+ ) -> str | None:
+ return None
+
+ def write(
+ self, micropartitions: Iterator[MicroPartition]
+ ) -> Iterator[WriteResult[list[Any]]]:
+ commit_messages = []
+ total_bytes = 0
+ for mp in micropartitions:
+ for rb in mp.get_record_batches():
+ batch = rb.to_arrow_record_batch()
+ column_index = batch.schema.get_field_index(
+ self._commit_messages_column
+ )
+ if column_index < 0:
+ raise ValueError(
+ "Missing internal column "
+ f"{self._commit_messages_column!r}"
+ )
+ for payload in batch.column(column_index).to_pylist():
+ if payload is None:
+ continue
+ total_bytes += len(payload)
+ decoded = pickle.loads(payload)
+ if not isinstance(decoded, list):
+ raise TypeError(
+ "Invalid Paimon group-write result: expected a
list "
+ f"of commit messages, got {type(decoded).__name__}"
+ )
+ commit_messages.extend(decoded)
Review Comment:
`pickle.loads(payload)` will execute arbitrary code if the internal
commit-messages column is ever user-controlled (e.g., a DataFrame manually
constructed with the same column name and written via `PaimonCommitDataSink`).
Even if this is intended to be internal-only, add a lightweight framing/magic
prefix and basic type checks so the sink refuses to unpickle unexpected
payloads.
This issue also appears on line 762 of the same file.
--
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]
