yugan95 commented on code in PR #9047:
URL: https://github.com/apache/paimon/pull/9047#discussion_r3725499236
##########
paimon-python/pypaimon/write/row_key_extractor.py:
##########
@@ -95,6 +104,91 @@ def extract_partitions_batch(self, data: pa.RecordBatch) ->
List[Tuple]:
"""Return partition tuples without calculating bucket hashes."""
return self._extract_partitions_batch(data)
+ def extract_partition_bucket_groups(
+ self, data: pa.RecordBatch) -> List[Tuple[Tuple, int,
Optional[pa.Array]]]:
+ """Group row indices by (partition, bucket) for the write path.
+
+ Returns a list of ``(partition, bucket, row_indices)`` where
+ ``row_indices`` is an Arrow ``int64`` array of the rows belonging to
the
+ group, or ``None`` when the whole batch is a single group (so callers
can
+ pass the original batch through without copying large values, e.g.
BLOBs).
+
+ The grouping is done in Arrow so only the distinct group keys are
+ materialized into Python objects, instead of one ``.as_py()`` scalar
per
+ row. The old per-row loop held the GIL for the entire batch, which
+ serialized multi-threaded writers down to ~1 core.
+
+ Buckets are computed once here, in row order, via
``_extract_buckets_batch``
+ so stateful extractors (dynamic bucket) keep their exact assignment
+ sequence and side effects regardless of which grouping path runs.
+ """
+ buckets = self._extract_buckets_batch(data)
+ if _ARROW_GROUP_BY_SUPPORTED:
+ try:
+ return self._group_indices_arrow(data, buckets)
+ except (pa.ArrowNotImplementedError, pa.ArrowInvalid):
+ # Only Arrow's own "can't group this column type" errors fall
+ # back to the legacy per-row grouping; any other exception is a
+ # real bug and must propagate rather than silently degrade to
the
+ # GIL-bound path. `buckets` is reused (never recomputed) so
+ # stateful extractors are not double-notified. Log so the
+ # (GIL-bound) fallback is visible.
+ logger.warning(
+ "Arrow group_by could not handle the partition/bucket key "
+ "types; falling back to per-row grouping (GIL-bound).",
+ exc_info=True)
+ # pyarrow < 7.0.0 has no group_by; use the per-row grouping directly.
+ return self._group_indices_python(data, buckets)
+
+ def _group_indices_arrow(
+ self, data: pa.RecordBatch,
+ buckets: List[int]) -> List[Tuple[Tuple, int, Optional[pa.Array]]]:
+ num_rows = data.num_rows
+ columns = {}
+ key_names = []
+ for k, pi in enumerate(self.partition_indices):
+ name = f"__p{k}"
+ columns[name] = data.column(pi)
+ key_names.append(name)
+ columns["__bucket"] = pa.array(buckets, type=pa.int32())
+ key_names.append("__bucket")
+ # Build the row index with numpy (C speed, releases the GIL). Using a
+ # Python range() here makes pyarrow iterate it element by element under
+ # the GIL, which dominates this method and kills multi-thread scaling.
+ columns["__idx"] = pa.array(np.arange(num_rows, dtype=np.int64))
+
+ grouped = pa.table(columns).group_by(key_names).aggregate([("__idx",
"list")])
Review Comment:
You're right, hasattr(pa.Table, "group_by") isn't sufficient — fixed in
d851cc34.
The flag now comes from _probe_arrow_group_by(), which actually runs a tiny
group_by(...).aggregate([("__idx", "list")]) once at import. That way pyarrow 7
(has group_by, but hash_list only landed in Arrow 8 → ArrowKeyError) is
detected as unsupported and falls through to per-row grouping, same as pyarrow
< 7 which has no group_by at all. I went with probing the kernel rather than
catching ArrowKeyError per-batch so there's no repeated warning/exception on
the hot path.
Regression test test_probe_arrow_group_by_false_when_hash_list_missing
simulates the missing kernel by making the aggregate raise ArrowKeyError and
asserts the probe returns False.
--
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]
