zeroshade commented on code in PR #1766:
URL: https://github.com/apache/iceberg-go/pull/1766#discussion_r3833787217
##########
table/partitioned_fanout_writer.go:
##########
@@ -541,6 +559,35 @@ func partitionBatchByKey(ctx context.Context)
partitionBatchFn {
}
}
+// contiguousSliceHasBoundedRetention reports whether a partial zero-copy slice
+// retains no more than twice as many input rows as it returns. The range must
+// already have been validated by contiguousRowRange.
+func contiguousSliceHasBoundedRetention(start, end, numRows int64) bool {
+ selectedRows := end - start
+
+ return selectedRows >= numRows-selectedRows
Review Comment:
This bounds retained rows, not retained bytes. For a 4,096-row string batch
where the omitted first half contains 16 KiB strings and the selected
contiguous second half contains one-byte strings, this condition accepts
`NewSlice`. After releasing the input, the result retained 67,125,824 bytes;
`compute.Take` for the same rows retained 10,560 bytes. A 64-record queue can
therefore still pin several gigabytes. Please use a byte-aware policy, or keep
partial zero-copy slices only for storage types where row count actually bounds
backing-buffer size.
##########
table/partitioned_fanout_writer_test.go:
##########
@@ -683,6 +683,144 @@ func (s *FanoutWriterTestSuite)
TestGetRecordPartitionsWithDroppedLeadingSourceC
s.Equal("foo=null/bar=7/baz=true",
spec.PartitionToPath(partitions[0].partitionRec, icebergSchema))
}
+func (s *FanoutWriterTestSuite) TestPartitionBatchByKeyFastPaths() {
+ arrowSchema := arrow.NewSchema([]arrow.Field{{Name: "value", Type:
arrow.PrimitiveTypes.Int64}}, nil)
+ record := s.createCustomTestRecord(arrowSchema, [][]any{{int64(0)},
{int64(1)}, {int64(2)}, {int64(3)}, {int64(4)}})
+ defer record.Release()
+
+ partitionBatch := partitionBatchByKey(s.ctx)
+
+ full, err := partitionBatch(record, []int64{0, 1, 2, 3, 4})
+ s.Require().NoError(err)
+ s.Same(record, full)
+ full.Release()
+
+ contiguous, err := partitionBatch(record, []int64{1, 2, 3})
+ s.Require().NoError(err)
+ defer contiguous.Release()
+ s.NotSame(record, contiguous)
+ s.Equal(int64(3), contiguous.NumRows())
+ contiguousValues := contiguous.Column(0).(*array.Int64)
+ s.Equal([]int64{1, 2, 3}, []int64{
+ contiguousValues.Value(0), contiguousValues.Value(1),
contiguousValues.Value(2),
+ })
+
+ scattered, err := partitionBatch(record, []int64{0, 2, 4})
+ s.Require().NoError(err)
+ defer scattered.Release()
+ s.Equal(int64(3), scattered.NumRows())
+ scatteredValues := scattered.Column(0).(*array.Int64)
+ s.Equal([]int64{0, 2, 4}, []int64{
+ scatteredValues.Value(0), scatteredValues.Value(1),
scatteredValues.Value(2),
+ })
+
+ empty, err := partitionBatch(record, nil)
+ s.Require().NoError(err)
+ defer empty.Release()
+ s.NotSame(record, empty)
+ s.Zero(empty.NumRows())
+
+ emptyRecord := s.createCustomTestRecord(arrowSchema, nil)
+ defer emptyRecord.Release()
+ emptyFull, err := partitionBatch(emptyRecord, nil)
+ s.Require().NoError(err)
+ s.Same(emptyRecord, emptyFull)
+ emptyFull.Release()
+
+ _, err = partitionBatch(record, []int64{4, 5})
+ s.Error(err)
+}
+
+func (s *FanoutWriterTestSuite)
TestPartitionBatchByKeyBoundsQueuedWriterMemory() {
+ const inputRows = 4096
+ const payloadSize = 128
+
+ arrSchema := arrow.NewSchema([]arrow.Field{
+ {Name: "id", Type: arrow.PrimitiveTypes.Int64},
+ {Name: "payload", Type: arrow.BinaryTypes.String},
+ }, nil)
+
+ probe := s.createLargeTestRecord(arrSchema, inputRows, 0, payloadSize)
+ fullBatchBytes := s.mem.CurrentAlloc()
+ probe.Release()
+
+ writer := &RollingDataWriter{
+ recordCh: make(chan arrow.RecordBatch,
rollingDataWriterQueueCapacity),
+ errorCh: make(chan error, 1),
+ ctx: s.ctx,
+ }
+ defer func() {
+ for len(writer.recordCh) > 0 {
+ (<-writer.recordCh).Release()
+ }
+ }()
+
+ partitionBatch := partitionBatchByKey(s.ctx)
+ peakBytes := 0
+ for batch := range rollingDataWriterQueueCapacity {
+ record := s.createLargeTestRecord(arrSchema, inputRows,
int64(batch*inputRows), payloadSize)
+ partitioned, err := partitionBatch(record, []int64{inputRows /
2})
+ s.Require().NoError(err)
+
+ s.Require().NoError(writer.Add(partitioned))
+ partitioned.Release()
+ record.Release()
+
+ peakBytes = max(peakBytes, s.mem.CurrentAlloc())
+ }
+
+ s.Less(peakBytes, fullBatchBytes*4, "queued partial batches should not
retain complete input batches")
Review Comment:
This regression uses uniform 128-byte payloads, so selected-row ratio
happens to approximate retained-byte ratio. Please add a variable-width skew
case where omitted rows contain large payloads and selected rows are small; the
current implementation retains about 67 MB instead of roughly 10 KB for that
shape.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
