This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new bb6432f4 feat(arrow/array): roll back JSON builder state after failed
rows (#1113)
bb6432f4 is described below
commit bb6432f423af52c9cb65d5ff6eb9e7cb4bdea997
Author: Minh Vu <[email protected]>
AuthorDate: Fri Aug 14 18:07:05 2026 +0200
feat(arrow/array): roll back JSON builder state after failed rows (#1113)
### Rationale for this change
RecordBuilder mutates field builders while decoding a JSON object. If a
later field fails, values appended for earlier fields remain and can
affect subsequent rows.
### What changes are included in this PR?
Build the nested checkpoint graph once when the RecordBuilder is
created. Each row captures reusable lengths and builder state, then
restores that state on decode errors. This covers lists, structs, maps,
fixed-size lists, unions, dictionaries, variable-width buffers, string
views, and run-end encoded children without rebuilding checkpoint trees
for every successful row.
Parent builders are restored before their children so parent Resize
calls cannot overwrite restored child state. This is especially
important for nested run-end encoded builders.
Custom builders with internal state can participate through the exported
CheckpointState and CheckpointableBuilder interfaces.
### Are these changes tested?
- go test ./arrow/array
- go test ./arrow/extensions
- go test ./arrow/array -run ^$ -bench
BenchmarkRecordFromJSON/Size_1000$ -benchtime=1x -benchmem
### Are there any user-facing changes?
Yes. Failed JSON rows are rolled back completely, including nested
builder state. This also adds two exported interfaces for custom
builders that need to restore internal state during row rollback.
---
arrow/array/binarybuilder.go | 57 ++-
arrow/array/bufferbuilder.go | 46 +++
arrow/array/bufferbuilder_test.go | 33 ++
arrow/array/builder.go | 23 +-
arrow/array/dictionary.go | 24 ++
arrow/array/fixedsize_binarybuilder.go | 5 +
arrow/array/list.go | 13 +
arrow/array/map.go | 2 +
arrow/array/record.go | 173 ++++++++-
arrow/array/record_test.go | 515 +++++++++++++++++++++++++
arrow/array/union.go | 9 +
arrow/extensions/timestamp_with_offset.go | 18 +
arrow/extensions/timestamp_with_offset_test.go | 47 +++
internal/hashing/xxh3_memo_table.go | 31 ++
internal/hashing/xxh3_memo_table_test.go | 226 +++++++++++
internal/hashing/xxh3_memo_table_types.go | 25 ++
16 files changed, 1240 insertions(+), 7 deletions(-)
diff --git a/arrow/array/binarybuilder.go b/arrow/array/binarybuilder.go
index 07114df9..2391bbe2 100644
--- a/arrow/array/binarybuilder.go
+++ b/arrow/array/binarybuilder.go
@@ -227,6 +227,23 @@ func (b *BinaryBuilder) init(capacity int) {
// DataLen returns the number of bytes in the data array.
func (b *BinaryBuilder) DataLen() int { return b.values.length }
+type binaryBuilderCheckpoint struct {
+ builder *BinaryBuilder
+ dataLen int
+}
+
+func (c *binaryBuilderCheckpoint) capture() {
+ c.dataLen = c.builder.DataLen()
+}
+
+func (c *binaryBuilderCheckpoint) restore() {
+ c.builder.ResizeData(c.dataLen)
+}
+
+func (b *BinaryBuilder) newCheckpoint() checkpointState {
+ return &binaryBuilderCheckpoint{builder: b}
+}
+
// DataCap returns the total number of bytes that can be stored
// without allocating additional memory.
func (b *BinaryBuilder) DataCap() int { return b.values.capacity }
@@ -248,13 +265,26 @@ func (b *BinaryBuilder) ReserveData(n int) {
// Resize adjusts the space allocated by b to n elements. If n is greater than
b.Cap(),
// additional memory will be allocated. If n is smaller, the allocated memory
may be reduced.
func (b *BinaryBuilder) Resize(n int) {
+ if n < b.length {
+ b.truncate(n)
+ }
b.offsets.resize((n + 1) * b.offsetByteWidth)
- if (n * b.offsetByteWidth) < b.offsets.Len() {
+ if n < b.offsets.Len() {
b.offsets.SetLength(n * b.offsetByteWidth)
}
b.resize(n, b.init)
}
+func (b *BinaryBuilder) truncate(n int) {
+ dataLen := b.values.Len()
+ if n < b.offsets.Len() {
+ dataLen = b.getOffsetVal(n)
+ }
+ b.builder.truncate(n)
+ b.offsets.SetLength(n * b.offsetByteWidth)
+ b.values.SetLength(dataLen)
+}
+
func (b *BinaryBuilder) ResizeData(n int) {
b.values.length = n
}
@@ -425,6 +455,31 @@ func (b *BinaryViewBuilder) SetBlockSize(sz uint) {
func (b *BinaryViewBuilder) Type() arrow.DataType { return b.dtype }
+type binaryViewBuilderCheckpoint struct {
+ builder *BinaryViewBuilder
+ length int
+ blockState *multiBufferCheckpoint
+}
+
+func (c *binaryViewBuilderCheckpoint) capture() {
+ c.length = c.builder.length
+ c.blockState.capture()
+}
+
+func (c *binaryViewBuilderCheckpoint) restore() {
+ c.blockState.restore()
+ for i := c.length; i < len(c.builder.rawData); i++ {
+ c.builder.rawData[i] = arrow.ViewHeader{}
+ }
+}
+
+func (b *BinaryViewBuilder) newCheckpoint() checkpointState {
+ return &binaryViewBuilderCheckpoint{
+ builder: b,
+ blockState: b.blockBuilder.newCheckpoint(),
+ }
+}
+
func (b *BinaryViewBuilder) Release() {
debug.Assert(b.refCount.Load() > 0, "too many releases")
diff --git a/arrow/array/bufferbuilder.go b/arrow/array/bufferbuilder.go
index a4d5e043..dbf1fbbb 100644
--- a/arrow/array/bufferbuilder.go
+++ b/arrow/array/bufferbuilder.go
@@ -163,6 +163,7 @@ type multiBufferBuilder struct {
mem memory.Allocator
blocks []*memory.Buffer
currentOutBuffer int
+ checkpoint *multiBufferCheckpoint
}
// Retain increases the reference count by 1.
@@ -237,9 +238,54 @@ func (b *multiBufferBuilder) Reset() {
}
}
+type multiBufferCheckpoint struct {
+ builder *multiBufferBuilder
+ blockCount int
+ blockLengths map[int]int
+ currentOutput int
+}
+
+func (b *multiBufferBuilder) newCheckpoint() *multiBufferCheckpoint {
+ checkpoint := &multiBufferCheckpoint{builder: b}
+ b.checkpoint = checkpoint
+ return checkpoint
+}
+
+func (c *multiBufferCheckpoint) capture() {
+ c.blockCount = len(c.builder.blocks)
+ clear(c.blockLengths)
+ c.currentOutput = c.builder.currentOutBuffer
+}
+
+func (c *multiBufferCheckpoint) recordBlock(index, length int) {
+ if index >= c.blockCount {
+ return
+ }
+ if c.blockLengths == nil {
+ c.blockLengths = make(map[int]int)
+ }
+ if _, ok := c.blockLengths[index]; !ok {
+ c.blockLengths[index] = length
+ }
+}
+
+func (c *multiBufferCheckpoint) restore() {
+ for _, block := range c.builder.blocks[c.blockCount:] {
+ block.Release()
+ }
+ c.builder.blocks = c.builder.blocks[:c.blockCount]
+ for i, length := range c.blockLengths {
+ c.builder.blocks[i].Resize(length)
+ }
+ c.builder.currentOutBuffer = c.currentOutput
+}
+
func (b *multiBufferBuilder) UnsafeAppend(hdr *arrow.ViewHeader, val []byte) {
buf := b.blocks[b.currentOutBuffer]
idx, offset := b.currentOutBuffer, buf.Len()
+ if b.checkpoint != nil {
+ b.checkpoint.recordBlock(idx, offset)
+ }
hdr.SetIndexOffset(int32(idx), int32(offset))
n := copy(buf.Buf()[offset:], val)
diff --git a/arrow/array/bufferbuilder_test.go
b/arrow/array/bufferbuilder_test.go
index 7acb5553..bd7b6baf 100644
--- a/arrow/array/bufferbuilder_test.go
+++ b/arrow/array/bufferbuilder_test.go
@@ -39,3 +39,36 @@ func
TestMultiBufferBuilderUnsafeAppendPanicsOnTruncatedCopy(t *testing.T) {
builder.UnsafeAppend(&hdr, make([]byte, 64))
})
}
+
+func TestMultiBufferCheckpointRestoresTouchedBlocks(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ builder := multiBufferBuilder{mem: mem, blockSize: 4}
+ builder.refCount.Add(1)
+ defer builder.Release()
+
+ first := memory.NewResizableBuffer(mem)
+ first.ResizeNoShrink(4)
+ first.Resize(2)
+ second := memory.NewResizableBuffer(mem)
+ second.ResizeNoShrink(4)
+ builder.blocks = []*memory.Buffer{first, second}
+ builder.currentOutBuffer = 1
+
+ var hdr arrow.ViewHeader
+
+ checkpoint := builder.newCheckpoint()
+ checkpoint.capture()
+
+ builder.Reserve(2)
+ builder.UnsafeAppend(&hdr, []byte("gh"))
+ builder.Reserve(1)
+ builder.UnsafeAppend(&hdr, []byte("i"))
+
+ checkpoint.restore()
+ assert.Len(t, builder.blocks, 2)
+ assert.Equal(t, 2, builder.blocks[0].Len())
+ assert.Equal(t, 4, builder.blocks[1].Len())
+ assert.Equal(t, 1, builder.currentOutBuffer)
+}
diff --git a/arrow/array/builder.go b/arrow/array/builder.go
index cdcd35cd..fe0112b7 100644
--- a/arrow/array/builder.go
+++ b/arrow/array/builder.go
@@ -79,6 +79,9 @@ type Builder interface {
// additional memory will be allocated. If n is smaller, the allocated
memory may reduced.
Resize(n int)
+ // truncate removes elements from the end of the builder without
changing its capacity.
+ truncate(n int)
+
// NewArray creates a new array from the memory buffers used
// by the builder and resets the Builder so it can be used to build
// a new array.
@@ -174,9 +177,25 @@ func (b *builder) resize(newBits int, init func(int)) {
memory.Set(b.nullBitmap.Buf()[oldBytesN:], 0)
}
if newBits < b.length {
- b.length = newBits
- b.nulls = newBits - bitutil.CountSetBits(b.nullBitmap.Buf(), 0,
newBits)
+ b.truncate(newBits)
+ }
+}
+
+func (b *builder) truncate(n int) {
+ if n < 0 || n > b.length {
+ panic("arrow/array: invalid builder truncation length")
+ }
+ if n == b.length {
+ return
+ }
+
+ if b.nullBitmap != nil {
+ bitutil.SetBitsTo(b.nullBitmap.Buf(), int64(n),
int64(b.length-n), false)
+ b.nulls = n - bitutil.CountSetBits(b.nullBitmap.Buf(), 0, n)
+ } else if b.nulls > n {
+ b.nulls = n
}
+ b.length = n
}
func (b *builder) reserve(elements int, resize func(int)) {
diff --git a/arrow/array/dictionary.go b/arrow/array/dictionary.go
index 435460e9..d151b209 100644
--- a/arrow/array/dictionary.go
+++ b/arrow/array/dictionary.go
@@ -686,6 +686,13 @@ func (b *dictionaryBuilder) Reserve(n int) {
func (b *dictionaryBuilder) Resize(n int) {
b.idxBuilder.Resize(n)
b.length = b.idxBuilder.Len()
+ b.nulls = b.idxBuilder.NullN()
+}
+
+func (b *dictionaryBuilder) truncate(n int) {
+ b.idxBuilder.truncate(n)
+ b.length = b.idxBuilder.Len()
+ b.nulls = b.idxBuilder.NullN()
}
func (b *dictionaryBuilder) ResetFull() {
@@ -694,6 +701,23 @@ func (b *dictionaryBuilder) ResetFull() {
b.memoTable.Reset()
}
+type dictionaryBuilderCheckpoint struct {
+ builder *dictionaryBuilder
+ size int
+}
+
+func (c *dictionaryBuilderCheckpoint) capture() {
+ c.size = c.builder.memoTable.Size()
+}
+
+func (c *dictionaryBuilderCheckpoint) restore() {
+ c.builder.memoTable.Truncate(c.size)
+}
+
+func (b *dictionaryBuilder) newCheckpoint() checkpointState {
+ return &dictionaryBuilderCheckpoint{builder: b}
+}
+
func (b *dictionaryBuilder) Cap() int { return b.idxBuilder.Cap() }
func (b *dictionaryBuilder) IsNull(i int) bool { return b.idxBuilder.IsNull(i)
}
diff --git a/arrow/array/fixedsize_binarybuilder.go
b/arrow/array/fixedsize_binarybuilder.go
index b6745be6..b11ef52d 100644
--- a/arrow/array/fixedsize_binarybuilder.go
+++ b/arrow/array/fixedsize_binarybuilder.go
@@ -150,6 +150,11 @@ func (b *FixedSizeBinaryBuilder) Resize(n int) {
b.resize(n, b.init)
}
+func (b *FixedSizeBinaryBuilder) truncate(n int) {
+ b.builder.truncate(n)
+ b.values.SetLength(n * b.dtype.ByteWidth)
+}
+
// NewArray creates a FixedSizeBinary array from the memory buffers used by the
// builder and resets the FixedSizeBinaryBuilder so it can be used to build a
new array.
func (b *FixedSizeBinaryBuilder) NewArray() arrow.Array {
diff --git a/arrow/array/list.go b/arrow/array/list.go
index 0379c786..b848192b 100644
--- a/arrow/array/list.go
+++ b/arrow/array/list.go
@@ -511,6 +511,13 @@ func (b *baseListBuilder) Resize(n int) {
b.offsets.Resize(n)
}
+func (b *baseListBuilder) truncate(n int) {
+ b.builder.truncate(n)
+ if b.offsets != nil {
+ b.offsets.truncate(n)
+ }
+}
+
func (b *baseListBuilder) resizeHelper(n int) {
if n < minBuilderCapacity {
n = minBuilderCapacity
@@ -1286,6 +1293,12 @@ func (b *baseListViewBuilder) Resize(n int) {
b.sizes.Resize(n)
}
+func (b *baseListViewBuilder) truncate(n int) {
+ b.builder.truncate(n)
+ b.offsets.truncate(n)
+ b.sizes.truncate(n)
+}
+
func (b *baseListViewBuilder) resizeHelper(n int) {
if n < minBuilderCapacity {
n = minBuilderCapacity
diff --git a/arrow/array/map.go b/arrow/array/map.go
index 1d13ea88..1a7d9e19 100644
--- a/arrow/array/map.go
+++ b/arrow/array/map.go
@@ -259,6 +259,8 @@ func (b *MapBuilder) Reserve(n int) {
b.listBuilder.Reserve(n) }
// b.Cap(), additional memory will be allocated. If n is smaller, the
allocated memory may be reduced.
func (b *MapBuilder) Resize(n int) { b.listBuilder.Resize(n) }
+func (b *MapBuilder) truncate(n int) { b.listBuilder.truncate(n) }
+
// AppendValues is for bulk appending a group of elements with offsets provided
// and validity booleans provided.
func (b *MapBuilder) AppendValues(offsets []int32, valid []bool) {
diff --git a/arrow/array/record.go b/arrow/array/record.go
index b7f84180..443d3492 100644
--- a/arrow/array/record.go
+++ b/arrow/array/record.go
@@ -298,10 +298,11 @@ func (rec *simpleRecord) MarshalJSON() ([]byte, error) {
// RecordBuilder eases the process of building a Record, iteratively, from
// a known Schema.
type RecordBuilder struct {
- refCount atomic.Int64
- mem memory.Allocator
- schema *arrow.Schema
- fields []Builder
+ refCount atomic.Int64
+ mem memory.Allocator
+ schema *arrow.Schema
+ fields []Builder
+ checkpoints []*builderCheckpoint
}
// NewRecordBuilder returns a builder, using the provided memory allocator and
a schema.
@@ -316,6 +317,10 @@ func NewRecordBuilder(mem memory.Allocator, schema
*arrow.Schema) *RecordBuilder
for i := 0; i < schema.NumFields(); i++ {
b.fields[i] = NewBuilder(b.mem, schema.Field(i).Type)
}
+ b.checkpoints = make([]*builderCheckpoint, len(b.fields))
+ for i, field := range b.fields {
+ b.checkpoints[i] = newBuilderCheckpoint(field)
+ }
return b
}
@@ -335,6 +340,7 @@ func (b *RecordBuilder) Release() {
f.Release()
}
b.fields = nil
+ b.checkpoints = nil
}
}
@@ -415,6 +421,139 @@ func (b *RecordBuilder) NewRecord() arrow.Record {
return b.NewRecordBatch()
}
+type checkpointableBuilder interface {
+ newCheckpoint() checkpointState
+}
+
+type checkpointState interface {
+ capture()
+ restore()
+}
+
+// CheckpointState captures and restores builder state that is not represented
by
+// the builder's length or storage builders. RecordBuilder reuses the same
+// checkpoint for each row, calling Capture before decoding and Restore after a
+// failed decode.
+type CheckpointState interface {
+ // Capture records the current state of the builder.
+ Capture()
+ // Restore returns the builder to the last captured state.
+ Restore()
+}
+
+// CheckpointableBuilder allows custom builders to participate in RecordBuilder
+// row rollback. The returned checkpoint is reused for every row.
+type CheckpointableBuilder interface {
+ // NewCheckpoint returns a reusable checkpoint for the builder.
+ NewCheckpoint() CheckpointState
+}
+
+type checkpointStateAdapter struct {
+ state CheckpointState
+}
+
+func (s *checkpointStateAdapter) capture() { s.state.Capture() }
+func (s *checkpointStateAdapter) restore() { s.state.Restore() }
+
+type storageBuilder interface {
+ StorageBuilder() Builder
+}
+
+type builderCheckpoint struct {
+ builder Builder
+ length int
+ children []*builderCheckpoint
+ state checkpointState
+ lastUnmarshalled interface{}
+ unmarshalled bool
+ lastStr *string
+}
+
+func newBuilderCheckpoint(builder Builder) *builderCheckpoint {
+ checkpoint := &builderCheckpoint{
+ builder: builder,
+ }
+ if checkpointable, ok := builder.(checkpointableBuilder); ok {
+ checkpoint.state = checkpointable.newCheckpoint()
+ } else if checkpointable, ok := builder.(CheckpointableBuilder); ok {
+ checkpoint.state = &checkpointStateAdapter{state:
checkpointable.NewCheckpoint()}
+ }
+
+ // Keep this switch in sync with builder types that own children. An
omitted
+ // nested builder would restore its own state but leave its children
changed.
+ switch builder := builder.(type) {
+ case *ListBuilder:
+ checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(builder.values))
+ case *LargeListBuilder:
+ checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(builder.values))
+ case *ListViewBuilder:
+ checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(builder.values))
+ case *LargeListViewBuilder:
+ checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(builder.values))
+ case *FixedSizeListBuilder:
+ checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(builder.values))
+ case *MapBuilder:
+ checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(builder.listBuilder))
+ case *StructBuilder:
+ for _, field := range builder.fields {
+ checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(field))
+ }
+ case *SparseUnionBuilder:
+ for _, child := range builder.children {
+ checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(child))
+ }
+ case *DenseUnionBuilder:
+ for _, child := range builder.children {
+ checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(child))
+ }
+ case storageBuilder:
+ checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(builder.StorageBuilder()))
+ case *RunEndEncodedBuilder:
+ checkpoint.children = append(checkpoint.children,
+ newBuilderCheckpoint(builder.runEnds),
+ newBuilderCheckpoint(builder.values),
+ )
+ }
+
+ checkpoint.capture()
+ return checkpoint
+}
+
+func (checkpoint *builderCheckpoint) capture() {
+ checkpoint.length = checkpoint.builder.Len()
+ if checkpoint.state != nil {
+ checkpoint.state.capture()
+ }
+ if builder, ok := checkpoint.builder.(*RunEndEncodedBuilder); ok {
+ checkpoint.lastUnmarshalled = builder.lastUnmarshalled
+ checkpoint.unmarshalled = builder.unmarshalled
+ checkpoint.lastStr = builder.lastStr
+ }
+ for _, child := range checkpoint.children {
+ child.capture()
+ }
+}
+
+func (checkpoint *builderCheckpoint) restore() {
+ if builder, ok := checkpoint.builder.(*RunEndEncodedBuilder); ok {
+ builder.length = checkpoint.length
+ builder.lastUnmarshalled = checkpoint.lastUnmarshalled
+ builder.unmarshalled = checkpoint.unmarshalled
+ builder.lastStr = checkpoint.lastStr
+ } else {
+ // Truncate the parent before restoring children. Some parent
builders
+ // truncate their children as part of truncation, so restoring
children
+ // first would overwrite their physical state again.
+ checkpoint.builder.truncate(checkpoint.length)
+ }
+ if checkpoint.state != nil {
+ checkpoint.state.restore()
+ }
+ for _, child := range checkpoint.children {
+ child.restore()
+ }
+}
+
// UnmarshalOne reads one row (a JSON object) from the supplied decoder and
// appends a value to each field in the RecordBuilder. Missing fields are
// appended as nulls and unrecognized keys are silently ignored.
@@ -424,6 +563,32 @@ func (b *RecordBuilder) NewRecord() arrow.Record {
// for nested field decoding. This is critical for preserving large integer
// values (>2^53) that cannot be represented exactly as float64.
func (b *RecordBuilder) UnmarshalOne(dec *json.Decoder) error {
+ if len(b.checkpoints) != len(b.fields) {
+ b.checkpoints = make([]*builderCheckpoint, len(b.fields))
+ for i, field := range b.fields {
+ b.checkpoints[i] = newBuilderCheckpoint(field)
+ }
+ } else {
+ for i, checkpoint := range b.checkpoints {
+ if checkpoint.builder != b.fields[i] {
+ b.checkpoints[i] =
newBuilderCheckpoint(b.fields[i])
+ }
+ }
+ }
+ for _, checkpoint := range b.checkpoints {
+ checkpoint.capture()
+ }
+ err := b.unmarshalOne(dec)
+ if err != nil {
+ for _, checkpoint := range b.checkpoints {
+ checkpoint.restore()
+ }
+ }
+ return err
+}
+
+func (b *RecordBuilder) unmarshalOne(dec *json.Decoder) (err error) {
+
// should start with a '{'
t, err := dec.Token()
if err != nil {
diff --git a/arrow/array/record_test.go b/arrow/array/record_test.go
index a3924382..6ad3c7e3 100644
--- a/arrow/array/record_test.go
+++ b/arrow/array/record_test.go
@@ -532,6 +532,521 @@ func TestRecordBuilder(t *testing.T) {
}
}
+func TestRecordBuilderRollsBackRowsAfterDecodeError(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "a", Type: arrow.PrimitiveTypes.Int32},
+ {Name: "b", Type: arrow.PrimitiveTypes.Int32},
+ }, nil)
+ b := array.NewRecordBuilder(mem, schema)
+ defer b.Release()
+
+ err := b.UnmarshalJSON([]byte("{\"a\":1,\"b\":\"invalid\"}"))
+ if err == nil {
+ t.Fatal("expected a decode error")
+ }
+ assert.Equal(t, 0, b.Field(0).Len())
+ assert.Equal(t, 0, b.Field(1).Len())
+
+ err = b.UnmarshalJSON([]byte("{\"a\":2,\"b\":3}"))
+ if err != nil {
+ t.Fatal(err)
+ }
+ rec := b.NewRecordBatch()
+ defer rec.Release()
+ assert.Equal(t, int64(1), rec.NumRows())
+
+}
+
+func TestRecordBuilderRollsBackVariableWidthState(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "value", Type: arrow.BinaryTypes.String},
+ {Name: "other", Type: arrow.PrimitiveTypes.Int32},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ large := strings.Repeat("x", 1<<20)
+ if err :=
builder.UnmarshalJSON([]byte(fmt.Sprintf(`{"value":%q,"other":"invalid"}`,
large))); err == nil {
+ t.Fatal("expected a decode error")
+ }
+ assert.Zero(t, builder.Field(0).(*array.StringBuilder).DataLen())
+
+ if err := builder.UnmarshalJSON([]byte(`{"value":"kept","other":1}`));
err != nil {
+ t.Fatal(err)
+ }
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+ assert.Equal(t, "kept", rec.Column(0).(*array.String).Value(0))
+ assert.NoError(t, array.ValidateFull(rec.Column(0)))
+}
+
+func TestRecordBuilderRollsBackVariableWidthData(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ for _, tc := range []struct {
+ name string
+ typ arrow.DataType
+ }{
+ {name: "string", typ: arrow.BinaryTypes.String},
+ {name: "large string", typ: arrow.BinaryTypes.LargeString},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "value", Type: tc.typ},
+ {Name: "other", Type:
arrow.PrimitiveTypes.Int32},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":"aaa","other":1}`)))
+ assert.Error(t,
builder.UnmarshalJSON([]byte(`{"value":"bbb","other":"invalid"}`)))
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":"ccc","other":2}`)))
+
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+ values := rec.Column(0).(array.StringLike)
+ assert.Equal(t, "aaa", values.Value(0))
+ assert.Equal(t, "ccc", values.Value(1))
+ assert.NoError(t, array.ValidateFull(rec.Column(0)))
+ })
+ }
+}
+
+func TestRecordBuilderRollsBackFixedSizeBinaryData(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "value", Type: &arrow.FixedSizeBinaryType{ByteWidth: 3}},
+ {Name: "other", Type: arrow.PrimitiveTypes.Int32},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":"YWFh","other":1}`)))
+ assert.Error(t,
builder.UnmarshalJSON([]byte(`{"value":"YmJi","other":"invalid"}`)))
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":"Y2Nj","other":2}`)))
+
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+ values := rec.Column(0).(*array.FixedSizeBinary)
+ assert.Equal(t, []byte("aaa"), values.Value(0))
+ assert.Equal(t, []byte("ccc"), values.Value(1))
+}
+
+func TestRecordBuilderRollsBackBooleanAndNullLengths(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ for _, tc := range []struct {
+ name string
+ typ arrow.DataType
+ }{
+ {name: "boolean", typ: arrow.FixedWidthTypes.Boolean},
+ {name: "null", typ: arrow.Null},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "value", Type: tc.typ},
+ {Name: "other", Type:
arrow.PrimitiveTypes.Int32},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ value := "true"
+ if tc.name == "null" {
+ value = "null"
+ }
+ assert.Error(t,
builder.UnmarshalJSON([]byte(fmt.Sprintf(`{"value":%s,"other":"invalid"}`,
value))))
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(fmt.Sprintf(`{"value":%s,"other":1}`, value))))
+
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+ assert.Equal(t, int64(1), rec.NumRows())
+ })
+ }
+}
+
+func TestRecordBuilderRollsBackDiscardedValidityBits(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "value", Type: arrow.PrimitiveTypes.Int32},
+ {Name: "other", Type: arrow.PrimitiveTypes.Int32},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":10,"other":1}`)))
+ assert.Error(t,
builder.UnmarshalJSON([]byte(`{"value":20,"other":"invalid"}`)))
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":null,"other":2}`)))
+
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+ values := rec.Column(0).(*array.Int32)
+ assert.Equal(t, int32(10), values.Value(0))
+ assert.True(t, values.IsNull(1))
+ assert.Equal(t, 1, values.NullN())
+ assert.NoError(t, array.ValidateFull(values))
+}
+
+func TestRecordBuilderRollsBackDictionaryState(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8,
ValueType: arrow.BinaryTypes.String}
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "value", Type: dictType},
+ {Name: "other", Type: arrow.PrimitiveTypes.Int32},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ if err :=
builder.UnmarshalJSON([]byte(`{"value":"discarded","other":"invalid"}`)); err
== nil {
+ t.Fatal("expected a decode error")
+ }
+ if err := builder.UnmarshalJSON([]byte(`{"value":"kept","other":1}`));
err != nil {
+ t.Fatal(err)
+ }
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+
+ dict := rec.Column(0).(*array.Dictionary)
+ assert.Equal(t, 1, dict.Dictionary().Len())
+ assert.Equal(t, "kept", dict.Dictionary().(*array.String).Value(0))
+ assert.NoError(t, array.ValidateFull(dict))
+}
+
+func TestRecordBuilderRollsBackDictionaryNullState(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8,
ValueType: arrow.BinaryTypes.String}
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "value", Type: dictType},
+ {Name: "other", Type: arrow.PrimitiveTypes.Int32},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ assert.Error(t,
builder.UnmarshalJSON([]byte(`{"value":null,"other":"invalid"}`)))
+ assert.Zero(t, builder.Field(0).Len())
+ assert.Zero(t, builder.Field(0).NullN())
+
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":"kept","other":1}`)))
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+ assert.Zero(t, rec.Column(0).NullN())
+}
+
+func TestRecordBuilderRollsBackExistingDictionaryState(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8,
ValueType: arrow.BinaryTypes.String}
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "value", Type: dictType},
+ {Name: "other", Type: arrow.PrimitiveTypes.Int32},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":"existing","other":1}`)))
+ assert.Error(t,
builder.UnmarshalJSON([]byte(`{"value":"discarded","other":"invalid"}`)))
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":"existing","other":2}`)))
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":"new","other":3}`)))
+
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+
+ dict := rec.Column(0).(*array.Dictionary)
+ dictValues := dict.Dictionary()
+ assert.Equal(t, 2, dictValues.Len())
+ assert.Equal(t, "existing", dictValues.(*array.String).Value(0))
+ assert.Equal(t, "new", dictValues.(*array.String).Value(1))
+ assert.Equal(t, 0, dict.GetValueIndex(0))
+ assert.Equal(t, 0, dict.GetValueIndex(1))
+ assert.Equal(t, 1, dict.GetValueIndex(2))
+ assert.Equal(t, "existing", dict.ValueStr(0))
+ assert.Equal(t, "existing", dict.ValueStr(1))
+ assert.Equal(t, "new", dict.ValueStr(2))
+ assert.NoError(t, array.ValidateFull(dict))
+}
+
+func TestRecordBuilderRollsBackStringViewBlocks(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ typ := arrow.ListOf(arrow.BinaryTypes.StringView)
+ schema := arrow.NewSchema([]arrow.Field{{Name: "value", Type: typ}},
nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ if err := builder.UnmarshalJSON([]byte(`{"value":["long string",1]}`));
err == nil {
+ t.Fatal("expected a decode error")
+ }
+ if err := builder.UnmarshalJSON([]byte(`{"value":["kept"]}`)); err !=
nil {
+ t.Fatal(err)
+ }
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+ list := rec.Column(0).(*array.List)
+ assert.Equal(t, "kept", list.ListValues().(*array.StringView).Value(0))
+ assert.NoError(t, array.ValidateFull(rec.Column(0)))
+}
+
+func TestRecordBuilderRollsBackAfterNewRecordBatch(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ typ := arrow.StructOf(arrow.Field{
+ Name: "items",
+ Type: arrow.ListOf(arrow.BinaryTypes.StringView),
+ })
+ schema := arrow.NewSchema([]arrow.Field{{Name: "row", Type: typ}}, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ if err := builder.UnmarshalJSON([]byte(`{"row":{"items":["first
retained value","second retained value"]}}`)); err != nil {
+ t.Fatal(err)
+ }
+ first := builder.NewRecordBatch()
+ defer first.Release()
+ firstBefore := first.Column(0).String()
+ if err := array.ValidateFull(first.Column(0)); err != nil {
+ t.Fatal(err)
+ }
+
+ if err := builder.UnmarshalJSON([]byte(`{"row":{"items":["discarded
retained value",1]}}`)); err == nil {
+ t.Fatal("expected a decode error")
+ }
+ if err := builder.UnmarshalJSON([]byte(`{"row":{"items":["replacement
retained value"]}}`)); err != nil {
+ t.Fatal(err)
+ }
+ second := builder.NewRecordBatch()
+ defer second.Release()
+
+ if err := array.ValidateFull(first.Column(0)); err != nil {
+ t.Fatal(err)
+ }
+ if err := array.ValidateFull(second.Column(0)); err != nil {
+ t.Fatal(err)
+ }
+ assert.Equal(t, firstBefore, first.Column(0).String())
+
+ firstRow := first.Column(0).(*array.Struct)
+ firstItems := firstRow.Field(0).(*array.List)
+ firstValues := firstItems.ListValues().(*array.StringView)
+ assert.Equal(t, 1, firstItems.Len())
+ assert.Equal(t, 2, firstValues.Len())
+ assert.Equal(t, "first retained value", firstValues.Value(0))
+ assert.Equal(t, "second retained value", firstValues.Value(1))
+
+ secondRow := second.Column(0).(*array.Struct)
+ secondItems := secondRow.Field(0).(*array.List)
+ secondValues := secondItems.ListValues().(*array.StringView)
+ assert.Equal(t, 1, secondItems.Len())
+ assert.Equal(t, 1, secondValues.Len())
+ assert.Equal(t, "replacement retained value", secondValues.Value(0))
+}
+
+func TestRecordBuilderRollsBackNestedRowsAfterDecodeError(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ tests := []struct {
+ name string
+ typ arrow.DataType
+ badValue string
+ goodValue string
+ check func(*testing.T, arrow.Array)
+ }{
+ {
+ name: "list view",
+ typ: arrow.ListViewOf(arrow.PrimitiveTypes.Int32),
+ badValue: `[1, "invalid"]`,
+ goodValue: `[7]`,
+ check: func(t *testing.T, value arrow.Array) {
+ list := value.(*array.ListView)
+ assert.Equal(t, 1, list.Len())
+ assert.Equal(t, int32(7),
list.ListValues().(*array.Int32).Value(0))
+ },
+ },
+ {
+ name: "large list",
+ typ:
arrow.LargeListOf(arrow.PrimitiveTypes.Int32),
+ badValue: `[1, "invalid"]`,
+ goodValue: `[8, 9]`,
+ check: func(t *testing.T, value arrow.Array) {
+ list := value.(*array.LargeList)
+ assert.Equal(t, 1, list.Len())
+ values := list.ListValues().(*array.Int32)
+ assert.Equal(t, int32(8), values.Value(0))
+ assert.Equal(t, int32(9), values.Value(1))
+ },
+ },
+ {
+ name: "fixed size list",
+ typ: arrow.FixedSizeListOf(2,
arrow.PrimitiveTypes.Int32),
+ badValue: `[1, "invalid"]`,
+ goodValue: `[10, 11]`,
+ check: func(t *testing.T, value arrow.Array) {
+ list := value.(*array.FixedSizeList)
+ assert.Equal(t, 1, list.Len())
+ values := list.ListValues().(*array.Int32)
+ assert.Equal(t, int32(10), values.Value(0))
+ assert.Equal(t, int32(11), values.Value(1))
+ },
+ },
+ {
+ name: "map",
+ typ: arrow.MapOf(arrow.BinaryTypes.String,
arrow.PrimitiveTypes.Int32),
+ badValue:
`[{"key":"discarded","value":1},{"key":"invalid","value":"invalid"}]`,
+ goodValue: `[{"key":"kept","value":12}]`,
+ check: func(t *testing.T, value arrow.Array) {
+ m := value.(*array.Map)
+ assert.Equal(t, 1, m.Len())
+ assert.Equal(t, "kept",
m.Keys().(*array.String).Value(0))
+ assert.Equal(t, int32(12),
m.Items().(*array.Int32).Value(0))
+ },
+ },
+ {
+ name: "sparse union",
+ typ: arrow.SparseUnionOf([]arrow.Field{{Name:
"value", Type: arrow.PrimitiveTypes.Int32}}, []arrow.UnionTypeCode{0}),
+ badValue: `[0, "invalid"]`,
+ goodValue: `[0, 13]`,
+ check: func(t *testing.T, value arrow.Array) {
+ union := value.(*array.SparseUnion)
+ assert.Equal(t, 1, union.Len())
+ assert.Equal(t, int32(13),
union.Field(0).(*array.Int32).Value(0))
+ },
+ },
+ {
+ name: "dense union",
+ typ: arrow.DenseUnionOf([]arrow.Field{{Name:
"value", Type: arrow.PrimitiveTypes.Int32}}, []arrow.UnionTypeCode{0}),
+ badValue: `[0, "invalid"]`,
+ goodValue: `[0, 14]`,
+ check: func(t *testing.T, value arrow.Array) {
+ union := value.(*array.DenseUnion)
+ assert.Equal(t, 1, union.Len())
+ assert.Equal(t, int32(14),
union.Field(0).(*array.Int32).Value(0))
+ },
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ schema := arrow.NewSchema([]arrow.Field{{Name: "value",
Type: tc.typ}}, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ if err := builder.UnmarshalJSON([]byte(`{"value":` +
tc.badValue + `}`)); err == nil {
+ t.Fatal("expected a decode error")
+ }
+ if err := builder.UnmarshalJSON([]byte(`{"value":` +
tc.goodValue + `}`)); err != nil {
+ t.Fatal(err)
+ }
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+ if err := array.ValidateFull(rec.Column(0)); err != nil
{
+ t.Fatal(err)
+ }
+ tc.check(t, rec.Column(0))
+ })
+ }
+}
+
+func TestRecordBuilderRollsBackRunEndStateAfterDecodeError(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ typ := arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int16,
arrow.PrimitiveTypes.Int32)
+ schema := arrow.NewSchema([]arrow.Field{{Name: "value", Type: typ}},
nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ for i := 0; i < 2; i++ {
+ if err := builder.UnmarshalJSON([]byte(`{"value":"invalid"}`));
err == nil {
+ t.Fatal("expected a decode error")
+ }
+ }
+ if got := builder.Field(0).Len(); got != 0 {
+ t.Fatalf("builder length = %d, want 0", got)
+ }
+}
+
+func TestRecordBuilderRefreshesCheckpointsAfterFieldReplacement(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{{Name: "value", Type:
arrow.PrimitiveTypes.Int32}}, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ fields := builder.Fields()
+ old := fields[0]
+ fields[0] = array.NewInt32Builder(mem)
+ old.Release()
+
+ if err := builder.UnmarshalJSON([]byte(`{"value":1}`)); err != nil {
+ t.Fatal(err)
+ }
+ if err := builder.UnmarshalJSON([]byte(`{"value":"invalid"}`)); err ==
nil {
+ t.Fatal("expected a decode error")
+ }
+ if err := builder.UnmarshalJSON([]byte(`{"value":2}`)); err != nil {
+ t.Fatal(err)
+ }
+
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+ values := rec.Column(0).(*array.Int32)
+ assert.Equal(t, 2, values.Len())
+ assert.Equal(t, int32(1), values.Value(0))
+ assert.Equal(t, int32(2), values.Value(1))
+}
+
+func TestRecordBuilderRollsBackNestedRunEndStateAfterDecodeError(t *testing.T)
{
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ typ := arrow.StructOf(
+ arrow.Field{Name: "value", Type:
arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int16, arrow.BinaryTypes.String)},
+ arrow.Field{Name: "other", Type: arrow.PrimitiveTypes.Int32},
+ )
+ schema := arrow.NewSchema([]arrow.Field{{Name: "row", Type: typ}}, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ if err :=
builder.UnmarshalJSON([]byte(`{"row":{"value":"a","other":1}}`)); err != nil {
+ t.Fatal(err)
+ }
+ if err :=
builder.UnmarshalJSON([]byte(`{"row":{"value":"a","other":"invalid"}}`)); err
== nil {
+ t.Fatal("expected a decode error")
+ }
+ if err :=
builder.UnmarshalJSON([]byte(`{"row":{"value":"a","other":2}}`)); err != nil {
+ t.Fatal(err)
+ }
+
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+ row := rec.Column(0).(*array.Struct)
+ rle := row.Field(0).(*array.RunEndEncoded)
+ assert.Equal(t, 2, rle.Len())
+ assert.Equal(t, 1, rle.RunEndsArr().Len())
+ assert.Equal(t, 1, rle.Values().Len())
+ assert.Equal(t, "a", rle.ValueStr(0))
+ assert.Equal(t, "a", rle.ValueStr(1))
+}
+
func TestRecordBuilderResize(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
defer mem.AssertSize(t, 0)
diff --git a/arrow/array/union.go b/arrow/array/union.go
index 8b86947d..ed894ab1 100644
--- a/arrow/array/union.go
+++ b/arrow/array/union.go
@@ -945,6 +945,10 @@ func (b *SparseUnionBuilder) Resize(n int) {
b.typesBuilder.resize(n)
}
+func (b *SparseUnionBuilder) truncate(n int) {
+ b.typesBuilder.SetLength(n)
+}
+
// AppendNull will append a null to the first child and an empty value
// (implementation-defined) to the rest of the children.
func (b *SparseUnionBuilder) AppendNull() {
@@ -1186,6 +1190,11 @@ func (b *DenseUnionBuilder) Resize(n int) {
b.offsetsBuilder.resize(n * arrow.Int32SizeBytes)
}
+func (b *DenseUnionBuilder) truncate(n int) {
+ b.typesBuilder.SetLength(n)
+ b.offsetsBuilder.SetLength(n * arrow.Int32SizeBytes)
+}
+
// AppendNull will only append a null value arbitrarily to the first child
// and use that offset for this element of the array.
func (b *DenseUnionBuilder) AppendNull() {
diff --git a/arrow/extensions/timestamp_with_offset.go
b/arrow/extensions/timestamp_with_offset.go
index 47f8d241..5cc17ed8 100644
--- a/arrow/extensions/timestamp_with_offset.go
+++ b/arrow/extensions/timestamp_with_offset.go
@@ -439,6 +439,24 @@ func (b *TimestampWithOffsetBuilder) NewArray()
arrow.Array {
return b.NewExtensionArray()
}
+type timestampWithOffsetCheckpoint struct {
+ builder *TimestampWithOffsetBuilder
+ lastOffset int16
+}
+
+func (c *timestampWithOffsetCheckpoint) Capture() {
+ c.lastOffset = c.builder.lastOffset
+}
+
+func (c *timestampWithOffsetCheckpoint) Restore() {
+ c.builder.lastOffset = c.lastOffset
+}
+
+// NewCheckpoint returns a checkpoint for the builder's run-end offset state.
+func (b *TimestampWithOffsetBuilder) NewCheckpoint() array.CheckpointState {
+ return ×tampWithOffsetCheckpoint{builder: b}
+}
+
// NewExtensionArray finalizes the current array and resets lastOffset so a
// reused builder starts a fresh run instead of continuing a run that belonged
// to the array just finalized (the underlying REE builder is reset too).
diff --git a/arrow/extensions/timestamp_with_offset_test.go
b/arrow/extensions/timestamp_with_offset_test.go
index efbdfd28..26221dc9 100644
--- a/arrow/extensions/timestamp_with_offset_test.go
+++ b/arrow/extensions/timestamp_with_offset_test.go
@@ -371,6 +371,23 @@ func
TestTimestampWithOffsetBuilderRunEndEncodedNullContinuesRun(t *testing.T) {
assert.Equal(t, testDate1, typedArr.Value(2))
}
+func TestTimestampWithOffsetBuilderRunEndEncodedResizeContinuesRun(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ builder, err := extensions.NewTimestampWithOffsetBuilder(mem,
testTimeUnit, ree(arrow.PrimitiveTypes.Int16))
+ require.NoError(t, err)
+
+ builder.Append(testDate1)
+ builder.Resize(builder.Cap() * 2)
+ builder.Append(testDate1)
+
+ arr := builder.NewArray()
+ defer arr.Release()
+ offsets :=
arr.(*extensions.TimestampWithOffsetArray).Storage().(*array.Struct).Field(1).(*array.RunEndEncoded)
+ assert.Equal(t, 1, offsets.Values().Len())
+}
+
func TestTimestampWithOffsetBuilderAppendValuesNilValids(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
defer mem.AssertSize(t, 0)
@@ -600,6 +617,36 @@ func TestTimestampWithOffsetExtensionRecordBuilder(t
*testing.T) {
}
}
+func TestTimestampWithOffsetExtensionRecordBuilderRollsBackState(t *testing.T)
{
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ dataType, err :=
extensions.NewTimestampWithOffsetTypeCustomOffset(testTimeUnit,
ree(arrow.PrimitiveTypes.Int16))
+ require.NoError(t, err)
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "timestamp_with_offset", Type: dataType},
+ {Name: "other", Type: arrow.PrimitiveTypes.Int32},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ require.NoError(t,
builder.UnmarshalJSON([]byte(`{"timestamp_with_offset":"2025-01-01T00:00:00+01:00","other":1}`)))
+ require.Error(t,
builder.UnmarshalJSON([]byte(`{"timestamp_with_offset":"2025-01-01T00:00:00+01:00","other":"invalid"}`)))
+ require.NoError(t,
builder.UnmarshalJSON([]byte(`{"timestamp_with_offset":"2025-01-01T00:00:00+01:00","other":2}`)))
+
+ rec := builder.NewRecordBatch()
+ defer rec.Release()
+
+ values := rec.Column(0).(*extensions.TimestampWithOffsetArray)
+ require.Equal(t, 2, values.Len())
+ _, offset := values.Value(0).Zone()
+ require.Equal(t, 60*60, offset)
+ _, offset = values.Value(1).Zone()
+ require.Equal(t, 60*60, offset)
+ offsets :=
values.Storage().(*array.Struct).Field(1).(*array.RunEndEncoded)
+ require.Equal(t, 1, offsets.RunEndsArr().Len())
+}
+
func TestTimestampWithOffsetTypeBatchIPCRoundTrip(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
defer mem.AssertSize(t, 0)
diff --git a/internal/hashing/xxh3_memo_table.go
b/internal/hashing/xxh3_memo_table.go
index 26b4ce64..728092e6 100644
--- a/internal/hashing/xxh3_memo_table.go
+++ b/internal/hashing/xxh3_memo_table.go
@@ -42,6 +42,8 @@ type MemoTable interface {
TypeTraits() TypeTraits
// Reset drops everything in the table allowing it to be reused
Reset()
+ // Truncate removes values with an index greater than or equal to size.
+ Truncate(size int)
// Size returns the current number of unique values stored in
// the table, including whether or not a null value has been
// inserted via GetOrInsertNull.
@@ -167,6 +169,35 @@ func (s *BinaryMemoTable) Reset() {
s.nullIdx = KeyNotFound
}
+func (s *BinaryMemoTable) Truncate(size int) {
+ if size < 0 {
+ panic("cannot truncate a memo table to a negative size")
+ }
+ if size >= s.Size() {
+ return
+ }
+
+ dataLen := 0
+ for i := 0; i < size; i++ {
+ dataLen += len(s.builder.Value(i))
+ }
+ s.builder.Resize(size)
+ s.builder.ResizeData(dataLen)
+
+ truncated := NewHashTable[int32](uint64(size))
+ s.tbl.VisitEntries(func(e *entry[int32]) {
+ if e.payload.val < 0 || uint64(e.payload.val) >= uint64(size) {
+ return
+ }
+ entry, _ := truncated.Lookup(e.h, func(int32) bool { return
false })
+ truncated.Insert(entry, e.h, e.payload.val, -1)
+ })
+ s.tbl = truncated
+ if s.nullIdx >= size {
+ s.nullIdx = KeyNotFound
+ }
+}
+
// GetNull returns the index of a null that has been inserted into the table or
// KeyNotFound. The bool returned will be true if there was a null inserted
into
// the table, and false otherwise.
diff --git a/internal/hashing/xxh3_memo_table_test.go
b/internal/hashing/xxh3_memo_table_test.go
new file mode 100644
index 00000000..d248fdc5
--- /dev/null
+++ b/internal/hashing/xxh3_memo_table_test.go
@@ -0,0 +1,226 @@
+// 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 hashing_test
+
+import (
+ "fmt"
+ "math"
+ "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/internal/hashing"
+ "github.com/stretchr/testify/assert"
+)
+
+func TestMemoTableTruncate(t *testing.T) {
+ table := hashing.NewMemoTable[int32](0)
+
+ idx, found, err := table.GetOrInsert(int32(10))
+ assert.NoError(t, err)
+ assert.False(t, found)
+ assert.Equal(t, 0, idx)
+ idx, found, err = table.GetOrInsert(int32(20))
+ assert.NoError(t, err)
+ assert.False(t, found)
+ assert.Equal(t, 1, idx)
+ idx, found, err = table.GetOrInsert(int32(30))
+ assert.NoError(t, err)
+ assert.False(t, found)
+ assert.Equal(t, 2, idx)
+
+ table.Truncate(2)
+ assert.Equal(t, 2, table.Size())
+ assertGetMemoValue(t, table, int32(10), 0)
+ assertGetMemoValue(t, table, int32(20), 1)
+ _, found = table.Get(int32(30))
+ assert.False(t, found)
+
+ idx, found, err = table.GetOrInsert(int32(30))
+ assert.NoError(t, err)
+ assert.False(t, found)
+ assert.Equal(t, 2, idx)
+
+ table.Truncate(0)
+ assert.Equal(t, 0, table.Size())
+ _, found = table.Get(int32(10))
+ assert.False(t, found)
+
+ idx, found, err = table.GetOrInsert(int32(40))
+ assert.NoError(t, err)
+ assert.False(t, found)
+ assert.Equal(t, 0, idx)
+}
+
+func TestMemoTableTruncateNoOp(t *testing.T) {
+ table := hashing.NewMemoTable[int32](0)
+
+ idx, found, err := table.GetOrInsert(int32(7))
+ assert.NoError(t, err)
+ assert.False(t, found)
+ assert.Equal(t, 0, idx)
+
+ table.Truncate(1)
+ table.Truncate(math.MaxInt)
+
+ assert.Equal(t, 1, table.Size())
+ assertGetMemoValue(t, table, int32(7), 0)
+}
+
+func assertGetMemoValue(t *testing.T, table *hashing.Table[int32], value
int32, want int) {
+ t.Helper()
+ got, found := table.Get(value)
+ assert.True(t, found)
+ assert.Equal(t, want, got)
+}
+
+func TestBinaryMemoTableTruncate(t *testing.T) {
+ t.Run("reinsert discarded values", func(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+ table := hashing.NewBinaryMemoTable(0, -1,
array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary))
+ defer table.Release()
+
+ insertBinaryMemoValue(t, table, "a", 0)
+ insertBinaryMemoNull(t, table, 1)
+ insertBinaryMemoValue(t, table, "discarded", 2)
+
+ table.Truncate(2)
+ assert.Equal(t, 2, table.Size())
+ assertGetBinaryMemoValue(t, table, "a", 0)
+ assertGetBinaryMemoNull(t, table, 1)
+ _, found := table.Get("discarded")
+ assert.False(t, found)
+
+ insertBinaryMemoValue(t, table, "discarded", 2)
+ assertGetBinaryMemoValue(t, table, "discarded", 2)
+ })
+
+ t.Run("retains or removes null at the truncation boundary", func(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+ table := hashing.NewBinaryMemoTable(0, -1,
array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary))
+ defer table.Release()
+
+ insertBinaryMemoValue(t, table, "kept", 0)
+ insertBinaryMemoNull(t, table, 1)
+ insertBinaryMemoValue(t, table, "removed", 2)
+
+ table.Truncate(1)
+ assert.Equal(t, 1, table.Size())
+ _, found := table.GetNull()
+ assert.False(t, found)
+ insertBinaryMemoNull(t, table, 1)
+ assertGetBinaryMemoNull(t, table, 1)
+
+ table.Truncate(0)
+ assert.Equal(t, 0, table.Size())
+ _, found = table.GetNull()
+ assert.False(t, found)
+ insertBinaryMemoValue(t, table, "after zero", 0)
+ assertGetBinaryMemoValue(t, table, "after zero", 0)
+ })
+
+ t.Run("rewinds binary data and offsets", func(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+ table := hashing.NewBinaryMemoTable(0, -1,
array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary))
+ defer table.Release()
+
+ insertBinaryMemoValue(t, table, "alpha", 0)
+ insertBinaryMemoValue(t, table, "bravo", 1)
+ insertBinaryMemoValue(t, table, "discarded", 2)
+
+ table.Truncate(2)
+ assert.Equal(t, len("alphabravo"), table.ValuesSize())
+ assert.Equal(t, []int32{0, 5, 10}, binaryMemoOffsets(table))
+
+ insertBinaryMemoValue(t, table, "charlie", 2)
+ assert.Equal(t, len("alphabravocharlie"), table.ValuesSize())
+ assert.Equal(t, []int32{0, 5, 10, 17}, binaryMemoOffsets(table))
+ values := make([]byte, table.ValuesSize())
+ table.CopyValues(values)
+ assert.Equal(t, "alphabravocharlie", string(values))
+ })
+
+ t.Run("preserves hash probing", func(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+ table := hashing.NewBinaryMemoTable(0, -1,
array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary))
+ defer table.Release()
+
+ first, second := findHashBucketCollision()
+ insertBinaryMemoValue(t, table, first, 0)
+ insertBinaryMemoValue(t, table, second, 1)
+
+ table.Truncate(1)
+ assertGetBinaryMemoValue(t, table, first, 0)
+ _, found := table.Get(second)
+ assert.False(t, found)
+
+ insertBinaryMemoValue(t, table, second, 1)
+ assertGetBinaryMemoValue(t, table, second, 1)
+ })
+}
+
+func findHashBucketCollision() (string, string) {
+ seen := make(map[uint64]string)
+ for i := 0; ; i++ {
+ value := fmt.Sprintf("collision-%d", i)
+ bucket := hashing.Hash([]byte(value), 0) & 31
+ if previous, ok := seen[bucket]; ok {
+ return previous, value
+ }
+ seen[bucket] = value
+ }
+}
+
+func insertBinaryMemoValue(t *testing.T, table *hashing.BinaryMemoTable, value
string, want int) {
+ t.Helper()
+ idx, found, err := table.GetOrInsert(value)
+ assert.NoError(t, err)
+ assert.False(t, found)
+ assert.Equal(t, want, idx)
+}
+
+func insertBinaryMemoNull(t *testing.T, table *hashing.BinaryMemoTable, want
int) {
+ t.Helper()
+ idx, found := table.GetOrInsertNull()
+ assert.False(t, found)
+ assert.Equal(t, want, idx)
+}
+
+func assertGetBinaryMemoValue(t *testing.T, table *hashing.BinaryMemoTable,
value string, want int) {
+ t.Helper()
+ idx, found := table.Get(value)
+ assert.True(t, found)
+ assert.Equal(t, want, idx)
+}
+
+func assertGetBinaryMemoNull(t *testing.T, table *hashing.BinaryMemoTable,
want int) {
+ t.Helper()
+ idx, found := table.GetNull()
+ assert.True(t, found)
+ assert.Equal(t, want, idx)
+}
+
+func binaryMemoOffsets(table *hashing.BinaryMemoTable) []int32 {
+ offsets := make([]int32, table.Size()+1)
+ table.CopyOffsets(offsets)
+ return offsets
+}
diff --git a/internal/hashing/xxh3_memo_table_types.go
b/internal/hashing/xxh3_memo_table_types.go
index 715be652..91e30fe5 100644
--- a/internal/hashing/xxh3_memo_table_types.go
+++ b/internal/hashing/xxh3_memo_table_types.go
@@ -67,6 +67,18 @@ func (h *HashTable[T]) Reset(cap uint64) {
h.entries = make([]entry[T], h.cap)
}
+func (h *HashTable[T]) Truncate(size uint64) {
+ truncated := NewHashTable[T](size)
+ h.VisitEntries(func(e *entry[T]) {
+ if uint64(e.payload.memoIdx) >= size {
+ return
+ }
+ entry, _ := truncated.Lookup(e.h, func(T) bool { return false })
+ truncated.Insert(entry, e.h, e.payload.val, e.payload.memoIdx)
+ })
+ *h = *truncated
+}
+
func (h *HashTable[T]) CopyValues(out []T) {
h.CopyValuesSubset(0, out)
}
@@ -190,6 +202,19 @@ func (t *Table[T]) Reset() {
t.nullIdx = KeyNotFound
}
+func (t *Table[T]) Truncate(size int) {
+ if size < 0 {
+ panic("cannot truncate a memo table to a negative size")
+ }
+ if size >= t.Size() {
+ return
+ }
+ t.tbl.Truncate(uint64(size))
+ if t.nullIdx >= int32(size) {
+ t.nullIdx = KeyNotFound
+ }
+}
+
func (t *Table[T]) Size() int {
sz := int(t.tbl.size)
if _, ok := t.GetNull(); ok {
