zeroshade commented on code in PR #1351:
URL: https://github.com/apache/iceberg-go/pull/1351#discussion_r3521453817
##########
table/rolling_data_writer.go:
##########
@@ -402,37 +441,155 @@ func (r *RollingDataWriter) stream(outputDataFilesCh
chan<- iceberg.DataFile) {
sorted, err := compute.SortRecordBatch(r.ctx,
converted, r.factory.sortKeys)
converted.Release()
if err != nil {
- record.Release()
- r.sendError(err)
-
- return
+ return err
}
converted = sorted
}
- err = currentWriter.Write(converted)
+ err := currentWriter.Write(converted)
converted.Release()
+ if err != nil {
+ return err
+ }
+
+ if allowRoll && currentWriter.BytesWritten() >=
r.factory.targetFileSize {
+ return closeWriter()
+ }
+
+ return nil
+ }
+
+ // flushBootstrap infers the schema, opens the file, and replays the
buffer.
+ flushBootstrap := func() error {
+ if len(buf) == 0 {
+ bootstrapping = false
+
+ return nil
+ }
+ fileArrowSchema =
tblutils.ShreddedArrowSchema(r.factory.arrowSchema,
+ inferShreddingFromBatches(buf,
r.factory.shredBufferRows))
+ if err := openCurrent(); err != nil {
+ releaseBuf()
+
+ return err
+ }
+ // Detach so the deferred releaseBuf can't double-release the
replayed batches.
+ batches := buf
+ buf = nil
+ bufRows = 0
+ bootstrapping = false
+ for i, b := range batches {
+ if err := writeConverted(b, false); err != nil {
+ for _, rest := range batches[i+1:] {
+ rest.Release()
+ }
+
+ return err
+ }
+ }
+ if currentWriter != nil && currentWriter.BytesWritten() >=
r.factory.targetFileSize {
+ return closeWriter()
+ }
+
+ return nil
+ }
+
+ for record := range r.recordCh {
+ converted, err := ToRequestedSchema(r.ctx, r.factory.fileSchema,
+ r.factory.taskSchema, record, SchemaOptions{
+ DowncastTimestamp: true,
+ IncludeFieldIDs: true,
+ UseWriteDefault: true,
+ })
record.Release()
if err != nil {
r.sendError(err)
return
}
- if currentWriter.BytesWritten() >= r.factory.targetFileSize {
- if err := closeWriter(); err != nil {
+ if bootstrapping {
+ buf = append(buf, converted)
Review Comment:
[MAJOR] `shredBufferRows` is documented as rows buffered per file, but this
appends the entire converted batch before checking the row limit. A single
large upstream batch (for example 1M rows with a small configured limit) will
still be fully retained until bootstrap flushes, causing the memory spike the
limit is meant to prevent and violating the property contract. Please either
slice batches so only the first `shredBufferRows` rows are retained for
bootstrap and write the remainder after the file opens, or document this as a
batch-granular limit with an additional byte/batch bound. A regression test
with one large batch and a small `write.parquet.variant-inference-buffer-size`
would catch this.
##########
table/rolling_data_writer.go:
##########
@@ -402,37 +441,155 @@ func (r *RollingDataWriter) stream(outputDataFilesCh
chan<- iceberg.DataFile) {
sorted, err := compute.SortRecordBatch(r.ctx,
converted, r.factory.sortKeys)
converted.Release()
if err != nil {
- record.Release()
- r.sendError(err)
-
- return
+ return err
}
converted = sorted
}
- err = currentWriter.Write(converted)
+ err := currentWriter.Write(converted)
converted.Release()
+ if err != nil {
+ return err
+ }
+
+ if allowRoll && currentWriter.BytesWritten() >=
r.factory.targetFileSize {
+ return closeWriter()
+ }
+
+ return nil
+ }
+
+ // flushBootstrap infers the schema, opens the file, and replays the
buffer.
+ flushBootstrap := func() error {
+ if len(buf) == 0 {
+ bootstrapping = false
+
+ return nil
+ }
+ fileArrowSchema =
tblutils.ShreddedArrowSchema(r.factory.arrowSchema,
+ inferShreddingFromBatches(buf,
r.factory.shredBufferRows))
+ if err := openCurrent(); err != nil {
+ releaseBuf()
+
+ return err
+ }
+ // Detach so the deferred releaseBuf can't double-release the
replayed batches.
+ batches := buf
+ buf = nil
+ bufRows = 0
+ bootstrapping = false
+ for i, b := range batches {
+ if err := writeConverted(b, false); err != nil {
+ for _, rest := range batches[i+1:] {
+ rest.Release()
+ }
+
+ return err
+ }
+ }
+ if currentWriter != nil && currentWriter.BytesWritten() >=
r.factory.targetFileSize {
+ return closeWriter()
+ }
+
+ return nil
+ }
+
+ for record := range r.recordCh {
+ converted, err := ToRequestedSchema(r.ctx, r.factory.fileSchema,
+ r.factory.taskSchema, record, SchemaOptions{
+ DowncastTimestamp: true,
+ IncludeFieldIDs: true,
+ UseWriteDefault: true,
+ })
record.Release()
if err != nil {
r.sendError(err)
return
}
- if currentWriter.BytesWritten() >= r.factory.targetFileSize {
- if err := closeWriter(); err != nil {
+ if bootstrapping {
+ buf = append(buf, converted)
+ bufRows += converted.NumRows()
+ if bufRows >= int64(r.factory.shredBufferRows) {
+ if err := flushBootstrap(); err != nil {
+ r.sendError(err)
+
+ return
+ }
+ }
+
+ continue
+ }
+
+ if currentWriter == nil {
+ if err := openCurrent(); err != nil {
+ converted.Release()
r.sendError(err)
return
}
}
+ if err := writeConverted(converted, true); err != nil {
+ r.sendError(err)
+
+ return
+ }
}
+ // Channel closed: flush any partial bootstrap buffer, then finalize.
+ if bootstrapping {
+ if err := flushBootstrap(); err != nil {
+ r.sendError(err)
+
+ return
+ }
+ }
if err := closeWriter(); err != nil {
r.sendError(err)
}
}
+// inferShreddingFromBatches infers the inner type per top-level variant
column,
+// sampling up to limit non-null values from the buffer. Keyed by column index.
+func inferShreddingFromBatches(buf []arrow.RecordBatch, limit int)
map[int]arrow.DataType {
+ if len(buf) == 0 {
+ return nil
+ }
+
+ inferred := make(map[int]arrow.DataType)
Review Comment:
[NIT] Since this map is keyed by column index, future changes beyond
top-level use may be easier to follow if the key carries the field ID or name
instead.
##########
table/rolling_data_writer.go:
##########
@@ -402,37 +441,155 @@ func (r *RollingDataWriter) stream(outputDataFilesCh
chan<- iceberg.DataFile) {
sorted, err := compute.SortRecordBatch(r.ctx,
converted, r.factory.sortKeys)
converted.Release()
if err != nil {
- record.Release()
- r.sendError(err)
-
- return
+ return err
}
converted = sorted
}
- err = currentWriter.Write(converted)
+ err := currentWriter.Write(converted)
converted.Release()
+ if err != nil {
+ return err
+ }
+
+ if allowRoll && currentWriter.BytesWritten() >=
r.factory.targetFileSize {
+ return closeWriter()
+ }
+
+ return nil
+ }
+
+ // flushBootstrap infers the schema, opens the file, and replays the
buffer.
+ flushBootstrap := func() error {
+ if len(buf) == 0 {
+ bootstrapping = false
+
+ return nil
+ }
+ fileArrowSchema =
tblutils.ShreddedArrowSchema(r.factory.arrowSchema,
+ inferShreddingFromBatches(buf,
r.factory.shredBufferRows))
+ if err := openCurrent(); err != nil {
+ releaseBuf()
+
+ return err
+ }
+ // Detach so the deferred releaseBuf can't double-release the
replayed batches.
+ batches := buf
+ buf = nil
+ bufRows = 0
+ bootstrapping = false
+ for i, b := range batches {
+ if err := writeConverted(b, false); err != nil {
+ for _, rest := range batches[i+1:] {
+ rest.Release()
+ }
+
+ return err
+ }
+ }
+ if currentWriter != nil && currentWriter.BytesWritten() >=
r.factory.targetFileSize {
+ return closeWriter()
+ }
+
+ return nil
+ }
+
+ for record := range r.recordCh {
+ converted, err := ToRequestedSchema(r.ctx, r.factory.fileSchema,
+ r.factory.taskSchema, record, SchemaOptions{
+ DowncastTimestamp: true,
+ IncludeFieldIDs: true,
+ UseWriteDefault: true,
+ })
record.Release()
if err != nil {
r.sendError(err)
return
}
- if currentWriter.BytesWritten() >= r.factory.targetFileSize {
- if err := closeWriter(); err != nil {
+ if bootstrapping {
+ buf = append(buf, converted)
+ bufRows += converted.NumRows()
+ if bufRows >= int64(r.factory.shredBufferRows) {
+ if err := flushBootstrap(); err != nil {
+ r.sendError(err)
+
+ return
+ }
+ }
+
+ continue
+ }
+
+ if currentWriter == nil {
+ if err := openCurrent(); err != nil {
+ converted.Release()
r.sendError(err)
return
}
}
+ if err := writeConverted(converted, true); err != nil {
+ r.sendError(err)
+
+ return
+ }
}
+ // Channel closed: flush any partial bootstrap buffer, then finalize.
+ if bootstrapping {
+ if err := flushBootstrap(); err != nil {
+ r.sendError(err)
+
+ return
+ }
+ }
if err := closeWriter(); err != nil {
r.sendError(err)
}
}
+// inferShreddingFromBatches infers the inner type per top-level variant
column,
+// sampling up to limit non-null values from the buffer. Keyed by column index.
+func inferShreddingFromBatches(buf []arrow.RecordBatch, limit int)
map[int]arrow.DataType {
+ if len(buf) == 0 {
+ return nil
+ }
+
+ inferred := make(map[int]arrow.DataType)
+ for col := 0; col < int(buf[0].NumCols()); col++ {
+ if _, ok := buf[0].Column(col).(*extensions.VariantArray); !ok {
+ continue
+ }
+ var sample []variant.Value
+ for _, b := range buf {
+ if len(sample) >= limit {
+ break
+ }
+ va, ok := b.Column(col).(*extensions.VariantArray)
+ if !ok {
+ continue
+ }
+ for i := 0; i < va.Len() && len(sample) < limit; i++ {
+ if va.Storage().IsNull(i) {
+ continue
+ }
+ v, err := va.Value(i)
+ if err != nil {
Review Comment:
[QUESTION] Should `va.Value(i)` errors be surfaced instead of silently
skipped? Ignoring malformed values can infer shredding from a biased subset of
rows. Returning an error, or at least logging/tracking skipped rows, would make
this easier to diagnose.
##########
table/internal/variant_shredding.go:
##########
@@ -0,0 +1,421 @@
+// 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 internal
+
+import (
+ "sort"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/decimal"
+ "github.com/apache/arrow-go/v18/arrow/extensions"
+ "github.com/apache/arrow-go/v18/parquet/variant"
+)
+
+// Variant shredding inference: most-common-type selection with explicit
+// tie-break, integer/decimal widening, and frequency/depth/field-count bounds.
+const (
+ minFieldFrequency = 0.10
+ maxShreddedFields = 300
+ maxShreddingDepth = 50
+ maxIntermediateFields = 1000
+)
+
+// integerPriority and decimalPriority pick the widest observed family member.
+var integerPriority = map[variant.Type]int{
+ variant.Int8: 0, variant.Int16: 1, variant.Int32: 2, variant.Int64: 3,
+}
+
+var decimalPriority = map[variant.Type]int{
+ variant.Decimal4: 0, variant.Decimal8: 1, variant.Decimal16: 2,
+}
+
+// tieBreakPriority breaks count ties in most-common-type selection; higher
wins.
+// Types absent here (Null, Array, Object) resolve to -1.
+var tieBreakPriority = map[variant.Type]int{
+ variant.Bool: 0, variant.Int8: 1, variant.Int16: 2, variant.Int32: 3,
+ variant.Int64: 4, variant.Float: 5, variant.Double: 6,
variant.Decimal4: 7,
+ variant.Decimal8: 8, variant.Decimal16: 9, variant.Date: 10,
variant.Time: 11,
+ variant.TimestampMicros: 12, variant.TimestampMicrosNTZ: 13,
variant.Binary: 14,
+ variant.String: 15, variant.TimestampNanos: 16,
variant.TimestampNanosNTZ: 17,
+ variant.UUID: 18,
+}
+
+type fieldInfo struct {
+ typeCounts map[variant.Type]int
+ observationCount int
+ maxDecimalScale int
+ maxDecimalIntDigits int
+}
+
+func newFieldInfo() *fieldInfo {
+ return &fieldInfo{typeCounts: make(map[variant.Type]int)}
+}
+
+// observe records one value's type at this node (per-value counting).
+func (f *fieldInfo) observe(v variant.Value) {
+ f.observationCount++
+
+ t := v.Type()
+ // arrow-go has a single Bool type, so no false-folds-to-true step.
+ f.typeCounts[t]++
+
+ if isDecimalType(t) {
+ intDigits, scale := decimalDigits(v.Value())
+ if scale > f.maxDecimalScale {
+ f.maxDecimalScale = scale
+ }
+ if intDigits > f.maxDecimalIntDigits {
+ f.maxDecimalIntDigits = intDigits
+ }
+ }
+}
+
+// mostCommonType collapses int/decimal families to the widest, then picks max
by count.
+func (f *fieldInfo) mostCommonType() (variant.Type, bool) {
+ combined := make(map[variant.Type]int)
+
+ intTotal, decTotal := 0, 0
+ var widestInt, widestDec variant.Type
+ haveInt, haveDec := false, false
+
+ for t, c := range f.typeCounts {
+ switch {
+ case isIntegerType(t):
+ intTotal += c
+ if !haveInt || integerPriority[t] >
integerPriority[widestInt] {
+ widestInt, haveInt = t, true
+ }
+ case isDecimalType(t):
+ decTotal += c
+ if !haveDec || decimalPriority[t] >
decimalPriority[widestDec] {
+ widestDec, haveDec = t, true
+ }
+ default:
+ combined[t] = c
+ }
+ }
+ if haveInt {
+ combined[widestInt] = intTotal
+ }
+ if haveDec {
+ combined[widestDec] = decTotal
+ }
+
+ if len(combined) == 0 {
+ return 0, false
+ }
+
+ // Max by count, then tieBreakPriority, then type value (stable, sorted
keys).
+ var best variant.Type
+ bestCount, bestPrio := -1, -2
+ first := true
+ keys := make([]variant.Type, 0, len(combined))
+ for t := range combined {
+ keys = append(keys, t)
+ }
+ sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] })
+ for _, t := range keys {
+ c := combined[t]
+ p := tiePriority(t)
+ if first || c > bestCount || (c == bestCount && p > bestPrio) {
+ best, bestCount, bestPrio, first = t, c, p, false
+ }
+ }
+
+ return best, true
+}
+
+func tiePriority(t variant.Type) int {
+ if p, ok := tieBreakPriority[t]; ok {
+ return p
+ }
+
+ return -1
+}
+
+func isIntegerType(t variant.Type) bool {
+ switch t {
+ case variant.Int8, variant.Int16, variant.Int32, variant.Int64:
+ return true
+ }
+
+ return false
+}
+
+func isDecimalType(t variant.Type) bool {
+ switch t {
+ case variant.Decimal4, variant.Decimal8, variant.Decimal16:
+ return true
+ }
+
+ return false
+}
+
+// decimalDigits returns the integer-digit count and scale of a variant
decimal.
+func decimalDigits(val any) (intDigits, scale int) {
+ switch d := val.(type) {
+ case variant.DecimalValue[decimal.Decimal32]:
+ return coefficientDigits(d.Value.ToString(0)) - int(d.Scale),
int(d.Scale)
+ case variant.DecimalValue[decimal.Decimal64]:
+ return coefficientDigits(d.Value.ToString(0)) - int(d.Scale),
int(d.Scale)
+ case variant.DecimalValue[decimal.Decimal128]:
+ return coefficientDigits(d.Value.ToString(0)) - int(d.Scale),
int(d.Scale)
+ }
+
+ return 0, 0
+}
+
+func coefficientDigits(s string) int {
+ if len(s) > 0 && s[0] == '-' {
+ s = s[1:]
+ }
+
+ return len(s)
+}
+
+type pathNode struct {
+ info *fieldInfo
+ objectChildren map[string]*pathNode
+ arrayElement *pathNode
+}
+
+func newPathNode() *pathNode {
+ return &pathNode{info: newFieldInfo(), objectChildren:
make(map[string]*pathNode)}
+}
+
+// AnalyzeVariantShredding infers the inner Arrow type to shred the sample by,
or
+// ok=false to not shred. The result is the INNER type for
NewShreddedVariantType.
+func AnalyzeVariantShredding(sample []variant.Value) (arrow.DataType, bool) {
+ if len(sample) == 0 {
+ return nil, false
+ }
+
+ root := newPathNode()
+ for _, v := range sample {
+ traverseVariant(root, v, 0)
+ }
+
+ rootType, ok := root.info.mostCommonType()
+ if !ok {
+ return nil, false
+ }
+
+ pruneInfrequent(root, root.info.observationCount)
+
+ dt := buildInnerType(root, rootType)
+ if dt == nil {
+ return nil, false
+ }
+
+ return dt, true
+}
+
+func traverseVariant(node *pathNode, v variant.Value, depth int) {
+ t := v.Type()
+ if t == variant.Null {
+ return
+ }
+
+ node.info.observe(v)
+
+ switch {
+ case t == variant.Object && depth < maxShreddingDepth:
+ obj, ok := v.Value().(variant.ObjectValue)
+ if !ok {
+ return
+ }
+ for name, fv := range obj.Values() {
+ child := node.objectChildren[name]
+ if child == nil {
+ if len(node.objectChildren) >=
maxIntermediateFields {
+ continue
+ }
+ child = newPathNode()
+ node.objectChildren[name] = child
+ }
+ traverseVariant(child, fv, depth+1)
+ }
+ case t == variant.Array && depth < maxShreddingDepth:
+ arr, ok := v.Value().(variant.ArrayValue)
+ if !ok {
+ return
+ }
+ if node.arrayElement == nil {
+ node.arrayElement = newPathNode()
+ }
+ for ev := range arr.Values() {
+ traverseVariant(node.arrayElement, ev, depth+1)
+ }
+ }
+}
+
+func pruneInfrequent(node *pathNode, totalRows int) {
+ if len(node.objectChildren) == 0 && node.arrayElement == nil {
+ return
+ }
+
+ // Frequency floor (strict <); exactly minFieldFrequency is kept.
+ for name, child := range node.objectChildren {
+ if float64(child.info.observationCount)/float64(totalRows) <
minFieldFrequency {
+ delete(node.objectChildren, name)
+ }
+ }
+
+ // Cap to maxShreddedFields, keeping the highest count; on a count tie
keep the
+ // alphabetically-smaller name.
+ if len(node.objectChildren) > maxShreddedFields {
+ names := make([]string, 0, len(node.objectChildren))
+ for name := range node.objectChildren {
+ names = append(names, name)
+ }
+ sort.Slice(names, func(i, j int) bool {
+ ci :=
node.objectChildren[names[i]].info.observationCount
+ cj :=
node.objectChildren[names[j]].info.observationCount
+ if ci != cj {
+ return ci > cj
+ }
+
+ return names[i] < names[j]
+ })
+ for _, name := range names[maxShreddedFields:] {
+ delete(node.objectChildren, name)
+ }
+ }
+
+ for _, child := range node.objectChildren {
+ pruneInfrequent(child, totalRows)
Review Comment:
[MINOR] The `maxIntermediateFields` and `maxShreddedFields` caps are applied
per object node (see the child creation/pruning above), not as a global
inferred-schema budget. Deep or broad samples can still infer a very large
recursive schema. Please add a global field/node budget, or rename/document
these as per-object caps and add a nested high-cardinality test.
##########
table/internal/parquet_files.go:
##########
@@ -305,6 +311,12 @@ func (parquetFormat) GetWriteProperties(props
iceberg.Properties) any {
writerProps = append(writerProps,
parquet.WithBloomFilterEnabledFor(colName, enabled))
}
+ // Shredded decimals need INT32/INT64/FLBA-by-precision
(VariantShredding.md);
+ // arrow-go emits those only with StoreDecimalAsInteger. Gated to keep
default writes unchanged.
+ if props.GetBool(ParquetShredVariantsKey, ParquetShredVariantsDefault) {
Review Comment:
[MAJOR] This enables `StoreDecimalAsInteger` solely from the table property,
not from the per-file inference result. As a result,
`write.parquet.shred-variants=true` changes the physical encoding of ordinary
decimal columns in unrelated files even when no shredded decimal `typed_value`
is present. Please select decimal-as-integer per file after inference
determines a shredded decimal is needed, or explicitly document the table-level
encoding side effect and add property-on/no-variant-column coverage.
--
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]
