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 f5b0c6ba feat(compute): add dictionary_encode (#1140)
f5b0c6ba is described below
commit f5b0c6ba5ac20b744cb53bedc6e1c4374e412612
Author: Minh Vu <[email protected]>
AuthorDate: Mon Aug 17 22:23:08 2026 +0200
feat(compute): add dictionary_encode (#1140)
### Rationale for this change
Arrow-Go currently has a vector hash kernel for unique, but not for
dictionary encoding. The existing memo table already assigns stable
positions to distinct values, so dictionary encoding can build on that
state and provide a useful output form for repeated values.
### What changes are included in this PR?
- Add the dictionary_encode compute function and
DictionaryEncodeOptions.
- Support primitive, binary, temporal, decimal, interval, fixed-size
binary, and boolean input types covered by the existing vector hash
kernels.
- Use int32 dictionary indexes and preserve first-seen dictionary order.
- Mask nulls by default. With NullEncodingEncode, null is included in
the dictionary and encoded as a regular index.
- Keep one shared dictionary across chunked input while returning one
index per input value.
- Pass through already dictionary-encoded input instead of creating a
nested dictionary type.
- Register boolean inputs in the shared vector hash kernels, which also
adds `unique(bool)` support while preserving its existing
null-as-a-distinct-value behavior.
- Correct the `unique` documentation to state that null is treated as a
distinct value.
### Are these changes tested?
- go test ./arrow/compute/... -count=1
- Added coverage for repeated and all-null values in both null modes,
chunked input, dictionary input, dictionary order, output cardinality,
boolean unique and dictionary encoding, sliced boolean input, and
memo-table resizing.
### Are there any user-facing changes?
Yes. This adds the public compute.DictionaryEncode function,
DictionaryEncodeOptions, and the NullEncodingMask and NullEncodingEncode
options. It also adds boolean input support to unique through the shared
vector hash registration.
---
arrow/array/util.go | 15 +-
arrow/compute/executor.go | 29 +-
arrow/compute/executor_test.go | 116 ++++
arrow/compute/expression.go | 9 +-
arrow/compute/expression_test.go | 53 ++
arrow/compute/internal/kernels/helpers.go | 10 +-
arrow/compute/internal/kernels/vector_hash.go | 378 ++++++++++-
arrow/compute/internal/kernels/vector_hash_test.go | 78 +++
arrow/compute/vector_hash.go | 70 +-
arrow/compute/vector_hash_test.go | 749 +++++++++++++++++++++
10 files changed, 1482 insertions(+), 25 deletions(-)
diff --git a/arrow/array/util.go b/arrow/array/util.go
index 1b303bf6..afd0f39b 100644
--- a/arrow/array/util.go
+++ b/arrow/array/util.go
@@ -327,9 +327,18 @@ func GetDictArrayData(mem memory.Allocator, valueType
arrow.DataType, memoTable
switch tbl := memoTable.(type) {
case hashing.NumericMemoTable:
- nbytes := tbl.TypeTraits().BytesRequired(dictLen)
- buffers[1].Resize(nbytes)
- tbl.WriteOutSubset(startOffset, buffers[1].Bytes())
+ if valueType.ID() == arrow.BOOL {
+ values := make([]uint8, dictLen)
+ tbl.CopyValuesSubset(startOffset, values)
+
buffers[1].Resize(int(bitutil.BytesForBits(int64(dictLen))))
+ for i, value := range values {
+ bitutil.SetBitTo(buffers[1].Bytes(), i, value
!= 0)
+ }
+ } else {
+ nbytes := tbl.TypeTraits().BytesRequired(dictLen)
+ buffers[1].Resize(nbytes)
+ tbl.WriteOutSubset(startOffset, buffers[1].Bytes())
+ }
case *hashing.BinaryMemoTable:
switch valueType.ID() {
case arrow.BINARY, arrow.STRING:
diff --git a/arrow/compute/executor.go b/arrow/compute/executor.go
index 095a7e15..1728e49e 100644
--- a/arrow/compute/executor.go
+++ b/arrow/compute/executor.go
@@ -1002,18 +1002,30 @@ func (v *vectorExecutor) WrapResults(ctx
context.Context, out <-chan Datum, hasC
output Datum
acc []arrow.Array
)
+ releaseAccumulated := func() {
+ for _, c := range acc {
+ c.Release()
+ }
+ acc = nil
+ }
toChunked := func() {
out := output.(ArrayLikeDatum).Chunks()
acc = make([]arrow.Array, 0, len(out))
+ isChunked := output.Kind() == KindChunked
for _, o := range out {
if o.Len() > 0 {
+ if isChunked {
+ // ChunkedDatum.Chunks returns borrowed
references.
+ o.Retain()
+ }
acc = append(acc, o)
+ } else if !isChunked {
+ // ArrayDatum.Chunks creates an owned array.
+ o.Release()
}
}
- if output.Kind() != KindChunked {
- output.Release()
- }
+ output.Release()
output = nil
}
@@ -1023,6 +1035,9 @@ func (v *vectorExecutor) WrapResults(ctx context.Context,
out <-chan Datum, hasC
return nil
case output = <-out:
if output == nil || ctx.Err() != nil {
+ if output != nil {
+ output.Release()
+ }
return nil
}
@@ -1038,6 +1053,10 @@ func (v *vectorExecutor) WrapResults(ctx
context.Context, out <-chan Datum, hasC
case <-ctx.Done():
// context is done, either cancelled or a timeout.
// either way, we end early and return what we've got
so far.
+ if output == nil {
+ releaseAccumulated()
+ return nil
+ }
return output
case o, ok := <-out:
if !ok { // channel closed, wrap it up
@@ -1045,9 +1064,7 @@ func (v *vectorExecutor) WrapResults(ctx context.Context,
out <-chan Datum, hasC
return output
}
- for _, c := range acc {
- defer c.Release()
- }
+ defer releaseAccumulated()
chkd := arrow.NewChunked(v.outType, acc)
defer chkd.Release()
diff --git a/arrow/compute/executor_test.go b/arrow/compute/executor_test.go
new file mode 100644
index 00000000..a8194108
--- /dev/null
+++ b/arrow/compute/executor_test.go
@@ -0,0 +1,116 @@
+// 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.
+
+//go:build go1.18
+
+package compute
+
+import (
+ "context"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/compute/exec"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/stretchr/testify/require"
+)
+
+type signalChunkedDatum struct {
+ *ChunkedDatum
+ chunksCalled chan struct{}
+}
+
+func (d *signalChunkedDatum) Chunks() []arrow.Array {
+ close(d.chunksCalled)
+ return d.Value.Chunks()
+}
+
+func TestVectorExecutorWrapResultsReleasesEmptyChunkedOutput(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ builder := array.NewInt32Builder(mem)
+ builder.Append(42)
+ value := builder.NewInt32Array()
+ builder.Release()
+ defer value.Release()
+
+ empty := array.NewSlice(value, 0, 0)
+ defer empty.Release()
+ nonEmpty := array.NewSlice(value, 0, 1)
+ defer nonEmpty.Release()
+
+ chunked := arrow.NewChunked(value.DataType(), []arrow.Array{empty,
nonEmpty})
+ output := make(chan Datum, 1)
+ output <- &ChunkedDatum{Value: chunked}
+ close(output)
+
+ executor := &vectorExecutor{
+ nonAggExecImpl: nonAggExecImpl{
+ kernel: &exec.VectorKernel{OutputChunked: true},
+ outType: value.DataType(),
+ },
+ }
+
+ result := executor.WrapResults(context.Background(), output, true)
+ require.NotNil(t, result)
+ require.Equal(t, KindChunked, result.Kind())
+
+ resultChunked := result.(*ChunkedDatum).Value
+ require.Len(t, resultChunked.Chunks(), 1)
+ require.Equal(t, int32(42),
resultChunked.Chunk(0).(*array.Int32).Value(0))
+ result.Release()
+}
+
+func TestVectorExecutorWrapResultsReleasesChunkedOutputOnCancellation(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ builder := array.NewInt32Builder(mem)
+ builder.Append(42)
+ value := builder.NewInt32Array()
+ builder.Release()
+ defer value.Release()
+
+ chunked := arrow.NewChunked(value.DataType(), []arrow.Array{value})
+ output := make(chan Datum)
+ ctx, cancel := context.WithCancel(context.Background())
+ defer cancel()
+
+ executor := &vectorExecutor{
+ nonAggExecImpl: nonAggExecImpl{
+ kernel: &exec.VectorKernel{OutputChunked: true},
+ outType: value.DataType(),
+ },
+ }
+
+ datum := &signalChunkedDatum{
+ ChunkedDatum: &ChunkedDatum{Value: chunked},
+ chunksCalled: make(chan struct{}),
+ }
+ result := make(chan Datum, 1)
+ go func() {
+ result <- executor.WrapResults(ctx, output, true)
+ }()
+
+ output <- datum
+ <-datum.chunksCalled
+ cancel()
+
+ require.Nil(t, <-result)
+ close(output)
+}
diff --git a/arrow/compute/expression.go b/arrow/compute/expression.go
index dcf572e7..a1927366 100644
--- a/arrow/compute/expression.go
+++ b/arrow/compute/expression.go
@@ -578,6 +578,7 @@ var (
funcOptsTypes = []FunctionOptions{
SetLookupOptions{}, ArithmeticOptions{}, CastOptions{},
FilterOptions{}, NullOptions{}, StrptimeOptions{},
MakeStructOptions{},
+ DictionaryEncodeOptions{},
CumulativeOptions{},
}
)
@@ -954,7 +955,13 @@ func DeserializeExpr(mem memory.Allocator, buf
*memory.Buffer) (Expression, erro
return nil,
errors.New("options scalar typename must be binary")
}
- optionsVal :=
reflect.New(funcOptionsMap[string(typname.(*scalar.Binary).Data())]).Interface()
+ typeName :=
string(typname.(*scalar.Binary).Data())
+ optionsType, ok :=
funcOptionsMap[typeName]
+ if !ok {
+ return nil,
fmt.Errorf("%w: unknown function options type %q", arrow.ErrInvalid, typeName)
+ }
+
+ optionsVal :=
reflect.New(optionsType).Interface()
if err :=
scalar.FromScalarWithAllocator(optsScalar.(*scalar.Struct), optionsVal, mem);
err != nil {
return nil, err
}
diff --git a/arrow/compute/expression_test.go b/arrow/compute/expression_test.go
index 37c58ad6..a28c795a 100644
--- a/arrow/compute/expression_test.go
+++ b/arrow/compute/expression_test.go
@@ -20,6 +20,7 @@
package compute_test
import (
+ "fmt"
"testing"
"github.com/apache/arrow-go/v18/arrow"
@@ -28,8 +29,13 @@ import (
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/arrow/scalar"
"github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
)
+type unknownFunctionOptions struct{}
+
+func (unknownFunctionOptions) TypeName() string { return
"UnknownFunctionOptions" }
+
type privateFunctionOptions struct {
value int
}
@@ -414,3 +420,50 @@ func TestExpressionSerializationRoundTrip(t *testing.T) {
})
}
}
+
+func TestDictionaryEncodeOptionsSerializationRoundTrip(t *testing.T) {
+ for _, behavior := range []compute.NullEncodingBehavior{
+ compute.NullEncodingMask,
+ compute.NullEncodingEncode,
+ } {
+ t.Run(fmt.Sprintf("null encoding %d", behavior), func(t
*testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ expr := compute.NewCall(
+ "dictionary_encode",
+
[]compute.Expression{compute.NewFieldRef("values")},
+ &compute.DictionaryEncodeOptions{NullEncoding:
behavior},
+ )
+ defer expr.Release()
+
+ serialized, err := compute.SerializeExpr(expr, mem)
+ require.NoError(t, err)
+ defer serialized.Release()
+
+ roundTripped, err := compute.DeserializeExpr(mem,
serialized)
+ require.NoError(t, err)
+ defer roundTripped.Release()
+ require.True(t, expr.Equals(roundTripped))
+ })
+ }
+}
+
+func TestDeserializeExprRejectsUnknownOptions(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ expr := compute.NewCall(
+ "dictionary_encode",
+ []compute.Expression{compute.NewFieldRef("values")},
+ unknownFunctionOptions{},
+ )
+ defer expr.Release()
+
+ serialized, err := compute.SerializeExpr(expr, mem)
+ require.NoError(t, err)
+ defer serialized.Release()
+
+ _, err = compute.DeserializeExpr(mem, serialized)
+ assert.ErrorIs(t, err, arrow.ErrInvalid)
+}
diff --git a/arrow/compute/internal/kernels/helpers.go
b/arrow/compute/internal/kernels/helpers.go
index ef5f0bb4..5daa31a1 100644
--- a/arrow/compute/internal/kernels/helpers.go
+++ b/arrow/compute/internal/kernels/helpers.go
@@ -851,6 +851,10 @@ func (v *validityBuilder) Finish() (buf *memory.Buffer) {
return
}
+func (v *validityBuilder) reset() {
+ v.bitLength, v.falseCount = 0, 0
+}
+
type execBufBuilder struct {
mem memory.Allocator
buffer *memory.Buffer
@@ -876,6 +880,10 @@ func (bldr *execBufBuilder) unsafeAppend(data []byte) {
bldr.sz += len(data)
}
+func (bldr *execBufBuilder) reset() {
+ bldr.sz = 0
+}
+
func (bldr *execBufBuilder) finish() (buf *memory.Buffer) {
if bldr.buffer == nil {
buf = memory.NewBufferBytes(nil)
@@ -883,7 +891,7 @@ func (bldr *execBufBuilder) finish() (buf *memory.Buffer) {
}
bldr.buffer.Resize(bldr.sz)
buf = bldr.buffer
- bldr.buffer, bldr.sz = nil, 0
+ bldr.buffer, bldr.data, bldr.sz = nil, nil, 0
return
}
diff --git a/arrow/compute/internal/kernels/vector_hash.go
b/arrow/compute/internal/kernels/vector_hash.go
index ca6184eb..00f48173 100644
--- a/arrow/compute/internal/kernels/vector_hash.go
+++ b/arrow/compute/internal/kernels/vector_hash.go
@@ -23,9 +23,11 @@ import (
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/bitutil"
"github.com/apache/arrow-go/v18/arrow/compute/exec"
"github.com/apache/arrow-go/v18/arrow/internal/debug"
"github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/arrow/scalar"
"github.com/apache/arrow-go/v18/internal/bitutils"
"github.com/apache/arrow-go/v18/internal/hashing"
)
@@ -78,11 +80,172 @@ func (emptyAction) ShouldEncodeNulls() bool {
return true }
type uniqueAction = emptyAction
+type NullEncodingBehavior int8
+
+const (
+ // NullEncodingMask keeps null input values null in the indices array.
+ // It is the zero value and default behavior.
+ NullEncodingMask NullEncodingBehavior = iota
+ // NullEncodingEncode adds null input values to the dictionary as a
regular entry.
+ NullEncodingEncode
+)
+
+type DictionaryEncodeOptions struct {
+ // NullEncoding controls how null input values are represented.
+ NullEncoding NullEncodingBehavior `compute:"null_encoding_behavior"`
+}
+
+func (DictionaryEncodeOptions) TypeName() string { return
"DictionaryEncodeOptions" }
+
+func (opts DictionaryEncodeOptions) ToScalar() (scalar.Scalar, error) {
+ var encoded uint32
+ switch opts.NullEncoding {
+ case NullEncodingEncode:
+ encoded = 0
+ case NullEncodingMask:
+ encoded = 1
+ default:
+ return nil, fmt.Errorf("%w: invalid null encoding behavior %d",
arrow.ErrInvalid, opts.NullEncoding)
+ }
+
+ return scalar.NewStructScalarWithNames(
+ []scalar.Scalar{
+ scalar.NewUint32Scalar(encoded),
+
scalar.NewBinaryScalar(memory.NewBufferBytes([]byte(opts.TypeName())),
arrow.BinaryTypes.Binary),
+ },
+ []string{"null_encoding_behavior", "_type_name"},
+ )
+}
+
+func (opts *DictionaryEncodeOptions) FromStructScalar(sc *scalar.Struct) error
{
+ value, err := sc.Field("null_encoding_behavior")
+ if err != nil {
+ return err
+ }
+
+ encoded, ok := value.(*scalar.Uint32)
+ if !ok || !encoded.IsValid() {
+ return fmt.Errorf("%w: null_encoding_behavior must be a valid
uint32 scalar", arrow.ErrInvalid)
+ }
+
+ var behavior NullEncodingBehavior
+ switch encoded.Value {
+ case 0:
+ behavior = NullEncodingEncode
+ case 1:
+ behavior = NullEncodingMask
+ default:
+ return fmt.Errorf("%w: invalid null encoding behavior %d",
arrow.ErrInvalid, encoded.Value)
+ }
+
+ opts.NullEncoding = behavior
+ return nil
+}
+
+type dictionaryEncodeAction struct {
+ nullEncoding NullEncodingBehavior
+ indices *bufferBuilder[int32]
+ validity validityBuilder
+ length int
+ nulls int
+ err error
+}
+
+func (a *dictionaryEncodeAction) Reset() error {
+ a.indices.reset()
+ a.validity.reset()
+ a.length = 0
+ a.nulls = 0
+ a.err = nil
+ return nil
+}
+
+func (a *dictionaryEncodeAction) Reserve(n int) error {
+ a.indices.reserve(n)
+ a.validity.Reserve(int64(n))
+ return nil
+}
+
+func (a *dictionaryEncodeAction) appendIndex(idx int, valid bool) {
+ if !valid {
+ idx = 0
+ }
+ if idx < 0 || int64(idx) > int64(1<<31-1) {
+ if a.err == nil {
+ a.err = fmt.Errorf("%w: dictionary index %d does not
fit in int32", arrow.ErrInvalid, idx)
+ }
+ return
+ }
+
+ a.indices.unsafeAppend(int32(idx))
+ a.validity.UnsafeAppend(valid)
+ a.length++
+ if !valid {
+ a.nulls++
+ }
+}
+
+func (a *dictionaryEncodeAction) Flush(out *exec.ExecResult) error {
+ if a.err != nil {
+ return a.err
+ }
+
+ out.Len = int64(a.length)
+ out.Nulls = int64(a.nulls)
+ if a.length != 0 {
+ out.Buffers[1].WrapBuffer(a.indices.finish())
+ } else {
+ a.indices.finish().Release()
+ }
+
+ validity := a.validity.Finish()
+ if a.nulls != 0 {
+ out.Buffers[0].WrapBuffer(validity)
+ } else if validity != nil {
+ validity.Release()
+ }
+
+ a.length = 0
+ a.nulls = 0
+ return nil
+}
+
+func (a *dictionaryEncodeAction) FlushFinal(*exec.ExecResult) error {
+ return nil
+}
+
+func (a *dictionaryEncodeAction) ObserveFound(idx int) {
+ a.appendIndex(idx, true)
+}
+
+func (a *dictionaryEncodeAction) ObserveNotFound(idx int) error {
+ a.appendIndex(idx, true)
+ return a.err
+}
+
+func (a *dictionaryEncodeAction) ObserveNullFound(idx int) {
+ if a.nullEncoding == NullEncodingMask {
+ a.appendIndex(0, false)
+ } else {
+ a.appendIndex(idx, true)
+ }
+}
+
+func (a *dictionaryEncodeAction) ObserveNullNotFound(idx int) error {
+ a.ObserveNullFound(idx)
+ return a.err
+}
+
+func (a *dictionaryEncodeAction) ShouldEncodeNulls() bool {
+ return a.nullEncoding == NullEncodingEncode
+}
+
type regularHashState struct {
- mem memory.Allocator
- typ arrow.DataType
- memoTable hashing.MemoTable
- action Action
+ mem memory.Allocator
+ typ arrow.DataType
+ memoTable hashing.MemoTable
+ action Action
+ memoReleased bool
doAppend func(Action, hashing.MemoTable, *exec.ArraySpan) error
}
@@ -92,7 +255,16 @@ func (rhs *regularHashState) Allocator() memory.Allocator {
return rhs.mem }
func (rhs *regularHashState) ValueType() arrow.DataType { return rhs.typ }
func (rhs *regularHashState) Reset() error {
- rhs.memoTable.Reset()
+ if rhs.memoReleased {
+ memoTable, err := newMemoTable(rhs.mem, rhs.typ.ID())
+ if err != nil {
+ return err
+ }
+ rhs.memoTable = memoTable
+ rhs.memoReleased = false
+ } else {
+ rhs.memoTable.Reset()
+ }
return rhs.action.Reset()
}
@@ -114,6 +286,10 @@ func (rhs *regularHashState) GetDictionary()
(arrow.ArrayData, error) {
}
func doAppendBinary[OffsetT int32 | int64](action Action, memo
hashing.MemoTable, arr *exec.ArraySpan) error {
+ if arr.Len == 0 {
+ return nil
+ }
+
var (
bitmap = arr.Buffers[0].Buf
offsets = exec.GetSpanOffsets[OffsetT](arr, 1)
@@ -142,6 +318,7 @@ func doAppendBinary[OffsetT int32 | int64](action Action,
memo hashing.MemoTable
idx, found := memo.GetOrInsertNull()
if found {
action.ObserveNullFound(idx)
+ return nil
}
return action.ObserveNullNotFound(idx)
})
@@ -173,6 +350,7 @@ func doAppendFixedSize(action Action, memo
hashing.MemoTable, arr *exec.ArraySpa
idx, found := memo.GetOrInsertNull()
if found {
action.ObserveNullFound(idx)
+ return nil
}
return action.ObserveNullNotFound(idx)
})
@@ -200,6 +378,43 @@ func doAppendNumeric[T arrow.IntType | arrow.UintType |
arrow.FloatType](action
idx, found := memo.GetOrInsertNull()
if found {
action.ObserveNullFound(idx)
+ return nil
+ }
+ return action.ObserveNullNotFound(idx)
+ })
+}
+
+func doAppendBoolean(action Action, memo hashing.MemoTable, arr
*exec.ArraySpan) error {
+ if arr.Len == 0 {
+ return nil
+ }
+
+ values := arr.Buffers[1].Buf
+ shouldEncodeNulls := action.ShouldEncodeNulls()
+ return bitutils.VisitBitBlocksShort(arr.Buffers[0].Buf, arr.Offset,
arr.Len,
+ func(pos int64) error {
+ value := uint8(0)
+ if bitutil.BitIsSet(values, int(arr.Offset+pos)) {
+ value = 1
+ }
+ idx, found, err := memo.GetOrInsert(value)
+ if err != nil {
+ return err
+ }
+ if found {
+ action.ObserveFound(idx)
+ return nil
+ }
+ return action.ObserveNotFound(idx)
+ }, func() error {
+ if !shouldEncodeNulls {
+ return action.ObserveNullNotFound(-1)
+ }
+
+ idx, found := memo.GetOrInsertNull()
+ if found {
+ action.ObserveNullFound(idx)
+ return nil
}
return action.ObserveNullNotFound(idx)
})
@@ -217,6 +432,7 @@ func (nhs *nullHashState) Allocator() memory.Allocator {
return nhs.mem }
func (nhs *nullHashState) ValueType() arrow.DataType { return nhs.typ }
func (nhs *nullHashState) Reset() error {
+ nhs.seenNull = false
return nhs.action.Reset()
}
@@ -254,6 +470,39 @@ func (nhs *nullHashState) GetDictionary()
(arrow.ArrayData, error) {
return data, nil
}
+func dictionaryEncodeIdentity(_ *exec.KernelCtx, batch *exec.ExecSpan, out
*exec.ExecResult) error {
+ data := batch.Values[0].Array.MakeData()
+ defer data.Release()
+ out.TakeOwnership(data)
+ return nil
+}
+
+func dictionaryEncodeIdentityChunked(_ *exec.KernelCtx, batch
[]*arrow.Chunked, _ *exec.ExecResult) ([]*exec.ExecResult, error) {
+ if len(batch) != 1 {
+ return nil, fmt.Errorf("%w: dictionary_encode expects one
input", arrow.ErrInvalid)
+ }
+
+ chunks := batch[0].Chunks()
+ if len(chunks) == 0 {
+ result := &exec.ExecResult{}
+ exec.FillZeroLength(batch[0].DataType(), result)
+ return []*exec.ExecResult{result}, nil
+ }
+
+ results := make([]*exec.ExecResult, 0, len(chunks))
+ for _, chunk := range chunks {
+ result := &exec.ExecResult{}
+ result.TakeOwnership(chunk.Data())
+ results = append(results, result)
+ }
+ return results, nil
+}
+
+func initDictionaryEncodeIdentity(_ *exec.KernelCtx, args exec.KernelInitArgs)
(exec.KernelState, error) {
+ _, err := parseDictionaryEncodeOptions(args.Options)
+ return nil, err
+}
+
type dictionaryHashState struct {
indicesKernel HashState
dictionary arrow.Array
@@ -329,10 +578,14 @@ func (dhs *dictionaryHashState) Append(ctx
*exec.KernelCtx, arr *exec.ArraySpan)
func nullHashInit(actionInit initAction) exec.KernelInitFn {
return func(ctx *exec.KernelCtx, args exec.KernelInitArgs)
(exec.KernelState, error) {
mem := exec.GetAllocator(ctx.Ctx)
+ action, err := actionInit(args.Inputs[0], args.Options, mem)
+ if err != nil {
+ return nil, err
+ }
ret := &nullHashState{
mem: mem,
typ: args.Inputs[0],
- action: actionInit(args.Inputs[0], args.Options, mem),
+ action: action,
}
ret.Reset()
return ret, nil
@@ -341,7 +594,7 @@ func nullHashInit(actionInit initAction) exec.KernelInitFn {
func newMemoTable(mem memory.Allocator, dt arrow.Type) (hashing.MemoTable,
error) {
switch dt {
- case arrow.INT8, arrow.UINT8:
+ case arrow.BOOL, arrow.INT8, arrow.UINT8:
return hashing.NewMemoTable[uint8](0), nil
case arrow.INT16, arrow.UINT16:
return hashing.NewMemoTable[uint16](0), nil
@@ -367,6 +620,10 @@ func newMemoTable(mem memory.Allocator, dt arrow.Type)
(hashing.MemoTable, error
func regularHashInit(dt arrow.DataType, actionInit initAction, appendFn
func(Action, hashing.MemoTable, *exec.ArraySpan) error) exec.KernelInitFn {
return func(ctx *exec.KernelCtx, args exec.KernelInitArgs)
(exec.KernelState, error) {
mem := exec.GetAllocator(ctx.Ctx)
+ action, err := actionInit(args.Inputs[0], args.Options, mem)
+ if err != nil {
+ return nil, err
+ }
memoTable, err := newMemoTable(mem, dt.ID())
if err != nil {
return nil, err
@@ -376,7 +633,7 @@ func regularHashInit(dt arrow.DataType, actionInit
initAction, appendFn func(Act
mem: mem,
typ: args.Inputs[0],
memoTable: memoTable,
- action: actionInit(args.Inputs[0], args.Options,
mem),
+ action: action,
doAppend: appendFn,
}
ret.Reset()
@@ -415,12 +672,14 @@ func dictionaryHashInit(actionInit initAction)
exec.KernelInitFn {
}
}
-type initAction func(arrow.DataType, any, memory.Allocator) Action
+type initAction func(arrow.DataType, any, memory.Allocator) (Action, error)
func getHashInit(typeID arrow.Type, actionInit initAction) exec.KernelInitFn {
switch typeID {
case arrow.NULL:
return nullHashInit(actionInit)
+ case arrow.BOOL:
+ return regularHashInit(arrow.FixedWidthTypes.Boolean,
actionInit, doAppendBoolean)
case arrow.INT8, arrow.UINT8:
return regularHashInit(arrow.PrimitiveTypes.Uint8, actionInit,
doAppendNumeric[uint8])
case arrow.INT16, arrow.UINT16:
@@ -514,6 +773,13 @@ func uniqueFinalizeDictionary(ctx *exec.KernelCtx, result
[]*exec.ArraySpan) (ou
func addHashKernels(base exec.VectorKernel, actionInit initAction, outTy
exec.OutputType) []exec.VectorKernel {
kernels := make([]exec.VectorKernel, 0)
+ base.Init = getHashInit(arrow.BOOL, actionInit)
+ base.Signature = &exec.KernelSignature{
+ InputTypes:
[]exec.InputType{exec.NewExactInput(arrow.FixedWidthTypes.Boolean)},
+ OutType: outTy,
+ }
+ kernels = append(kernels, base)
+
for _, ty := range primitiveTypes {
base.Init = getHashInit(ty.ID(), actionInit)
base.Signature = &exec.KernelSignature{
@@ -538,8 +804,81 @@ func addHashKernels(base exec.VectorKernel, actionInit
initAction, outTy exec.Ou
return kernels
}
-func initUnique(dt arrow.DataType, _ any, mem memory.Allocator) Action {
- return uniqueAction{mem: mem, dt: dt}
+func initUnique(dt arrow.DataType, _ any, mem memory.Allocator) (Action,
error) {
+ return uniqueAction{mem: mem, dt: dt}, nil
+}
+
+func parseDictionaryEncodeOptions(options any) (DictionaryEncodeOptions,
error) {
+ opts := DictionaryEncodeOptions{}
+ switch v := options.(type) {
+ case nil:
+ case DictionaryEncodeOptions:
+ opts = v
+ case *DictionaryEncodeOptions:
+ if v != nil {
+ opts = *v
+ }
+ default:
+ return opts, fmt.Errorf("%w: expected DictionaryEncodeOptions,
got %T", arrow.ErrInvalid, options)
+ }
+
+ if opts.NullEncoding != NullEncodingMask && opts.NullEncoding !=
NullEncodingEncode {
+ return opts, fmt.Errorf("%w: invalid null encoding behavior
%d", arrow.ErrInvalid, opts.NullEncoding)
+ }
+ return opts, nil
+}
+
+func initDictionaryEncode(_ arrow.DataType, options any, mem memory.Allocator)
(Action, error) {
+ opts, err := parseDictionaryEncodeOptions(options)
+ if err != nil {
+ return nil, err
+ }
+
+ return &dictionaryEncodeAction{
+ nullEncoding: opts.NullEncoding,
+ indices: newBufferBuilder[int32](mem),
+ validity: validityBuilder{mem: mem},
+ }, nil
+}
+
+var outputDictionaryType = exec.NewComputedOutputType(func(_ *exec.KernelCtx,
args []arrow.DataType) (arrow.DataType, error) {
+ if len(args) != 1 {
+ return nil, fmt.Errorf("%w: dictionary_encode expects one input
type", arrow.ErrInvalid)
+ }
+ return &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int32,
+ ValueType: args[0],
+ }, nil
+})
+
+func dictionaryEncodeFinalize(ctx *exec.KernelCtx, results []*exec.ArraySpan)
([]*exec.ArraySpan, error) {
+ impl, ok := ctx.State.(HashState)
+ if !ok {
+ return nil, fmt.Errorf("%w: HashState in invalid state",
arrow.ErrInvalid)
+ }
+ defer releaseHashMemo(impl)
+
+ dict, err := impl.GetDictionary()
+ if err != nil {
+ return nil, err
+ }
+ defer dict.Release()
+
+ for _, result := range results {
+ var dictSpan exec.ArraySpan
+ dictSpan.TakeOwnership(dict)
+ result.SetDictionary(&dictSpan)
+ }
+ return results, nil
+}
+
+func releaseHashMemo(hash HashState) {
+ if state, ok := hash.(*regularHashState); ok {
+ if memo, ok := state.memoTable.(*hashing.BinaryMemoTable); ok
&& !state.memoReleased {
+ memo.Release()
+ state.memoReleased = true
+ }
+ }
}
func GetVectorHashKernels() (unique, valueCounts, dictEncode
[]exec.VectorKernel) {
@@ -561,5 +900,22 @@ func GetVectorHashKernels() (unique, valueCounts,
dictEncode []exec.VectorKernel
}
unique = append(unique, base)
+ // dictionary encode
+ base.Finalize = dictionaryEncodeFinalize
+ base.OutputChunked = true
+ base.NullHandling = exec.NullComputedNoPrealloc
+ base.MemAlloc = exec.MemNoPrealloc
+ dictEncode = addHashKernels(base, initDictionaryEncode,
outputDictionaryType)
+ identity := exec.NewVectorKernelWithSig(
+ &exec.KernelSignature{
+ InputTypes:
[]exec.InputType{exec.NewIDInput(arrow.DICTIONARY)},
+ OutType: OutputFirstType,
+ },
+ dictionaryEncodeIdentity,
+ initDictionaryEncodeIdentity)
+ identity.CanExecuteChunkWise = false
+ identity.ExecChunked = dictionaryEncodeIdentityChunked
+ dictEncode = append(dictEncode, identity)
+
return
}
diff --git a/arrow/compute/internal/kernels/vector_hash_test.go
b/arrow/compute/internal/kernels/vector_hash_test.go
new file mode 100644
index 00000000..9f0549d4
--- /dev/null
+++ b/arrow/compute/internal/kernels/vector_hash_test.go
@@ -0,0 +1,78 @@
+// 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.
+
+//go:build go1.18
+
+package kernels
+
+import (
+ "context"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/compute/exec"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/stretchr/testify/require"
+)
+
+func TestDictionaryEncodeStateResetAfterFinalize(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ ctx := &exec.KernelCtx{Ctx: exec.WithAllocator(context.Background(),
mem)}
+ inputBuilder := array.NewStringBuilder(mem)
+ inputBuilder.AppendValues([]string{"foo", "bar", "foo"}, nil)
+ input := inputBuilder.NewStringArray()
+ inputBuilder.Release()
+ defer input.Release()
+
+ state, err := getHashInit(arrow.STRING, initDictionaryEncode)(ctx,
exec.KernelInitArgs{
+ Inputs: []arrow.DataType{arrow.BinaryTypes.String},
+ Options: DictionaryEncodeOptions{},
+ })
+ require.NoError(t, err)
+ ctx.State = state
+ hash := state.(HashState)
+
+ inputSpan := &exec.ArraySpan{}
+ inputSpan.SetMembers(input.Data())
+ outputType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int32,
+ ValueType: input.DataType(),
+ }
+
+ encodeOnce := func() {
+ require.NoError(t, hash.Append(ctx, inputSpan))
+ result := &exec.ArraySpan{Type: outputType}
+ require.NoError(t, hash.Flush(result))
+
+ results, err := dictionaryEncodeFinalize(ctx,
[]*exec.ArraySpan{result})
+ require.NoError(t, err)
+ for _, result := range results {
+ result.Release()
+ }
+ }
+
+ encodeOnce()
+ stateImpl := state.(*regularHashState)
+ require.True(t, stateImpl.memoReleased)
+
+ require.NoError(t, hash.Reset())
+ require.False(t, stateImpl.memoReleased)
+
+ encodeOnce()
+}
diff --git a/arrow/compute/vector_hash.go b/arrow/compute/vector_hash.go
index facdd1ff..017bf4de 100644
--- a/arrow/compute/vector_hash.go
+++ b/arrow/compute/vector_hash.go
@@ -20,19 +20,44 @@ package compute
import (
"context"
+ "fmt"
"github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
"github.com/apache/arrow-go/v18/arrow/compute/internal/kernels"
)
var (
uniqueDoc = FunctionDoc{
- Summary: "Compute unique elements",
- Description: "Return an array with distinct values. Nulls in
the input are ignored",
+ Summary: "Compute unique elements",
+ Description: "Return an array with distinct values.\n" +
+ "Nulls in the input are considered a distinct value",
+ ArgNames: []string{"array"},
+ }
+ dictionaryEncodeDoc = FunctionDoc{
+ Summary: "Dictionary encode an array",
+ Description: "Return a dictionary-encoded array with the
distinct values in the dictionary.\n" +
+ "If the input is already dictionary encoded, it is
returned unchanged,\n" +
+ "including its index type and null representation.\n" +
+ "Newly encoded arrays use int32 dictionary indices.",
ArgNames: []string{"array"},
+ OptionsType: "DictionaryEncodeOptions",
}
)
+// NullEncodingBehavior controls how null input values are represented.
+type NullEncodingBehavior = kernels.NullEncodingBehavior
+
+const (
+ // NullEncodingMask keeps null input values null in the indices array.
+ NullEncodingMask = kernels.NullEncodingMask
+ // NullEncodingEncode adds null input values to the dictionary as a
regular entry.
+ NullEncodingEncode = kernels.NullEncodingEncode
+)
+
+// DictionaryEncodeOptions controls dictionary encoding behavior.
+type DictionaryEncodeOptions = kernels.DictionaryEncodeOptions
+
func Unique(ctx context.Context, values Datum) (Datum, error) {
return CallFunction(ctx, "unique", nil, values)
}
@@ -47,8 +72,38 @@ func UniqueArray(ctx context.Context, values arrow.Array)
(arrow.Array, error) {
return out.(*ArrayDatum).MakeArray(), nil
}
+// DictionaryEncode returns a dictionary-encoded version of values.
+// Newly encoded arrays use int32 indices. For newly encoded arrays, nulls are
+// masked unless NullEncodingEncode is selected. Existing dictionary arrays are
+// returned unchanged, including their index type and null representation.
+func DictionaryEncode(ctx context.Context, opts DictionaryEncodeOptions,
values Datum) (Datum, error) {
+ return CallFunction(ctx, "dictionary_encode", &opts, values)
+}
+
+// DictionaryEncodeArray returns a dictionary-encoded version of values.
+func DictionaryEncodeArray(ctx context.Context, opts DictionaryEncodeOptions,
values arrow.Array) (arrow.Array, error) {
+ datum, err := DictionaryEncode(ctx, opts, &ArrayDatum{Value:
values.Data()})
+ if err != nil {
+ return nil, err
+ }
+ defer datum.Release()
+
+ switch out := datum.(type) {
+ case *ArrayDatum:
+ return out.MakeArray(), nil
+ case *ChunkedDatum:
+ return array.Concatenate(out.Chunks(), GetAllocator(ctx))
+ default:
+ return nil, fmt.Errorf(
+ "%w: dictionary_encode returned unexpected datum kind
%s",
+ arrow.ErrInvalid,
+ datum.Kind(),
+ )
+ }
+}
+
func RegisterVectorHash(reg FunctionRegistry) {
- unique, _, _ := kernels.GetVectorHashKernels()
+ unique, _, dictEncode := kernels.GetVectorHashKernels()
uniqFn := NewVectorFunction("unique", Unary(), uniqueDoc)
for _, vd := range unique {
if err := uniqFn.AddKernel(vd); err != nil {
@@ -56,4 +111,13 @@ func RegisterVectorHash(reg FunctionRegistry) {
}
}
reg.AddFunction(uniqFn, false)
+
+ dictFn := NewVectorFunction("dictionary_encode", Unary(),
dictionaryEncodeDoc)
+ dictFn.SetDefaultOptions(&DictionaryEncodeOptions{})
+ for _, vd := range dictEncode {
+ if err := dictFn.AddKernel(vd); err != nil {
+ panic(err)
+ }
+ }
+ reg.AddFunction(dictFn, false)
}
diff --git a/arrow/compute/vector_hash_test.go
b/arrow/compute/vector_hash_test.go
index 1c1b0179..2dc44f8c 100644
--- a/arrow/compute/vector_hash_test.go
+++ b/arrow/compute/vector_hash_test.go
@@ -29,6 +29,7 @@ import (
"github.com/apache/arrow-go/v18/arrow/decimal128"
"github.com/apache/arrow-go/v18/arrow/decimal256"
"github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/arrow/scalar"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
@@ -296,6 +297,27 @@ func TestHashKernels(t *testing.T) {
suite.Run(t, &BinaryTypeHashKernelSuite[[]byte]{dt:
arrow.BinaryTypes.LargeBinary})
}
+func TestUniqueBoolean(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ input, _, err := array.FromJSON(mem, arrow.FixedWidthTypes.Boolean,
+ strings.NewReader(`[false, null, true, false, null, true]`))
+ require.NoError(t, err)
+ defer input.Release()
+ expected, _, err := array.FromJSON(mem, arrow.FixedWidthTypes.Boolean,
+ strings.NewReader(`[false, null, true]`))
+ require.NoError(t, err)
+ defer expected.Release()
+
+ result, err := compute.UniqueArray(ctx, input)
+ require.NoError(t, err)
+ defer result.Release()
+
+ assert.True(t, array.Equal(expected, result))
+}
+
func TestUniqueTimeTimestamp(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
defer mem.AssertSize(t, 0)
@@ -515,3 +537,730 @@ func TestDictionaryUnique(t *testing.T) {
})
}
}
+
+func TestDictionaryEncode(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ values, _, err := array.FromJSON(mem, arrow.BinaryTypes.String,
+ strings.NewReader(`["foo", "bar", "foo", null, "bar", null]`))
+ require.NoError(t, err)
+ defer values.Release()
+
+ out, err := compute.DictionaryEncode(ctx,
compute.DictionaryEncodeOptions{},
+ &compute.ArrayDatum{Value: values.Data()})
+ require.NoError(t, err)
+ defer out.Release()
+
+ result := out.(*compute.ArrayDatum).MakeArray().(*array.Dictionary)
+ defer result.Release()
+
+ require.Equal(t, values.Len(), result.Len())
+ require.Equal(t, &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int32,
+ ValueType: arrow.BinaryTypes.String,
+ }, result.DataType())
+ require.Equal(t, 2, result.Dictionary().Len())
+ require.Equal(t, "foo", result.Dictionary().ValueStr(0))
+ require.Equal(t, "bar", result.Dictionary().ValueStr(1))
+ assert.Equal(t, []int32{0, 1, 0, 0, 1, 0},
+
arrow.Int32Traits.CastFromBytes(result.Indices().Data().Buffers()[1].Bytes()))
+ assert.True(t, result.IsNull(3))
+ assert.True(t, result.IsNull(5))
+
+ encoded, err := compute.DictionaryEncode(ctx,
compute.DictionaryEncodeOptions{
+ NullEncoding: compute.NullEncodingEncode,
+ }, &compute.ArrayDatum{Value: values.Data()})
+ require.NoError(t, err)
+ defer encoded.Release()
+
+ encodedResult :=
encoded.(*compute.ArrayDatum).MakeArray().(*array.Dictionary)
+ defer encodedResult.Release()
+
+ require.Equal(t, values.Len(), encodedResult.Len())
+ assert.Equal(t, 0, encodedResult.NullN())
+ require.Equal(t, 3, encodedResult.Dictionary().Len())
+ assert.Equal(t, 1, encodedResult.Dictionary().NullN())
+ assert.Equal(t, []int32{0, 1, 0, 2, 1, 2},
+
arrow.Int32Traits.CastFromBytes(encodedResult.Indices().Data().Buffers()[1].Bytes()))
+ for i := 0; i < values.Len(); i++ {
+ assert.True(t, encodedResult.IsValid(i))
+ idx := encodedResult.GetValueIndex(i)
+ if values.IsNull(i) {
+ assert.True(t, encodedResult.Dictionary().IsNull(idx))
+ } else {
+ assert.Equal(t, values.ValueStr(i),
encodedResult.Dictionary().ValueStr(idx))
+ }
+ }
+
+ _, err = compute.DictionaryEncode(ctx, compute.DictionaryEncodeOptions{
+ NullEncoding: compute.NullEncodingBehavior(99),
+ }, &compute.ArrayDatum{Value: values.Data()})
+ require.ErrorIs(t, err, arrow.ErrInvalid)
+}
+
+func TestDictionaryEncodeArrayWithSmallExecChunkSize(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ values, _, err := array.FromJSON(
+ mem,
+ arrow.PrimitiveTypes.Int32,
+ strings.NewReader(`[1, 2, 1, 3, 2]`),
+ )
+ require.NoError(t, err)
+ defer values.Release()
+
+ execCtx := compute.DefaultExecCtx()
+ execCtx.ChunkSize = 2
+ ctx := compute.SetExecCtx(
+ compute.WithAllocator(context.Background(), mem),
+ execCtx,
+ )
+
+ result, err := compute.DictionaryEncodeArray(
+ ctx,
+ compute.DictionaryEncodeOptions{},
+ values,
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ encoded := result.(*array.Dictionary)
+ require.Equal(t, values.Len(), encoded.Len())
+ require.Equal(t, 3, encoded.Dictionary().Len())
+ assert.Equal(t, []int32{0, 1, 0, 2, 1},
+
arrow.Int32Traits.CastFromBytes(encoded.Indices().Data().Buffers()[1].Bytes()))
+}
+
+func TestDictionaryEncodeBoolean(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ values, _, err := array.FromJSON(mem, arrow.FixedWidthTypes.Boolean,
+ strings.NewReader(`[false, true, false, null, true]`))
+ require.NoError(t, err)
+ defer values.Release()
+
+ tests := []struct {
+ name string
+ nullMode compute.NullEncodingBehavior
+ dict []bool
+ dictLen int
+ indices []int32
+ nullCount int
+ }{
+ {
+ name: "masked nulls",
+ nullMode: compute.NullEncodingMask,
+ dict: []bool{false, true},
+ dictLen: 2,
+ indices: []int32{0, 1, 0, 0, 1},
+ nullCount: 1,
+ },
+ {
+ name: "encoded nulls",
+ nullMode: compute.NullEncodingEncode,
+ dict: []bool{false, true},
+ dictLen: 3,
+ indices: []int32{0, 1, 0, 2, 1},
+ nullCount: 0,
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ result, err := compute.DictionaryEncodeArray(ctx,
compute.DictionaryEncodeOptions{
+ NullEncoding: tc.nullMode,
+ }, values)
+ require.NoError(t, err)
+ defer result.Release()
+
+ encoded := result.(*array.Dictionary)
+ dict := encoded.Dictionary().(*array.Boolean)
+ require.Equal(t, tc.dictLen, dict.Len())
+ for i, value := range tc.dict {
+ assert.Equal(t, value, dict.Value(i))
+ }
+ assert.Equal(t, tc.indices,
+
arrow.Int32Traits.CastFromBytes(encoded.Indices().Data().Buffers()[1].Bytes()))
+ assert.Equal(t, tc.nullCount, encoded.NullN())
+ if tc.nullMode == compute.NullEncodingMask {
+ assert.True(t, encoded.IsNull(3))
+ } else {
+ assert.True(t, encoded.IsValid(3))
+ assert.True(t,
dict.IsNull(encoded.GetValueIndex(3)))
+ }
+ })
+ }
+}
+
+func TestDictionaryEncodeBooleanSlicedInput(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ values, _, err := array.FromJSON(mem, arrow.FixedWidthTypes.Boolean,
+ strings.NewReader(`[true, false, true, null, true, false,
true]`))
+ require.NoError(t, err)
+ defer values.Release()
+
+ input := array.NewSlice(values, 1, 6)
+ defer input.Release()
+ require.Equal(t, 1, input.Data().Offset())
+
+ tests := []struct {
+ name string
+ nullMode compute.NullEncodingBehavior
+ dict []bool
+ dictLen int
+ indices []int32
+ nullCount int
+ }{
+ {
+ name: "masked nulls",
+ nullMode: compute.NullEncodingMask,
+ dict: []bool{false, true},
+ dictLen: 2,
+ indices: []int32{0, 1, 0, 1, 0},
+ nullCount: 1,
+ },
+ {
+ name: "encoded nulls",
+ nullMode: compute.NullEncodingEncode,
+ dict: []bool{false, true},
+ dictLen: 3,
+ indices: []int32{0, 1, 2, 1, 0},
+ nullCount: 0,
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ result, err := compute.DictionaryEncodeArray(ctx,
compute.DictionaryEncodeOptions{
+ NullEncoding: tc.nullMode,
+ }, input)
+ require.NoError(t, err)
+ defer result.Release()
+
+ encoded := result.(*array.Dictionary)
+ dict := encoded.Dictionary().(*array.Boolean)
+ require.Equal(t, tc.dictLen, dict.Len())
+ for i, value := range tc.dict {
+ assert.Equal(t, value, dict.Value(i))
+ }
+ assert.Equal(t, tc.indices,
+
arrow.Int32Traits.CastFromBytes(encoded.Indices().Data().Buffers()[1].Bytes()))
+ assert.Equal(t, tc.nullCount, encoded.NullN())
+ if tc.nullMode == compute.NullEncodingMask {
+ assert.True(t, encoded.IsNull(2))
+ } else {
+ assert.True(t, encoded.IsValid(2))
+ assert.True(t, dict.IsNull(len(tc.dict)))
+ assert.True(t,
dict.IsNull(encoded.GetValueIndex(2)))
+ }
+ })
+ }
+}
+
+func TestDictionaryEncodeResizesMemoTable(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ inputValues := make([]int32, 200)
+ for i := range inputValues {
+ inputValues[i] = int32(i % 100)
+ }
+
+ builder := array.NewInt32Builder(mem)
+ builder.AppendValues(inputValues, nil)
+ input := builder.NewInt32Array()
+ builder.Release()
+ defer input.Release()
+
+ result, err := compute.DictionaryEncodeArray(ctx,
compute.DictionaryEncodeOptions{}, input)
+ require.NoError(t, err)
+ defer result.Release()
+
+ encoded := result.(*array.Dictionary)
+ dictionary := encoded.Dictionary().(*array.Int32)
+ indices := encoded.Indices().(*array.Int32)
+
+ expectedDictionary := inputValues[:100]
+ require.Equal(t, expectedDictionary, dictionary.Int32Values())
+ require.Equal(t, inputValues, indices.Int32Values())
+}
+
+func TestDictionaryEncodeArraySlicedInput(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ values, _, err := array.FromJSON(mem, arrow.BinaryTypes.String,
+ strings.NewReader(`["ignored", "foo", null, "bar", "foo",
"ignored"]`))
+ require.NoError(t, err)
+ defer values.Release()
+
+ input := array.NewSlice(values, 1, 5)
+ defer input.Release()
+
+ result, err := compute.DictionaryEncodeArray(ctx,
compute.DictionaryEncodeOptions{}, input)
+ require.NoError(t, err)
+ defer result.Release()
+
+ encoded := result.(*array.Dictionary)
+ require.Equal(t, 4, encoded.Len())
+ require.Equal(t, 2, encoded.Dictionary().Len())
+ assert.Equal(t, "foo", encoded.Dictionary().ValueStr(0))
+ assert.Equal(t, "bar", encoded.Dictionary().ValueStr(1))
+ assert.Equal(t, []int32{0, 0, 1, 0},
+
arrow.Int32Traits.CastFromBytes(encoded.Indices().Data().Buffers()[1].Bytes()))
+ assert.True(t, encoded.IsNull(1))
+}
+
+func TestDictionaryEncodePreservesTimestampType(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ dt := &arrow.TimestampType{Unit: arrow.Second, TimeZone: "UTC"}
+ builder := array.NewTimestampBuilder(mem, dt)
+ builder.AppendValues([]arrow.Timestamp{1, 2, 1}, nil)
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ result, err := compute.DictionaryEncodeArray(ctx,
compute.DictionaryEncodeOptions{}, input)
+ require.NoError(t, err)
+ defer result.Release()
+
+ encoded := result.(*array.Dictionary)
+ require.True(t, arrow.TypeEqual(dt, encoded.Dictionary().DataType()))
+ require.Equal(t, 2, encoded.Dictionary().Len())
+}
+
+func TestDictionaryEncodeZeroChunkedArray(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ input := arrow.NewChunked(arrow.BinaryTypes.String, nil)
+ defer input.Release()
+
+ out, err := compute.DictionaryEncode(ctx,
compute.DictionaryEncodeOptions{},
+ &compute.ChunkedDatum{Value: input})
+ require.NoError(t, err)
+ defer out.Release()
+
+ result := out.(*compute.ChunkedDatum).Value
+ assert.Equal(t, 0, result.Len())
+ assert.Empty(t, result.Chunks())
+ assert.True(t, arrow.TypeEqual(&arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int32,
+ ValueType: arrow.BinaryTypes.String,
+ }, result.DataType()))
+}
+
+func TestDictionaryEncodeZeroChunkedDictionaryInput(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ input := arrow.NewChunked(dictType, nil)
+ defer input.Release()
+
+ out, err := compute.DictionaryEncode(ctx,
compute.DictionaryEncodeOptions{},
+ &compute.ChunkedDatum{Value: input})
+ require.NoError(t, err)
+ defer out.Release()
+
+ result := out.(*compute.ChunkedDatum).Value
+ require.Equal(t, 0, result.Len())
+ require.Empty(t, result.Chunks())
+ require.True(t, arrow.TypeEqual(dictType, result.DataType()))
+}
+
+func TestDictionaryEncodeChunked(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ first, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int32,
strings.NewReader(`[1, 2, 1]`))
+ require.NoError(t, err)
+ defer first.Release()
+ second, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int32,
strings.NewReader(`[2, 3, null]`))
+ require.NoError(t, err)
+ defer second.Release()
+
+ input := arrow.NewChunked(arrow.PrimitiveTypes.Int32,
[]arrow.Array{first, second})
+ defer input.Release()
+
+ out, err := compute.DictionaryEncode(ctx,
compute.DictionaryEncodeOptions{},
+ &compute.ChunkedDatum{Value: input})
+ require.NoError(t, err)
+ defer out.Release()
+
+ result := out.(*compute.ChunkedDatum).Value
+ require.Len(t, result.Chunks(), 2)
+ require.Equal(t, first.Len(), result.Chunk(0).Len())
+ require.Equal(t, second.Len(), result.Chunk(1).Len())
+ assert.Equal(t, []int32{0, 1, 0},
+
arrow.Int32Traits.CastFromBytes(result.Chunk(0).(*array.Dictionary).Indices().Data().Buffers()[1].Bytes()))
+ assert.Equal(t, []int32{1, 2, 0},
+
arrow.Int32Traits.CastFromBytes(result.Chunk(1).(*array.Dictionary).Indices().Data().Buffers()[1].Bytes()))
+ assert.True(t, result.Chunk(1).IsNull(2))
+ for _, chunk := range result.Chunks() {
+ encoded := chunk.(*array.Dictionary)
+ assert.True(t, array.Equal(encoded.Dictionary(),
result.Chunk(0).(*array.Dictionary).Dictionary()))
+ }
+}
+
+func TestDictionaryEncodeChunkedBinary(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ types := []struct {
+ name string
+ typ arrow.DataType
+ }{
+ {name: "string", typ: arrow.BinaryTypes.String},
+ {name: "large string", typ: arrow.BinaryTypes.LargeString},
+ }
+ modes := []struct {
+ name string
+ nullMode compute.NullEncodingBehavior
+ dict []string
+ indices [][]int32
+ nullCount []int
+ }{
+ {
+ name: "masked nulls",
+ nullMode: compute.NullEncodingMask,
+ dict: []string{"a", "b", "c"},
+ indices: [][]int32{{0, 1, 0, 0}, {1, 2, 0, 0}},
+ nullCount: []int{1, 1},
+ },
+ {
+ name: "encoded nulls",
+ nullMode: compute.NullEncodingEncode,
+ dict: []string{"a", "b", "(null)", "c"},
+ indices: [][]int32{{0, 1, 2, 0}, {1, 3, 2, 0}},
+ nullCount: []int{0, 0},
+ },
+ }
+
+ for _, typ := range types {
+ for _, mode := range modes {
+ t.Run(typ.name+"/"+mode.name, func(t *testing.T) {
+ first, _, err := array.FromJSON(mem, typ.typ,
strings.NewReader(`["a", "b", null, "a"]`))
+ require.NoError(t, err)
+ defer first.Release()
+ second, _, err := array.FromJSON(mem, typ.typ,
strings.NewReader(`["b", "c", null, "a"]`))
+ require.NoError(t, err)
+ defer second.Release()
+
+ input := arrow.NewChunked(typ.typ,
[]arrow.Array{first, second})
+ defer input.Release()
+
+ out, err := compute.DictionaryEncode(ctx,
compute.DictionaryEncodeOptions{
+ NullEncoding: mode.nullMode,
+ }, &compute.ChunkedDatum{Value: input})
+ require.NoError(t, err)
+ defer out.Release()
+
+ result := out.(*compute.ChunkedDatum).Value
+ require.Len(t, result.Chunks(), 2)
+ require.True(t,
arrow.TypeEqual(result.DataType(), &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int32,
+ ValueType: typ.typ,
+ }))
+ for i, chunk := range result.Chunks() {
+ encoded := chunk.(*array.Dictionary)
+ require.Equal(t, len(mode.dict),
encoded.Dictionary().Len())
+ for j, value := range mode.dict {
+ assert.Equal(t, value,
encoded.Dictionary().ValueStr(j))
+ }
+ assert.Equal(t, mode.indices[i],
+
arrow.Int32Traits.CastFromBytes(encoded.Indices().Data().Buffers()[1].Bytes()))
+ assert.Equal(t, mode.nullCount[i],
encoded.NullN())
+ if mode.nullMode ==
compute.NullEncodingMask {
+ assert.True(t,
encoded.IsNull(2))
+ } else {
+ assert.True(t,
encoded.IsValid(2))
+ assert.True(t,
encoded.Dictionary().IsNull(encoded.GetValueIndex(2)))
+ }
+ }
+ })
+ }
+ }
+}
+
+func TestDictionaryEncodeNullArray(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ input := array.NewNull(3)
+ defer input.Release()
+
+ tests := []struct {
+ name string
+ nullEncoding compute.NullEncodingBehavior
+ dictLen int
+ nullCount int
+ }{
+ {name: "masked", nullEncoding: compute.NullEncodingMask,
dictLen: 1, nullCount: 3},
+ {name: "encoded", nullEncoding: compute.NullEncodingEncode,
dictLen: 1, nullCount: 0},
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ out, err := compute.DictionaryEncode(ctx,
compute.DictionaryEncodeOptions{
+ NullEncoding: tc.nullEncoding,
+ }, &compute.ArrayDatum{Value: input.Data()})
+ require.NoError(t, err)
+ defer out.Release()
+
+ result :=
out.(*compute.ArrayDatum).MakeArray().(*array.Dictionary)
+ defer result.Release()
+ require.Equal(t, input.Len(), result.Len())
+ assert.Equal(t, tc.dictLen, result.Dictionary().Len())
+ assert.Equal(t, tc.nullCount, result.NullN())
+ if tc.nullEncoding == compute.NullEncodingEncode {
+ for i := 0; i < result.Len(); i++ {
+ assert.True(t, result.IsValid(i))
+ }
+ }
+ })
+ }
+}
+
+func TestDictionaryEncodeDictionaryInput(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ input, err := array.DictArrayFromJSON(mem, dictType, `[0, null, 1]`,
`["foo", "bar"]`)
+ require.NoError(t, err)
+ defer input.Release()
+
+ out, err := compute.DictionaryEncode(ctx,
compute.DictionaryEncodeOptions{},
+ &compute.ArrayDatum{Value: input.Data()})
+ require.NoError(t, err)
+ defer out.Release()
+
+ result := out.(*compute.ArrayDatum).MakeArray().(*array.Dictionary)
+ defer result.Release()
+ require.True(t, arrow.TypeEqual(dictType, result.DataType()))
+ require.Equal(t, input.Len(), result.Len())
+ assert.True(t, array.Equal(input, result))
+}
+
+func TestDictionaryEncodeOptionsUseCanonicalFieldName(t *testing.T) {
+ for _, tt := range []struct {
+ name string
+ behavior compute.NullEncodingBehavior
+ encoded uint32
+ }{
+ {name: "encode", behavior: compute.NullEncodingEncode, encoded:
0},
+ {name: "mask", behavior: compute.NullEncodingMask, encoded: 1},
+ } {
+ t.Run(tt.name, func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ encoded, err :=
scalar.ToScalar(compute.DictionaryEncodeOptions{
+ NullEncoding: tt.behavior,
+ }, mem)
+ require.NoError(t, err)
+ if releasable, ok := encoded.(interface{ Release() });
ok {
+ defer releasable.Release()
+ }
+
+ options := encoded.(*scalar.Struct)
+ field, err := options.Field("null_encoding_behavior")
+ require.NoError(t, err)
+ assert.Equal(t, tt.encoded,
field.(*scalar.Uint32).Value)
+ _, err = options.Field("null_encoding")
+ assert.Error(t, err)
+ })
+ }
+}
+
+func TestDictionaryEncodeOptionsDecodeCanonicalValues(t *testing.T) {
+ for _, tt := range []struct {
+ name string
+ encoded uint32
+ behavior compute.NullEncodingBehavior
+ }{
+ {name: "encode", encoded: 0, behavior:
compute.NullEncodingEncode},
+ {name: "mask", encoded: 1, behavior: compute.NullEncodingMask},
+ } {
+ t.Run(tt.name, func(t *testing.T) {
+ encoded, err := scalar.NewStructScalarWithNames(
+
[]scalar.Scalar{scalar.NewUint32Scalar(tt.encoded)},
+ []string{"null_encoding_behavior"},
+ )
+ require.NoError(t, err)
+ defer encoded.Release()
+
+ var options compute.DictionaryEncodeOptions
+ require.NoError(t, scalar.FromScalar(encoded, &options))
+ assert.Equal(t, tt.behavior, options.NullEncoding)
+ })
+ }
+}
+
+func TestDictionaryEncodeOptionsRejectInvalidCanonicalValue(t *testing.T) {
+ encoded, err := scalar.NewStructScalarWithNames(
+ []scalar.Scalar{scalar.NewUint32Scalar(2)},
+ []string{"null_encoding_behavior"},
+ )
+ require.NoError(t, err)
+ defer encoded.Release()
+
+ var options compute.DictionaryEncodeOptions
+ err = scalar.FromScalar(encoded, &options)
+ assert.ErrorIs(t, err, arrow.ErrInvalid)
+}
+
+func TestDictionaryEncodeOptionsRejectInvalidSerialization(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ _, err := scalar.ToScalar(compute.DictionaryEncodeOptions{
+ NullEncoding: compute.NullEncodingBehavior(42),
+ }, mem)
+ assert.ErrorIs(t, err, arrow.ErrInvalid)
+}
+
+func TestDictionaryEncodeDictionaryInputWithEmptyIndices(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ input, err := array.DictArrayFromJSON(mem, dictType, `[]`, `["foo",
"bar"]`)
+ require.NoError(t, err)
+ defer input.Release()
+
+ out, err := compute.DictionaryEncode(ctx,
compute.DictionaryEncodeOptions{},
+ &compute.ArrayDatum{Value: input.Data()})
+ require.NoError(t, err)
+ defer out.Release()
+
+ result := out.(*compute.ArrayDatum).MakeArray().(*array.Dictionary)
+ defer result.Release()
+ require.Equal(t, 0, result.Len())
+ require.True(t, arrow.TypeEqual(dictType, result.DataType()))
+ require.Equal(t, 2, result.Dictionary().Len())
+ assert.Equal(t, "foo", result.Dictionary().ValueStr(0))
+ assert.Equal(t, "bar", result.Dictionary().ValueStr(1))
+}
+
+func TestDictionaryEncodeChunkedDictionaryInput(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ first, err := array.DictArrayFromJSON(mem, dictType, `[0, 1]`, `["foo",
"bar"]`)
+ require.NoError(t, err)
+ defer first.Release()
+ second, err := array.DictArrayFromJSON(mem, dictType, `[0, 1]`,
`["bar", "baz"]`)
+ require.NoError(t, err)
+ defer second.Release()
+
+ input := arrow.NewChunked(dictType, []arrow.Array{first, second})
+ defer input.Release()
+
+ out, err := compute.DictionaryEncode(ctx,
compute.DictionaryEncodeOptions{},
+ &compute.ChunkedDatum{Value: input})
+ require.NoError(t, err)
+ defer out.Release()
+
+ result := out.(*compute.ChunkedDatum).Value
+ require.Len(t, result.Chunks(), 2)
+ assert.True(t, array.Equal(first, result.Chunk(0)))
+ assert.True(t, array.Equal(second, result.Chunk(1)))
+ assert.Equal(t, "foo",
result.Chunk(0).(*array.Dictionary).Dictionary().ValueStr(0))
+ assert.Equal(t, "baz",
result.Chunk(1).(*array.Dictionary).Dictionary().ValueStr(1))
+}
+
+func TestDictionaryEncodeArrayDictionaryInputWithSmallExecChunkSize(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ input, err := array.DictArrayFromJSON(mem, dictType, `[0, 1, 0, null,
1]`, `["foo", "bar"]`)
+ require.NoError(t, err)
+ defer input.Release()
+
+ execCtx := compute.DefaultExecCtx()
+ execCtx.ChunkSize = 2
+ ctx := compute.SetExecCtx(
+ compute.WithAllocator(context.Background(), mem),
+ execCtx,
+ )
+
+ result, err := compute.DictionaryEncodeArray(
+ ctx,
+ compute.DictionaryEncodeOptions{},
+ input,
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ require.True(t, arrow.TypeEqual(dictType, result.DataType()))
+ assert.True(t, array.Equal(input, result))
+}
+
+func TestDictionaryEncodeDictionaryInputRejectsInvalidOptions(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ input, err := array.DictArrayFromJSON(mem, dictType, `[0, 1]`, `["foo",
"bar"]`)
+ require.NoError(t, err)
+ defer input.Release()
+
+ _, err = compute.DictionaryEncode(ctx, compute.DictionaryEncodeOptions{
+ NullEncoding: compute.NullEncodingBehavior(99),
+ }, &compute.ArrayDatum{Value: input.Data()})
+ require.ErrorIs(t, err, arrow.ErrInvalid)
+}
+
+func TestDictionaryEncodeFunctionValidation(t *testing.T) {
+ for _, name := range []string{"unique", "dictionary_encode"} {
+ t.Run(name, func(t *testing.T) {
+ fn, ok :=
compute.GetFunctionRegistry().GetFunction(name)
+ require.True(t, ok)
+ require.NoError(t, fn.Validate())
+ })
+ }
+}
