zeroshade commented on code in PR #1395:
URL: https://github.com/apache/arrow-go/pull/1395#discussion_r4339836247
##########
arrow/compute/internal/kernels/vector_hash.go:
##########
@@ -526,56 +545,113 @@ func (dhs *dictionaryHashState) GetDictionary()
(arrow.ArrayData, error) {
func (dhs *dictionaryHashState) ValueType() arrow.DataType { return
dhs.indicesKernel.ValueType() }
func (dhs *dictionaryHashState) DictionaryValueType() arrow.DataType { return
dhs.dictValueType }
func (dhs *dictionaryHashState) Dictionary() arrow.Array { return
dhs.dictionary }
-func (dhs *dictionaryHashState) Append(ctx *exec.KernelCtx, arr
*exec.ArraySpan) error {
- arrDict := arr.Dictionary().MakeArray()
- if dhs.dictionary == nil || array.Equal(dhs.dictionary, arrDict) {
- dhs.dictionary = arrDict
- return dhs.indicesKernel.Append(ctx, arr)
- }
+func (dhs *dictionaryHashState) Append(ctx *exec.KernelCtx, arr
*exec.ArraySpan) (err error) {
+ // A failed append must not retain the unifier's allocator-backed
buffers.
+ defer func() {
+ if err != nil {
+ dhs.releaseDictionary()
+ }
+ }()
+ arrDict := arr.Dictionary().MakeArray()
defer arrDict.Release()
- // NOTE: this approach computes a new dictionary unification per chunk
- // this is in effect O(n*k) where n is the total chunked array length
- // and k is the number of chunks (therefore O(n**2) if chunks have a
fixed size).
- //
- // A better approach may be to run the kernel over each individual
chunk,
- // and then hash-aggregate all results (for example sum-group-by for
- // the "value_counts" kernel)
- unifier, err :=
array.NewDictionaryUnifier(dhs.indicesKernel.Allocator(), dhs.dictValueType)
- if err != nil {
- return err
- }
- defer unifier.Release()
+ if dhs.unifier == nil {
+ if dhs.dictionary == nil {
+ // Preserve the original indices and dictionary for the
common case
+ // where all chunks use the same dictionary.
+ arrDict.Retain()
+ dhs.dictionary = arrDict
+ return dhs.indicesKernel.Append(ctx, arr)
+ }
+ if array.Equal(dhs.dictionary, arrDict) {
+ return dhs.indicesKernel.Append(ctx, arr)
+ }
- if err := unifier.Unify(dhs.dictionary); err != nil {
- return err
+ // Seed the unifier only once. Re-unifying the full accumulated
+ // dictionary for every chunk makes many changing dictionaries
quadratic.
+ dhs.unifier, err = array.NewDictionaryUnifier(dhs.Allocator(),
dhs.dictValueType)
+ if err != nil {
+ return err
+ }
+ if err = dhs.unifier.Unify(dhs.dictionary); err != nil {
+ return err
+ }
+ dhs.dictionary.Release()
+ dhs.dictionary = nil
}
- transposeMap, err := unifier.UnifyAndTranspose(arrDict)
+
+ transpose, err := dhs.unifier.UnifyAndTranspose(arrDict)
if err != nil {
return err
}
- defer transposeMap.Release()
- _, outDict, err := unifier.GetResult()
- if err != nil {
- return err
+ defer transpose.Release()
+ mapping := arrow.Int32Traits.CastFromBytes(transpose.Bytes())
+
+ // The unique kernel keeps the input dictionary's index type. Reject
+ // transpositions which would truncate a newly assigned dictionary
index.
+ var maxIndex int32 = math.MaxInt32
+ switch dhs.indexType.ID() {
+ case arrow.INT8:
+ maxIndex = math.MaxInt8
+ case arrow.UINT8:
+ maxIndex = math.MaxUint8
+ case arrow.INT16:
+ maxIndex = math.MaxInt16
+ case arrow.UINT16:
+ maxIndex = math.MaxUint16
+ }
+ if maxIndex != math.MaxInt32 {
+ for _, mapped := range mapping {
Review Comment:
This checks every merged entry, including ones no row refers to. With `int8`
chunks whose dictionaries are `["v0" … "v127"]` and `["unused", "v0"]` and
whose indices are `[0]` and `[1]`, `main` returns `"v0"`, but this errors
because the unreferenced `"unused"` is assigned ID 128. Could you check
capacity only for the IDs that valid rows actually map to, and add this case as
a test?
##########
arrow/compute/vector_hash_dictionary_incremental_test.go:
##########
@@ -0,0 +1,183 @@
+// 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_test
+
+import (
+ "context"
+ "fmt"
+ "strings"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/compute"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestUniqueChunkedChangingDictionaries(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int16,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ makeDict := func(indices, values string) *array.Dictionary {
+ t.Helper()
+ arr, err := array.DictArrayFromJSON(mem, dictType, indices,
values)
+ require.NoError(t, err)
+ return arr
+ }
+
+ first := makeDict("[0, 1, 0]", `["a", "b"]`)
+ defer first.Release()
+ second := makeDict("[0, 1, 0]", `["b", "c"]`)
+ defer second.Release()
+ // Distinct arrays with equal dictionaries, after unification began.
+ third := makeDict("[1, 0, 1]", `["b", "c"]`)
+ defer third.Release()
+ full := makeDict("[null, 0, 1, 2, null]", `["d", "a", "e"]`)
+ defer full.Release()
+ sliced := array.NewSlice(full, 1, 5)
+ defer sliced.Release()
+ require.Equal(t, 1, sliced.Data().Offset())
+ revisit := makeDict("[1, 0]", `["a", "b"]`)
+ defer revisit.Release()
+
+ input := arrow.NewChunked(dictType, []arrow.Array{first, second, third,
sliced, revisit})
+ defer input.Release()
+ out, err := compute.Unique(compute.WithAllocator(context.Background(),
mem),
+ &compute.ChunkedDatum{Value: input})
+ require.NoError(t, err)
+ defer out.Release()
+
+ got := out.(*compute.ArrayDatum).MakeArray().(*array.Dictionary)
+ defer got.Release()
+ require.Equal(t, 6, got.Len()) // five distinct values plus the masked
null
+ require.Equal(t, 5, got.Dictionary().Len())
+ for i, value := range []string{"a", "b", "c", "d", "e"} {
+ assert.Equal(t, value, got.Dictionary().ValueStr(i))
+ assert.Equal(t, i, got.GetValueIndex(i))
+ }
+ assert.True(t, got.IsNull(5))
+}
+
+func TestUniqueChunkedDictionaryValueNull(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int16,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ first, err := array.DictArrayFromJSON(mem, dictType, "[0, 1]", `["one",
null]`)
+ require.NoError(t, err)
+ defer first.Release()
+ second, err := array.DictArrayFromJSON(mem, dictType, "[0, 1, 2, null]",
+ `[null, "one", "two"]`)
+ require.NoError(t, err)
+ defer second.Release()
+
+ input := arrow.NewChunked(dictType, []arrow.Array{first, second})
+ defer input.Release()
+ out, err := compute.Unique(compute.WithAllocator(context.Background(),
mem),
+ &compute.ChunkedDatum{Value: input})
+ require.NoError(t, err)
+ defer out.Release()
+
+ got := out.(*compute.ArrayDatum).MakeArray().(*array.Dictionary)
+ defer got.Release()
+ require.Equal(t, 4, got.Len()) // "one", dictionary null, "two", index
null
+ require.Equal(t, 3, got.Dictionary().Len())
+ assert.Equal(t, "one", got.Dictionary().ValueStr(0))
+ assert.True(t, got.Dictionary().IsNull(1))
+ assert.Equal(t, "two", got.Dictionary().ValueStr(2))
+ assert.True(t, got.IsValid(1)) // valid index pointing to a null
dictionary value
+ assert.True(t, got.IsNull(3))
+}
+
+func TestUniqueChunkedDictionarySkipsOutOfBoundsNullIndex(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ first, err := array.DictArrayFromJSON(mem, dictType, "[0, 1]", `["a",
"b"]`)
+ require.NoError(t, err)
+ defer first.Release()
+
+ dict, _, err := array.FromJSON(mem, arrow.BinaryTypes.String,
+ strings.NewReader(`["b", "c"]`))
+ require.NoError(t, err)
+ defer dict.Release()
+
+ builder := array.NewInt8Builder(mem)
+ builder.Append(0)
+ builder.AppendNull()
+ builder.Append(1)
+ indices := builder.NewInt8Array()
+ builder.Release()
+ defer indices.Release()
+
+ // Validity bits mask this slot, so its physical value is unspecified.
+ indices.Data().Buffers()[1].Bytes()[1] = 127
+ second := array.NewDictionaryArray(dictType, indices, dict)
+ defer second.Release()
+
+ input := arrow.NewChunked(dictType, []arrow.Array{first, second})
+ defer input.Release()
+ out, err := compute.Unique(compute.WithAllocator(context.Background(),
mem),
+ &compute.ChunkedDatum{Value: input})
+ require.NoError(t, err)
+ defer out.Release()
+
+ got := out.(*compute.ArrayDatum).MakeArray().(*array.Dictionary)
+ defer got.Release()
+ require.Equal(t, 4, got.Len())
+ assert.True(t, got.IsNull(2))
+ assert.Equal(t, "c", got.ValueStr(3))
+}
+
+func TestUniqueChunkedDictionaryIndexOverflow(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ chunks := make([]arrow.Array, 129)
+ for i := range chunks {
+ arr, err := array.DictArrayFromJSON(mem, dictType, "[0]",
+ fmt.Sprintf(`["value-%d"]`, i))
+ require.NoError(t, err)
+ chunks[i] = arr
+ defer arr.Release()
+ }
+ input := arrow.NewChunked(dictType, chunks)
+ defer input.Release()
+
+ _, err := compute.Unique(compute.WithAllocator(context.Background(),
mem),
+ &compute.ChunkedDatum{Value: input})
+ require.ErrorIs(t, err, arrow.ErrInvalid)
+}
Review Comment:
After fixing the helper's return type, this `AssertSize` fails with 16,512
bytes outstanding (also with `noasm`, `-race` and on 386). Returning the new
error skips `uniqueFinalize`, so the output buffers allocated for earlier
chunks, and the current chunk's, are never released. The cleanup lives in the
executor (`executor.go` around lines 990 and 1163), so it may need a change
there, or avoiding the output preallocation for this kernel. Please keep the
checked-allocator assertion.
##########
arrow/compute/vector_hash_dictionary_incremental_test.go:
##########
@@ -0,0 +1,183 @@
+// 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_test
+
+import (
+ "context"
+ "fmt"
+ "strings"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/compute"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestUniqueChunkedChangingDictionaries(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int16,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ makeDict := func(indices, values string) *array.Dictionary {
+ t.Helper()
+ arr, err := array.DictArrayFromJSON(mem, dictType, indices,
values)
+ require.NoError(t, err)
+ return arr
Review Comment:
`array.DictArrayFromJSON` returns `arrow.Array`, so it can't be returned
from a helper declared as `*array.Dictionary`. This is why the compute test
package fails to compile in CI. Returning `arrow.Array` (which the callers
already accept) or adding a type assertion fixes it.
##########
arrow/compute/vector_hash_dictionary_incremental_bench_test.go:
##########
@@ -0,0 +1,90 @@
+// 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_test
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+ "strings"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/compute"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+)
+
+// BenchmarkUniqueChunkedChangingDictionaries measures the many-chunk case.
+// Each chunk has new dictionary entries; total dictionary size grows with
chunks.
+func BenchmarkUniqueChunkedChangingDictionaries(b *testing.B) {
+ const (
+ valuesPerChunk = 128
+ dictSize = 16
+ )
+ mem := memory.DefaultAllocator
+ ctx := compute.WithAllocator(context.Background(), mem)
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int32,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ indices := make([]string, valuesPerChunk)
+ for i := range indices {
+ indices[i] = fmt.Sprint(i % dictSize)
+ }
+ indexJSON := "[" + strings.Join(indices, ",") + "]"
+
+ for _, numChunks := range []int{8, 32, 128, 256} {
+ b.Run(fmt.Sprintf("chunks=%d", numChunks), func(b *testing.B) {
+ chunks := make([]arrow.Array, numChunks)
+ for i := range chunks {
+ values := make([]string, dictSize)
+ for j := range values {
+ values[j] =
fmt.Sprintf("chunk-%04d-value-%02d", i, j)
+ }
+ dictJSON, err := json.Marshal(values)
+ if err != nil {
+ b.Fatal(err)
+ }
+ arr, err := array.DictArrayFromJSON(mem,
dictType, indexJSON, string(dictJSON))
+ if err != nil {
+ b.Fatal(err)
+ }
+ chunks[i] = arr
+ }
+ input := arrow.NewChunked(dictType, chunks)
+ for _, chunk := range chunks {
+ chunk.Release()
+ }
+ defer input.Release()
+
+ datum := &compute.ChunkedDatum{Value: input}
+ b.ReportAllocs()
+ b.SetBytes(int64(numChunks * valuesPerChunk))
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
Review Comment:
Nit: please use `for b.Loop()`, and bump this file's `//go:build go1.18` to
`go1.24` in the same change so the vet check that `go test` runs accepts it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
