yugan95 commented on code in PR #9047:
URL: https://github.com/apache/paimon/pull/9047#discussion_r3725496849
##########
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:
Good catch, thanks — confirmed and fixed in d851cc34.
Rather than sorting at the take call site, I made grouping stable at the
source: _group_indices_arrow now sorts each group's row indices back to
ascending input order (np.sort on the aggregated list, so it stays in C and
releases the GIL). Both grouping paths therefore return indices in input order
— the per-row fallback already appended them ascending.
With that invariant, write_arrow_batch no longer needs min_max; it derives
the contiguous span from the first/last positions directly.
Regression test test_group_indices_arrow_sorts_unordered_aggregation forces
the aggregation to report a group out of order ([4,0,2] / [3,1]) and asserts
the output comes back [0,2,4] / [1,3], and the existing grouping tests now
assert order-sensitively instead of by membership.
--
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]
