arrow/array/data.go - arrow-go - Git at Google

 // 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 array

 import (
 	"encoding/binary"
 	"sync/atomic"

 	"github.com/apache/arrow-go/v18/arrow"
 	"github.com/apache/arrow-go/v18/arrow/internal/debug"
 	"github.com/apache/arrow-go/v18/arrow/memory"
 	"github.com/apache/arrow-go/v18/internal/utils/maphash"
 )

 // Data represents the memory and metadata of an Arrow array.
 type Data struct {
 	refCount atomic.Int64
 	dtype    arrow.DataType
 	nulls    int
 	offset   int
 	length   int

 	// for dictionary arrays: buffers will be the null validity bitmap and the indexes that reference
 	// values in the dictionary member. childData would be empty in a dictionary array
 	buffers    []*memory.Buffer  // TODO(sgc): should this be an interface?
 	childData  []arrow.ArrayData // TODO(sgc): managed by ListArray, StructArray and UnionArray types
 	dictionary *Data             // only populated for dictionary arrays
 }

 // NewData creates a new Data.
 func NewData(dtype arrow.DataType, length int, buffers []*memory.Buffer, childData []arrow.ArrayData, nulls, offset int) *Data {
 	for _, b := range buffers {
 		if b != nil {
 			b.Retain()
 		}
 	}

 	for _, child := range childData {
 		if child != nil {
 			child.Retain()
 		}
 	}

 	d := &Data{
 		dtype:     dtype,
 		nulls:     nulls,
 		length:    length,
 		offset:    offset,
 		buffers:   buffers,
 		childData: childData,
 	}
 	d.refCount.Add(1)
 	return d
 }

 // NewDataWithDictionary creates a new data object, but also sets the provided dictionary into the data if it's not nil
 func NewDataWithDictionary(dtype arrow.DataType, length int, buffers []*memory.Buffer, nulls, offset int, dict *Data) *Data {
 	data := NewData(dtype, length, buffers, nil, nulls, offset)
 	if dict != nil {
 		dict.Retain()
 	}
 	data.dictionary = dict
 	return data
 }

 func (d *Data) Copy() *Data {
 	// don't pass the slices directly, otherwise it retains the connection
 	// we need to make new slices and populate them with the same pointers
 	bufs := make([]*memory.Buffer, len(d.buffers))
 	copy(bufs, d.buffers)
 	children := make([]arrow.ArrayData, len(d.childData))
 	copy(children, d.childData)

 	data := NewData(d.dtype, d.length, bufs, children, d.nulls, d.offset)
 	data.SetDictionary(d.dictionary)
 	return data
 }

 // Reset sets the Data for re-use.
 func (d *Data) Reset(dtype arrow.DataType, length int, buffers []*memory.Buffer, childData []arrow.ArrayData, nulls, offset int) {
 	// Retain new buffers before releasing existing buffers in-case they're the same ones to prevent accidental premature
 	// release.
 	for _, b := range buffers {
 		if b != nil {
 			b.Retain()
 		}
 	}
 	for _, b := range d.buffers {
 		if b != nil {
 			b.Release()
 		}
 	}
 	d.buffers = buffers

 	// Retain new children data before releasing existing children data in-case they're the same ones to prevent accidental
 	// premature release.
 	for _, d := range childData {
 		if d != nil {
 			d.Retain()
 		}
 	}
 	for _, d := range d.childData {
 		if d != nil {
 			d.Release()
 		}
 	}
 	d.childData = childData

 	if d.dictionary != nil {
 		d.dictionary.Release()
 		d.dictionary = nil
 	}

 	d.dtype = dtype
 	d.length = length
 	d.nulls = nulls
 	d.offset = offset
 }

 // Retain increases the reference count by 1.
 // Retain may be called simultaneously from multiple goroutines.
 func (d *Data) Retain() {
 	d.refCount.Add(1)
 }

 // Release decreases the reference count by 1.
 // When the reference count goes to zero, the memory is freed.
 // Release may be called simultaneously from multiple goroutines.
 func (d *Data) Release() {
 	debug.Assert(d.refCount.Load() > 0, "too many releases")

 	if d.refCount.Add(-1) == 0 {
 		for _, b := range d.buffers {
 			if b != nil {
 				b.Release()
 			}
 		}

 		for _, b := range d.childData {
 			if b != nil {
 				b.Release()
 			}
 		}

 		if d.dictionary != nil {
 			d.dictionary.Release()
 		}
 		d.dictionary, d.buffers, d.childData = nil, nil, nil
 	}
 }

 // DataType returns the DataType of the data.
 func (d *Data) DataType() arrow.DataType { return d.dtype }

 func (d *Data) SetNullN(n int) { d.nulls = n }

 // NullN returns the number of nulls.
 func (d *Data) NullN() int { return d.nulls }

 // Len returns the length.
 func (d *Data) Len() int { return d.length }

 // Offset returns the offset.
 func (d *Data) Offset() int { return d.offset }

 // Buffers returns the buffers.
 func (d *Data) Buffers() []*memory.Buffer { return d.buffers }

 func (d *Data) Children() []arrow.ArrayData { return d.childData }

 // Dictionary returns the ArrayData object for the dictionary member, or nil
 func (d *Data) Dictionary() arrow.ArrayData { return d.dictionary }

 // SetDictionary allows replacing the dictionary for this particular Data object
 func (d *Data) SetDictionary(dict arrow.ArrayData) {
 	var newDict *Data
 	if dict != nil {
 		var ok bool
 		newDict, ok = dict.(*Data)
 		if !ok {
 			panic("arrow/array: dictionary data must be *array.Data")
 		}
 	}

 	if newDict != nil {
 		newDict.Retain()
 	}
 	oldDict := d.dictionary
 	d.dictionary = newDict
 	if oldDict != nil {
 		oldDict.Release()
 	}
 }

 // SizeInBytes returns the size of the Data and any children and/or dictionary in bytes by
 // recursively examining the nested structures of children and/or dictionary.
 // The value returned is an upper-bound since offset is not taken into account.
 func (d *Data) SizeInBytes() uint64 {
 	var size uint64

 	if d == nil {
 		return 0
 	}

 	for _, b := range d.Buffers() {
 		if b != nil {
 			size += uint64(b.Len())
 		}
 	}
 	for _, c := range d.Children() {
 		size += c.SizeInBytes()
 	}
 	if d.dictionary != nil {
 		size += d.dictionary.SizeInBytes()
 	}

 	return size
 }

 // NewSliceData returns a new slice that shares backing data with the input.
 // The returned Data slice starts at i and extends j-i elements, such as:
 //
 //	slice := data[i:j]
 //
 // The returned value must be Release'd after use.
 //
 // NewSliceData panics if the slice is outside the valid range of the input Data.
 // NewSliceData panics if j < i.
 func NewSliceData(data arrow.ArrayData, i, j int64) arrow.ArrayData {
 	if i < 0 || j < 0 || j > int64(data.Len()) || i > j || data.Offset()+int(i) > data.Offset()+data.Len() {
 		panic("arrow/array: index out of range")
 	}

 	for _, b := range data.Buffers() {
 		if b != nil {
 			b.Retain()
 		}
 	}

 	for _, child := range data.Children() {
 		if child != nil {
 			child.Retain()
 		}
 	}

 	if data.(*Data).dictionary != nil {
 		data.(*Data).dictionary.Retain()
 	}

 	o := &Data{
 		dtype:      data.DataType(),
 		nulls:      UnknownNullCount,
 		length:     int(j - i),
 		offset:     data.Offset() + int(i),
 		buffers:    data.Buffers(),
 		childData:  data.Children(),
 		dictionary: data.(*Data).dictionary,
 	}
 	o.refCount.Add(1)

 	if data.NullN() == 0 {
 		o.nulls = 0
 	}

 	return o
 }

 func Hash(h *maphash.MapHash, data arrow.ArrayData) {
 	a := data.(*Data)

 	hashString(h, a.dtype.Fingerprint())
 	hashInt(h, a.length)
 	hashInt(h, a.offset)

 	hashInt(h, len(a.buffers))
 	for _, b := range a.buffers {
 		if b == nil {
 			hashByte(h, 0)
 			continue
 		}

 		hashByte(h, 1)
 		hashBytes(h, b.Bytes())
 	}

 	hashInt(h, len(a.childData))
 	for _, c := range a.childData {
 		if c == nil {
 			hashByte(h, 0)
 			continue
 		}

 		hashByte(h, 1)
 		Hash(h, c)
 	}

 	if a.dictionary == nil {
 		hashByte(h, 0)
 		return
 	}

 	hashByte(h, 1)
 	Hash(h, a.dictionary)
 }

 func hashByte(h *maphash.MapHash, v byte) {
 	var b [1]byte
 	b[0] = v
 	h.Write(b[:])
 }

 func hashInt(h *maphash.MapHash, v int) {
 	var b [8]byte
 	binary.LittleEndian.PutUint64(b[:], uint64(int64(v)))
 	h.Write(b[:])
 }

 func hashBytes(h *maphash.MapHash, b []byte) {
 	hashInt(h, len(b))
 	h.Write(b)
 }

 func hashString(h *maphash.MapHash, s string) {
 	hashInt(h, len(s))
 	h.WriteString(s)
 }
	// 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 array

	import (
	"encoding/binary"
	"sync/atomic"

	"github.com/apache/arrow-go/v18/arrow"
	"github.com/apache/arrow-go/v18/arrow/internal/debug"
	"github.com/apache/arrow-go/v18/arrow/memory"
	"github.com/apache/arrow-go/v18/internal/utils/maphash"
	)

	// Data represents the memory and metadata of an Arrow array.
	type Data struct {
	refCount atomic.Int64
	dtype arrow.DataType
	nulls int
	offset int
	length int

	// for dictionary arrays: buffers will be the null validity bitmap and the indexes that reference
	// values in the dictionary member. childData would be empty in a dictionary array
	buffers []*memory.Buffer // TODO(sgc): should this be an interface?
	childData []arrow.ArrayData // TODO(sgc): managed by ListArray, StructArray and UnionArray types
	dictionary *Data // only populated for dictionary arrays
	}

	// NewData creates a new Data.
	func NewData(dtype arrow.DataType, length int, buffers []memory.Buffer, childData []arrow.ArrayData, nulls, offset int) Data {
	for _, b := range buffers {
	if b != nil {
	b.Retain()
	}
	}

	for _, child := range childData {
	if child != nil {
	child.Retain()
	}
	}

	d := &Data{
	dtype: dtype,
	nulls: nulls,
	length: length,
	offset: offset,
	buffers: buffers,
	childData: childData,
	}
	d.refCount.Add(1)
	return d
	}

	// NewDataWithDictionary creates a new data object, but also sets the provided dictionary into the data if it's not nil
	func NewDataWithDictionary(dtype arrow.DataType, length int, buffers []memory.Buffer, nulls, offset int, dict Data) *Data {
	data := NewData(dtype, length, buffers, nil, nulls, offset)
	if dict != nil {
	dict.Retain()
	}
	data.dictionary = dict
	return data
	}

	func (d Data) Copy() Data {
	// don't pass the slices directly, otherwise it retains the connection
	// we need to make new slices and populate them with the same pointers
	bufs := make([]*memory.Buffer, len(d.buffers))
	copy(bufs, d.buffers)
	children := make([]arrow.ArrayData, len(d.childData))
	copy(children, d.childData)

	data := NewData(d.dtype, d.length, bufs, children, d.nulls, d.offset)
	data.SetDictionary(d.dictionary)
	return data
	}

	// Reset sets the Data for re-use.
	func (d Data) Reset(dtype arrow.DataType, length int, buffers []memory.Buffer, childData []arrow.ArrayData, nulls, offset int) {
	// Retain new buffers before releasing existing buffers in-case they're the same ones to prevent accidental premature
	// release.
	for _, b := range buffers {
	if b != nil {
	b.Retain()
	}
	}
	for _, b := range d.buffers {
	if b != nil {
	b.Release()
	}
	}
	d.buffers = buffers

	// Retain new children data before releasing existing children data in-case they're the same ones to prevent accidental
	// premature release.
	for _, d := range childData {
	if d != nil {
	d.Retain()
	}
	}
	for _, d := range d.childData {
	if d != nil {
	d.Release()
	}
	}
	d.childData = childData

	if d.dictionary != nil {
	d.dictionary.Release()
	d.dictionary = nil
	}

	d.dtype = dtype
	d.length = length
	d.nulls = nulls
	d.offset = offset
	}

	// Retain increases the reference count by 1.
	// Retain may be called simultaneously from multiple goroutines.
	func (d *Data) Retain() {
	d.refCount.Add(1)
	}

	// Release decreases the reference count by 1.
	// When the reference count goes to zero, the memory is freed.
	// Release may be called simultaneously from multiple goroutines.
	func (d *Data) Release() {
	debug.Assert(d.refCount.Load() > 0, "too many releases")

	if d.refCount.Add(-1) == 0 {
	for _, b := range d.buffers {
	if b != nil {
	b.Release()
	}
	}

	for _, b := range d.childData {
	if b != nil {
	b.Release()
	}
	}

	if d.dictionary != nil {
	d.dictionary.Release()
	}
	d.dictionary, d.buffers, d.childData = nil, nil, nil
	}
	}

	// DataType returns the DataType of the data.
	func (d *Data) DataType() arrow.DataType { return d.dtype }

	func (d *Data) SetNullN(n int) { d.nulls = n }

	// NullN returns the number of nulls.
	func (d *Data) NullN() int { return d.nulls }

	// Len returns the length.
	func (d *Data) Len() int { return d.length }

	// Offset returns the offset.
	func (d *Data) Offset() int { return d.offset }

	// Buffers returns the buffers.
	func (d Data) Buffers() []memory.Buffer { return d.buffers }

	func (d *Data) Children() []arrow.ArrayData { return d.childData }

	// Dictionary returns the ArrayData object for the dictionary member, or nil
	func (d *Data) Dictionary() arrow.ArrayData { return d.dictionary }

	// SetDictionary allows replacing the dictionary for this particular Data object
	func (d *Data) SetDictionary(dict arrow.ArrayData) {
	var newDict *Data
	if dict != nil {
	var ok bool
	newDict, ok = dict.(*Data)
	if !ok {
	panic("arrow/array: dictionary data must be *array.Data")
	}
	}

	if newDict != nil {
	newDict.Retain()
	}
	oldDict := d.dictionary
	d.dictionary = newDict
	if oldDict != nil {
	oldDict.Release()
	}
	}

	// SizeInBytes returns the size of the Data and any children and/or dictionary in bytes by
	// recursively examining the nested structures of children and/or dictionary.
	// The value returned is an upper-bound since offset is not taken into account.
	func (d *Data) SizeInBytes() uint64 {
	var size uint64

	if d == nil {
	return 0
	}

	for _, b := range d.Buffers() {
	if b != nil {
	size += uint64(b.Len())
	}
	}
	for _, c := range d.Children() {
	size += c.SizeInBytes()
	}
	if d.dictionary != nil {
	size += d.dictionary.SizeInBytes()
	}

	return size
	}

	// NewSliceData returns a new slice that shares backing data with the input.
	// The returned Data slice starts at i and extends j-i elements, such as:
	//
	// slice := data[i:j]
	//
	// The returned value must be Release'd after use.
	//
	// NewSliceData panics if the slice is outside the valid range of the input Data.
	// NewSliceData panics if j < i.
	func NewSliceData(data arrow.ArrayData, i, j int64) arrow.ArrayData {
	if i < 0 \|\| j < 0 \|\| j > int64(data.Len()) \|\| i > j \|\| data.Offset()+int(i) > data.Offset()+data.Len() {
	panic("arrow/array: index out of range")
	}

	for _, b := range data.Buffers() {
	if b != nil {
	b.Retain()
	}
	}

	for _, child := range data.Children() {
	if child != nil {
	child.Retain()
	}
	}

	if data.(*Data).dictionary != nil {
	data.(*Data).dictionary.Retain()
	}

	o := &Data{
	dtype: data.DataType(),
	nulls: UnknownNullCount,
	length: int(j - i),
	offset: data.Offset() + int(i),
	buffers: data.Buffers(),
	childData: data.Children(),
	dictionary: data.(*Data).dictionary,
	}
	o.refCount.Add(1)

	if data.NullN() == 0 {
	o.nulls = 0
	}

	return o
	}

	func Hash(h *maphash.MapHash, data arrow.ArrayData) {
	a := data.(*Data)

	hashString(h, a.dtype.Fingerprint())
	hashInt(h, a.length)
	hashInt(h, a.offset)

	hashInt(h, len(a.buffers))
	for _, b := range a.buffers {
	if b == nil {
	hashByte(h, 0)
	continue
	}

	hashByte(h, 1)
	hashBytes(h, b.Bytes())
	}

	hashInt(h, len(a.childData))
	for _, c := range a.childData {
	if c == nil {
	hashByte(h, 0)
	continue
	}

	hashByte(h, 1)
	Hash(h, c)
	}

	if a.dictionary == nil {
	hashByte(h, 0)
	return
	}

	hashByte(h, 1)
	Hash(h, a.dictionary)
	}

	func hashByte(h *maphash.MapHash, v byte) {
	var b [1]byte
	b[0] = v
	h.Write(b[:])
	}

	func hashInt(h *maphash.MapHash, v int) {
	var b [8]byte
	binary.LittleEndian.PutUint64(b[:], uint64(int64(v)))
	h.Write(b[:])
	}

	func hashBytes(h *maphash.MapHash, b []byte) {
	hashInt(h, len(b))
	h.Write(b)
	}

	func hashString(h *maphash.MapHash, s string) {
	hashInt(h, len(s))
	h.WriteString(s)
	}