nssalian commented on code in PR #1351:
URL: https://github.com/apache/iceberg-go/pull/1351#discussion_r3555277085
##########
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:
Thanks for catching this. I bounded the bootstrap buffer to
`shredBufferRows`: an oversized batch is sliced, only the head is buffered, the
tail writes steady-state. Wrote a test
`TestShreddedVariantWriteLargeBootstrapBatch` (3000-row batch, 100-row buffer,
1-byte target) that asserts the batch splits across files - whole-batch
buffering would produce one - plus row conservation under the checked allocator.
##########
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:
Fixed this too. `StoreDecimalAsInteger` is now decided per file in
`NewFileWriter` from the inferred schema, not the table property, so unrelated
decimal columns keep FLBA. The props slice is `slices.Clone`d before the append
so concurrent partition writers don't race
(`TestShreddedVariantPartitionedDecimalRace` fails under `-race` without it).
Covered: INT32/INT64/FLBA by precision, and wrote a test
`TestShreddedVariantPartitionedGoWriteSparkRead` that reads a Go-written
partitioned shredded decimal back via Spark.
##########
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:
Per-node on purpose, matching `VariantShreddingAnalyzer` in Java Iceberg
(per-node cap in `pruneInfrequentFields`, frequency measured against root row
count at every depth, same 0.10 / 300 / 50 / 1000 constants). A global budget
would diverge from Java. Fixed the doc comment and added
TestAnalyzeNestedHighCardinality + TestAnalyzeDepthCap.
##########
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:
Now this counts and logs undecodable sample values instead of dropping them
silently; kept non-fatal since a malformed value is already hard-errored at
shred time. `TestInferShreddingSkipsUndecodable` covers it in testing.
##########
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:
Keyed by column index on purpose: `ShreddedArrowSchema` applies each type to
`fields[idx]` by position. Clarified the doc; happy to switch to field-id if
you prefer.
--
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]
