laskoviymishka commented on code in PR #1993:
URL: https://github.com/apache/iceberg-go/pull/1993#discussion_r3972043797
##########
table/arrow_scanner.go:
##########
@@ -2139,6 +2144,17 @@ func (as *arrowScan) GetRecords(ctx context.Context,
tasks []FileScanTask) (*arr
}
tableProperties := as.metadata.Properties()
+ batchSize := as.options.Get(ParquetBatchSizeKey, "")
+ if as.arrowBatchSize > 0 {
+ batchSize = strconv.Itoa(as.arrowBatchSize)
+ }
+ if batchSize != "" {
+ tableProperties = maps.Clone(tableProperties)
Review Comment:
`as.metadata.Properties()` already hands back a fresh clone, or
`iceberg.Properties{}` when the map is nil, so `tableProperties` is never nil
here and is already safe to mutate. That makes this `maps.Clone` a second
allocation on every batch-size scan, and the `if tableProperties == nil` branch
below it dead (it also nudges the next reader into thinking `Properties()` can
return nil). I'd drop both and set the key in place. Minor, non-blocking.
##########
table/write_records.go:
##########
@@ -95,6 +97,32 @@ func WithClusteredWrite() WriteRecordOption {
}
}
+// WithRecordBatchBufferSize sets the capacity, in record batches, of
+// each rolling data writer's input channel. Every buffered batch is
+// retained in memory until its writer consumes it, so this bound times
+// the batch row count caps the memory a stalled writer can hold. The
+// default is 64 batches. A non-positive value is ignored.
+func WithRecordBatchBufferSize(n int) WriteRecordOption {
+ return func(c *writeRecordConfig) {
+ if n > 0 {
+ c.recordBatchBufferSize = n
+ }
+ }
+}
+
+// WithParquetRowGroupLimit overrides the table's
+// write.parquet.row-group-limit property for this write, capping the
+// number of rows per Parquet row group in the output files. Smaller row
+// groups bound the writer's buffered memory before each flush. A
+// non-positive value is ignored.
+func WithParquetRowGroupLimit(n int) WriteRecordOption {
Review Comment:
Same interop note as `WithArrowBatchSize`, the other direction:
`write.parquet.row-group-limit` is honored by iceberg-go and PyIceberg
(PyIceberg defines the identical key), but Java has no row-count row-group cap.
Its control is byte-based (`write.parquet.row-group-size-bytes`). Worth a
one-line doc note so nobody expects a Java reader to respect a table that leans
on this. Not blocking.
##########
table/table.go:
##########
@@ -1304,6 +1304,24 @@ func WithRowLineage() ScanOption {
}
}
+// WithArrowBatchSize caps the number of rows decoded per Arrow record
+// batch when reading data files, overriding the table's
+// read.parquet.batch-size property for this scan. Smaller batches bound
+// the memory a scan holds per decoded batch, which matters when the
+// consumer buffers batches (e.g. a compaction's read+write pipeline).
+// The cap is stored on the scan itself rather than in the options map,
+// so it applies regardless of ordering relative to [WithOptions]. A
+// non-positive value is ignored.
+func WithArrowBatchSize(n int) ScanOption {
Review Comment:
Doc nit, not blocking: `read.parquet.batch-size` is an iceberg-go-internal
key rather than a spec property. Java uses
`read.parquet.vectorization.batch-size` (default 5000, versus 131072 here) and
PyIceberg has no read-batch-size key at all. The cap is per-scan and never
persisted so nothing breaks, but a one-line note that this key doesn't match
the Java convention would save someone reaching for `SetProperties` expecting
cross-client behavior.
##########
table/write_read_tuning_test.go:
##########
@@ -0,0 +1,247 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package table
+
+import (
+ "fmt"
+ "path/filepath"
+ "strings"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet/file"
+ "github.com/apache/iceberg-go"
+ iceio "github.com/apache/iceberg-go/io"
+ "github.com/google/uuid"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func manyRowsJSON(n int) string {
+ var sb strings.Builder
+ sb.WriteString("[")
+ for i := range n {
+ if i > 0 {
+ sb.WriteString(",")
+ }
+ fmt.Fprintf(&sb, `{"id":%d,"data":"row-%d"}`, i, i)
+ }
+ sb.WriteString("]")
+
+ return sb.String()
+}
+
+func TestWithArrowBatchSizeCapsDecodedBatches(t *testing.T) {
+ const numRows = 100
+ const batchSize = 7
+ tbl := buildV3TableWithRows(t, manyRowsJSON(numRows))
+
+ _, records, err :=
tbl.Scan(WithArrowBatchSize(batchSize)).ToArrowRecords(t.Context())
+ require.NoError(t, err)
+
+ var totalRows, batches int64
+ for rec, err := range records {
+ require.NoError(t, err)
+ assert.LessOrEqual(t, rec.NumRows(), int64(batchSize))
+ totalRows += rec.NumRows()
+ batches++
+ rec.Release()
+ }
+ assert.Equal(t, int64(numRows), totalRows)
+ assert.Greater(t, batches, int64(numRows/batchSize))
Review Comment:
Small robustness thing, and it predates this round: `numRows/batchSize` is
integer division, so this really asserts `batches > floor(100/7) = 14`. It
passes today because 100 isn't divisible by 7, but if the constants ever drift
to an evenly-divisible pair (say 70 rows at 7 per batch, exactly 10 batches)
then `10 > 10` is false and the test fails for the wrong reason. I'd use the
ceil, `assert.GreaterOrEqual(t, batches,
int64((numRows+batchSize-1)/batchSize))`, or just assert the exact count.
##########
table/write_records.go:
##########
@@ -95,6 +97,32 @@ func WithClusteredWrite() WriteRecordOption {
}
}
+// WithRecordBatchBufferSize sets the capacity, in record batches, of
+// each rolling data writer's input channel. Every buffered batch is
+// retained in memory until its writer consumes it, so this bound times
+// the batch row count caps the memory a stalled writer can hold. The
+// default is 64 batches. A non-positive value is ignored.
+func WithRecordBatchBufferSize(n int) WriteRecordOption {
Review Comment:
The stated memory bound holds for unpartitioned and clustered writes, but
the fanout writer opens one of these channels per active partition, so peak on
a partitioned table is really `N_active_partitions * this bound *
rows-per-batch`. Might be worth saying the bound is per-partition-writer, so
nobody sizes it assuming a single channel.
--
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]
