JingsongLi commented on code in PR #9047:
URL: https://github.com/apache/paimon/pull/9047#discussion_r3721683459
##########
paimon-python/pypaimon/write/table_write.py:
##########
@@ -61,25 +60,30 @@ def write_arrow(self, table: pa.Table):
def write_arrow_batch(self, data: pa.RecordBatch):
self._validate_pyarrow_schema(data.schema)
- partitions, buckets =
self.row_key_extractor.extract_partition_bucket_batch(data)
- partition_bucket_groups = defaultdict(list)
- for i in range(data.num_rows):
- partition_bucket_groups[(tuple(partitions[i]),
buckets[i])].append(i)
-
- for (partition, bucket), row_indices in
partition_bucket_groups.items():
- if len(row_indices) == data.num_rows:
+ for partition, bucket, row_indices in \
+ self.row_key_extractor.extract_partition_bucket_groups(data):
+ if row_indices is None:
# Every input row belongs to the same partition/bucket.
Passing the
# original batch through avoids copying large BLOB values
through
# Arrow take before the dedicated BLOB writer consumes them.
sub_table = data
- elif row_indices[-1] - row_indices[0] + 1 == len(row_indices):
- # Contiguous groups can share the original Arrow buffers
instead of
- # gathering their rows into newly allocated buffers with take.
- sub_table = data.slice(row_indices[0], len(row_indices))
else:
- indices_array = pa.array(row_indices, type=pa.int64())
- sub_table = pa.compute.take(data, indices_array)
+ # row_indices is an int64 array of this group's rows. Arrow's
+ # grouped list aggregation runs multi-threaded and does NOT
+ # guarantee ascending order within a group, so the span must be
+ # derived from min/max, not the first/last positions.
+ bounds = pa.compute.min_max(row_indices)
+ lo = bounds["min"].as_py()
+ hi = bounds["max"].as_py()
+ count = len(row_indices)
+ if hi - lo + 1 == count:
+ # Distinct row indices spanning exactly `count` values are
+ # contiguous, so share the original Arrow buffers instead
of
+ # gathering their rows into newly allocated buffers with
take.
+ sub_table = data.slice(lo, count)
+ else:
+ sub_table = pa.compute.take(data, row_indices)
Review Comment:
Could we preserve the original row order before this `take`? `hash_list`
with the default threaded group-by can return a group's indices out of order.
For contiguous groups the `slice` above restores input order, but this path
writes rows in aggregation order. `KeyValueDataWriter._add_system_fields` then
assigns sequence numbers in that order, so an interleaved group containing
repeated primary keys can make an earlier input row receive the highest
sequence and win deduplication or partial update. I reproduced this with
PyArrow 19.0.1: a group whose last input index was `2,999,988` ended with
`1,048,575`. Please sort the row-index values before gathering (or otherwise
make grouping stable) and assert row order/latest-wins in the regression test.
##########
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:
`Table.group_by` alone is not a sufficient capability check. PyArrow 7 has
it, but `hash_list` was added in Arrow 8, while the Python 3.7 dependency range
still permits `pyarrow>=7,<13`. On PyArrow 7 this aggregate lookup raises
`ArrowKeyError`, which is not caught above, so every `write_arrow_batch` fails
instead of using the fallback. Could we probe for `hash_list` when initializing
the capability flag (or explicitly handle only the missing-kernel case) and add
a PyArrow 7 regression test?
--
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]
