laskoviymishka commented on code in PR #1595:
URL: https://github.com/apache/iceberg-go/pull/1595#discussion_r3675823298
##########
table/arrow_utils.go:
##########
@@ -1880,6 +1881,15 @@ func unpartitionedWrite(ctx context.Context, factory
*writerFactory, records ite
return
}
+ select {
+ case <-writerCtx.Done():
+ errCh <- context.Cause(writerCtx)
Review Comment:
I think there's a subtle wrinkle here. This select fires after we've already
pulled `rec` off the iterator but before `writer.Add(rec)`, so on cancellation
that batch is dropped between the two lines — harmless today since the
`releasing` wrapper frees it, but worth noting.
The bigger thing is what we send. `writerCtx` is a plain `WithCancel`, so
`context.Cause(writerCtx)` is always `context.Canceled`, and on the early-stop
path the iterator returns before it ever reads `errCh` — so that value is
silently discarded. Anyone who later wires up a read of `errCh` after early
stop would see a spurious `context.Canceled` reported as the outcome of what
was actually a clean stop.
`writer.Add`'s own blocking `<-ctx.Done()` check is already a cancellation
point, so one option is to drop this non-blocking select entirely and lean on
that. If we want to keep the eager check, I'd switch to `WithCancelCause`, call
`cancel(nil)` on clean stop, and only send to `errCh` when `context.Cause` is
non-nil. wdyt?
##########
table/partitioned_fanout_writer.go:
##########
@@ -175,32 +180,37 @@ func (p *partitionedFanoutWriter) processRecord(ctx
context.Context, writerCtx c
return nil
}
-func (p *partitionedFanoutWriter) yieldDataFiles(fanoutWorkers
*errgroup.Group, outputDataFilesCh chan iceberg.DataFile, writerCancel
context.CancelFunc) iter.Seq2[iceberg.DataFile, error] {
+func (p *partitionedFanoutWriter) yieldDataFiles(fanoutWorkers
*errgroup.Group, inputRecordsCh chan arrow.RecordBatch, outputDataFilesCh chan
iceberg.DataFile, cancel context.CancelFunc) iter.Seq2[iceberg.DataFile, error]
{
return yieldDataFiles(
p.writerFactory,
fanoutWorkers,
+ inputRecordsCh,
outputDataFilesCh,
p.writerFactory.closeAll,
p.writerFactory.abortAll,
- writerCancel,
+ cancel,
)
}
func yieldDataFiles(
writerFactory *writerFactory,
fanoutWorkers *errgroup.Group,
+ inputRecordsCh chan arrow.RecordBatch,
outputDataFilesCh chan iceberg.DataFile,
closeAll func() error,
abortAll func(),
- writerCancel context.CancelFunc,
+ cancel context.CancelFunc,
) iter.Seq2[iceberg.DataFile, error] {
// Use a channel to safely communicate the error from the goroutine
// to avoid a data race between writing err in the goroutine and
reading it in the iterator.
errCh := make(chan error, 1)
go func() {
defer close(outputDataFilesCh)
- defer writerCancel()
+ defer cancel()
err := fanoutWorkers.Wait()
+ for record := range inputRecordsCh {
Review Comment:
Could we add a one-liner on why this drain is safe? It took me a second:
`startRecordFeeder` is part of `fanoutWorkers`, so by the time `Wait()` returns
the feeder has already exited and closed `inputRecordsCh`. This only ever
releases batches the feeder `Retain()`ed but no worker dequeued — dequeued ones
are freed by the worker's own `defer Release()`, so there's no double-release.
Mostly the comment guards against someone later pulling the feeder out of the
errgroup and quietly deadlocking this loop.
##########
table/arrow_utils.go:
##########
@@ -1880,6 +1881,15 @@ func unpartitionedWrite(ctx context.Context, factory
*writerFactory, records ite
return
}
+ select {
Review Comment:
Small one, assuming the select stays: there's no blank line between the
closing `}` of the `if err != nil` block and this `select`, which wsl will flag
— I'm fairly sure CI fails on it. Worth adding one.
##########
table/partitioned_fanout_writer.go:
##########
@@ -175,32 +180,37 @@ func (p *partitionedFanoutWriter) processRecord(ctx
context.Context, writerCtx c
return nil
}
-func (p *partitionedFanoutWriter) yieldDataFiles(fanoutWorkers
*errgroup.Group, outputDataFilesCh chan iceberg.DataFile, writerCancel
context.CancelFunc) iter.Seq2[iceberg.DataFile, error] {
+func (p *partitionedFanoutWriter) yieldDataFiles(fanoutWorkers
*errgroup.Group, inputRecordsCh chan arrow.RecordBatch, outputDataFilesCh chan
iceberg.DataFile, cancel context.CancelFunc) iter.Seq2[iceberg.DataFile, error]
{
return yieldDataFiles(
p.writerFactory,
fanoutWorkers,
+ inputRecordsCh,
outputDataFilesCh,
p.writerFactory.closeAll,
p.writerFactory.abortAll,
- writerCancel,
+ cancel,
)
}
func yieldDataFiles(
writerFactory *writerFactory,
fanoutWorkers *errgroup.Group,
+ inputRecordsCh chan arrow.RecordBatch,
Review Comment:
Since `yieldDataFiles` only ever drains this, could we type it `<-chan
arrow.RecordBatch` here (and in the two wrapper methods + the pos-delete one)?
Makes the read-only intent explicit and stops anyone accidentally sending into
a channel `startRecordFeeder` owns.
##########
table/write_records_test.go:
##########
@@ -233,6 +234,63 @@ func (s *WriteRecordsTestSuite)
TestSmallTargetFileSizeProducesMultipleFiles() {
s.Equal(int64(1000), totalRows)
}
+func (s *WriteRecordsTestSuite) TestEarlyStopCancelsRecordProduction() {
+ tests := []struct {
+ name string
+ partitioned bool
+ }{
+ {name: "unpartitioned"},
+ {name: "partitioned", partitioned: true},
+ }
+
+ for _, tt := range tests {
+ s.Run(tt.name, func() {
+ loc := filepath.ToSlash(s.T().TempDir())
+ iceSch := iceberg.NewSchema(1,
+ iceberg.NestedField{ID: 1, Name: "id", Type:
iceberg.PrimitiveTypes.Int32},
+ iceberg.NestedField{ID: 2, Name: "name", Type:
iceberg.PrimitiveTypes.String},
+ )
+ spec := iceberg.NewPartitionSpec()
+ if tt.partitioned {
+ spec =
iceberg.NewPartitionSpec(iceberg.PartitionField{
+ SourceIDs: []int{1}, FieldID: 1000,
Name: "id", Transform: iceberg.IdentityTransform{},
+ })
+ }
+ meta, err := table.NewMetadata(iceSch, &spec,
table.UnsortedSortOrder, loc, iceberg.Properties{})
+ s.Require().NoError(err)
+ tbl := table.New(
+ table.Identifier{"test", tt.name}, meta,
filepath.Join(loc, "metadata", "v1.metadata.json"),
+ func(context.Context) (iceio.IO, error) {
return iceio.LocalFS{}, nil }, nil,
+ )
+
+ schema := s.arrowSchema()
+ var produced atomic.Int32
+ records := func(yield func(arrow.RecordBatch, error)
bool) {
+ for range 1000 {
+ produced.Add(1)
+ record := s.buildRecords(schema, 100)
+ accepted := yield(record, nil)
+ if !accepted {
+ return
+ }
+ }
+ }
+
+ for df, writeErr := range table.WriteRecords(
+ s.ctx, tbl, schema, records,
table.WithTargetFileSize(1), table.WithMaxWriteWorkers(1),
+ ) {
+ s.Require().NoError(writeErr)
+ s.Require().NotNil(df)
+
+ break
+ }
+
+ s.Positive(produced.Load())
+ s.Less(produced.Load(), int32(1000))
Review Comment:
`< 1000` is a pretty loose bound — with `WithMaxWriteWorkers(1)` and
`WithTargetFileSize(1)` we break after the first file, so real production is
only 1-3 batches. As written this still passes even if cancellation were
delayed by 990 records, which is most of what we're trying to catch. I'd
tighten it to something like `< 10`. Same assertion at
`pos_delete_partitioned_fanout_writer_test.go:463`.
##########
table/partitioned_fanout_writer.go:
##########
@@ -215,6 +225,7 @@ func yieldDataFiles(
for range outputDataFilesCh {
}
}()
+ defer cancel()
Review Comment:
`cancel()` is deferred here and again in the goroutine at line 209 — both
idempotent so it's fine, but the intent isn't obvious (this one fires first via
LIFO and is the early-stop signal; the goroutine's is the fallback).
`clustered_writer.go:183-186` already has a comment for exactly this — mind
mirroring it?
##########
table/pos_delete_partitioned_fanout_writer_test.go:
##########
@@ -402,6 +403,66 @@ func
TestPositionDeletePartitionedFanoutWriterRoutesPartitionsIndependently(t *t
assert.Equal(t, int64(1), byPart[2].Count(), "id=2 delete file must
contain only the one row targeting pathB")
}
+func
TestPositionDeletePartitionedFanoutWriterEarlyStopCancelsRecordProduction(t
*testing.T) {
Review Comment:
Every other top-level test in this file opens with `t.Parallel()` — this
one's missing it. Worth adding for consistency.
##########
table/pos_delete_partitioned_fanout_writer_test.go:
##########
@@ -402,6 +403,66 @@ func
TestPositionDeletePartitionedFanoutWriterRoutesPartitionsIndependently(t *t
assert.Equal(t, int64(1), byPart[2].Count(), "id=2 delete file must
contain only the one row targeting pathB")
}
+func
TestPositionDeletePartitionedFanoutWriterEarlyStopCancelsRecordProduction(t
*testing.T) {
+ const path = "file://t/id=1/a.parquet"
+
+ tableSchema := iceberg.NewSchema(
+ 0,
+ iceberg.NestedField{ID: 1, Name: "id", Type:
iceberg.PrimitiveTypes.Int32, Required: true},
+ )
+ partitionSpec := iceberg.NewPartitionSpec(iceberg.PartitionField{
+ FieldID: 1000, SourceIDs: []int{1}, Name: "id", Transform:
iceberg.IdentityTransform{},
+ })
+ metadataBuilder, err := NewMetadataBuilder(2)
+ require.NoError(t, err)
+ require.NoError(t, metadataBuilder.AddSchema(tableSchema))
+ require.NoError(t, metadataBuilder.SetCurrentSchemaID(0))
+ require.NoError(t, metadataBuilder.AddPartitionSpec(&partitionSpec,
true))
+ require.NoError(t, metadataBuilder.SetDefaultSpecID(0))
+ require.NoError(t, metadataBuilder.AddSortOrder(&UnsortedSortOrder))
+ require.NoError(t, metadataBuilder.SetDefaultSortOrderID(0))
+ latestMeta, err := metadataBuilder.Build()
+ require.NoError(t, err)
+
+ var produced atomic.Int32
+ records := func(yield func(arrow.RecordBatch, error) bool) {
+ for range 1000 {
+ produced.Add(1)
+ batch := mustLoadRecordBatchFromJSON(
+ PositionalDeleteArrowSchema,
+ fmt.Sprintf(`[{"file_path":%q,"pos":0}]`, path),
+ )
+ if !yield(batch, nil) {
+ return
+ }
+ }
+ }
+
+ writeUUID := uuid.New()
+ factory, err := newWriterFactory(t.TempDir(), recordWritingArgs{
+ fs: &io.LocalFS{}, sc: PositionalDeleteArrowSchema, writeUUID:
&writeUUID, counter: internal.Counter(0),
+ }, metadataBuilder, iceberg.PositionalDeleteSchema, 1,
+ withContentType(iceberg.EntryContentPosDeletes),
+ withFactoryFileSchema(iceberg.PositionalDeleteSchema))
+ require.NoError(t, err)
+ writer := newPositionDeletePartitionedFanoutWriter(
+ latestMeta,
+ map[string]partitionContext{path: {partitionData:
map[int]any{1000: int32(1)}, specID: 0}},
+ records,
+ factory,
+ )
+
+ for dataFile, writeErr := range writer.Write(t.Context(), 1) {
+ require.NoError(t, writeErr)
+ require.NotNil(t, dataFile)
+
+ break
+ }
+
+ assert.Positive(t, produced.Load())
+ assert.Less(t, produced.Load(), int32(1000))
Review Comment:
The PR description calls out that retained Arrow memory gets released,
verified by a checked allocator — but this direct-writer test doesn't assert
that (only the table-level path does). Could we thread a
`memory.CheckedAllocator` through the context and assert `CurrentAllocated() ==
0` after the loop? Cheap way to pin the memory claim on this path too.
--
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]
