feat: filter and take for dictionary arrays

Add 'take' and 'filter' implementations for dictionary arrays. These
implementations simply filter/take on the index column and shallow-copy
the dictionary.

Author: lesam

arrow/compute/selection.go

@@ -504,6 +504,33 @@ func structFilter(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResul
 	return nil
 }
 
+func dictionaryFilter(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult) error {
+	indices, err := kernels.GetTakeIndices(exec.GetAllocator(ctx.Ctx),
+		&batch.Values[1].Array, ctx.State.(kernels.FilterState).NullSelection)
+	if err != nil {
+		return err
+	}
+	defer indices.Release()
+
+	filter := NewDatum(indices)
+	defer filter.Release()
+
+	valData := batch.Values[0].Array.MakeData()
+	defer valData.Release()
+
+	vals := NewDatum(valData)
+	defer vals.Release()
+
+	result, err := Take(ctx.Ctx, kernels.TakeOptions{BoundsCheck: false}, vals, filter)
+	if err != nil {
+		return err
+	}
+	defer result.Release()
+
+	out.TakeOwnership(result.(*ArrayDatum).Value)
+	return nil
+}
+
 func structTake(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult) error {
 	// generate top level validity bitmap
 	if err := kernels.TakeExec(kernels.StructImpl)(ctx, batch, out); err != nil {
@@ -538,6 +565,26 @@ func structTake(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult)
 	return eg.Wait()
 }
 
+func dictionaryTake(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult) error {
+	dictArr := batch.Values[0].Array.MakeArray().(*array.Dictionary)
+	defer dictArr.Release()
+
+	selection := batch.Values[1].Array.MakeArray()
+	defer selection.Release()
+
+	takenIndices, err := TakeArrayOpts(ctx.Ctx, dictArr.Indices(), selection, kernels.TakeOptions{BoundsCheck: ctx.State.(kernels.TakeState).BoundsCheck})
+	if err != nil {
+		return err
+	}
+	defer takenIndices.Release()
+
+	result := array.NewDictionaryArray(dictArr.DataType(), takenIndices, dictArr.Dictionary())
+	defer result.Release()
+
+	out.TakeOwnership(result.Data())
+	return nil
+}
+
 // RegisterVectorSelection registers functions that select specific
 // values from arrays such as Take and Filter
 func RegisterVectorSelection(reg FunctionRegistry) {
@@ -552,6 +599,7 @@ func RegisterVectorSelection(reg FunctionRegistry) {
 		{In: exec.NewIDInput(arrow.LARGE_LIST), Exec: selectListImpl(kernels.FilterExec(kernels.ListImpl[int64]))},
 		{In: exec.NewIDInput(arrow.FIXED_SIZE_LIST), Exec: selectListImpl(kernels.FilterExec(kernels.FSLImpl))},
 		{In: exec.NewIDInput(arrow.DENSE_UNION), Exec: denseUnionImpl(kernels.FilterExec(kernels.DenseUnionImpl))},
+		{In: exec.NewIDInput(arrow.DICTIONARY), Exec: dictionaryFilter},
 		{In: exec.NewIDInput(arrow.EXTENSION), Exec: extensionFilterImpl},
 		{In: exec.NewIDInput(arrow.STRUCT), Exec: structFilter},
 	}...)
@@ -562,6 +610,7 @@ func RegisterVectorSelection(reg FunctionRegistry) {
 		{In: exec.NewIDInput(arrow.FIXED_SIZE_LIST), Exec: selectListImpl(kernels.TakeExec(kernels.FSLImpl))},
 		{In: exec.NewIDInput(arrow.MAP), Exec: selectMapImpl(kernels.TakeExec(kernels.MapImpl))},
 		{In: exec.NewIDInput(arrow.DENSE_UNION), Exec: denseUnionImpl(kernels.TakeExec(kernels.DenseUnionImpl))},
+		{In: exec.NewIDInput(arrow.DICTIONARY), Exec: dictionaryTake},
 		{In: exec.NewIDInput(arrow.EXTENSION), Exec: extensionTakeImpl},
 		{In: exec.NewIDInput(arrow.STRUCT), Exec: structTake},
 	}...)

arrow/compute/vector_selection_test.go

@@ -66,6 +66,21 @@ func (f *FilterKernelTestSuite) getArr(dt arrow.DataType, str string) arrow.Arra
 	return arr
 }
 
+// assertDictionaryLogicalEqual compares two dictionary arrays by decoding both to the
+// value type and comparing. Use when logical equality is desired (same decoded values)
+// rather than physical (same indices and dictionary). ctx must have an allocator.
+func assertDictionaryLogicalEqual(t *testing.T, ctx context.Context, expected, actual arrow.Array, opts ...array.EqualOption) bool {
+	valueType := expected.DataType().(*arrow.DictionaryType).ValueType
+	castOpts := compute.NewCastOptions(valueType, true)
+	decodedExpected, err := compute.CastArray(ctx, expected, castOpts)
+	require.NoError(t, err)
+	defer decodedExpected.Release()
+	decodedActual, err := compute.CastArray(ctx, actual, castOpts)
+	require.NoError(t, err)
+	defer decodedActual.Release()
+	return assertArraysEqual(t, decodedExpected, decodedActual, opts...)
+}
+
 func (f *FilterKernelTestSuite) doAssertFilter(values, filter, expected arrow.Array) {
 	ctx := compute.WithAllocator(context.TODO(), f.mem)
 	valDatum := compute.NewDatum(values)
@@ -79,7 +94,11 @@ func (f *FilterKernelTestSuite) doAssertFilter(values, filter, expected arrow.Ar
 		defer out.Release()
 		actual := out.(*compute.ArrayDatum).MakeArray()
 		defer actual.Release()
-		f.Truef(array.Equal(expected, actual), "expected: %s\ngot: %s", expected, actual)
+		if expected.DataType().ID() == arrow.DICTIONARY {
+			assertDictionaryLogicalEqual(f.T(), ctx, expected, actual)
+		} else {
+			f.Truef(array.Equal(expected, actual), "expected: %s\ngot: %s", expected, actual)
+		}
 	})
 
 	// f.Run("drop", func() {
@@ -174,7 +193,12 @@ func (tk *TakeKernelTestSuite) assertTakeArrays(values, indices, expected arrow.
 	actual, err := compute.TakeArray(tk.ctx, values, indices)
 	tk.Require().NoError(err)
 	defer actual.Release()
-	assertArraysEqual(tk.T(), expected, actual)
+
+	if expected.DataType().ID() == arrow.DICTIONARY {
+		assertDictionaryLogicalEqual(tk.T(), tk.ctx, expected, actual)
+	} else {
+		assertArraysEqual(tk.T(), expected, actual)
+	}
 }
 
 func (tk *TakeKernelTestSuite) takeJSON(dt arrow.DataType, values string, idxType arrow.DataType, indices string) (arrow.Array, error) {
@@ -1711,6 +1735,7 @@ func TestTakeKernels(t *testing.T) {
 	for _, dt := range baseBinaryTypes {
 		suite.Run(t, &TakeKernelTestString{TakeKernelTestTyped: TakeKernelTestTyped{dt: dt}})
 	}
+	suite.Run(t, &TakeKernelTestString{TakeKernelTestTyped: TakeKernelTestTyped{dt: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: arrow.BinaryTypes.String}}})
 	suite.Run(t, new(TakeKernelLists))
 	suite.Run(t, new(TakeKernelDenseUnion))
 	suite.Run(t, new(TakeKernelTestExtension))
@@ -1732,6 +1757,7 @@ func TestFilterKernels(t *testing.T) {
 	for _, dt := range baseBinaryTypes {
 		suite.Run(t, &FilterKernelWithString{dt: dt})
 	}
+	suite.Run(t, &FilterKernelWithString{dt: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: arrow.BinaryTypes.String}})
 	suite.Run(t, new(FilterKernelWithList))
 	suite.Run(t, new(FilterKernelWithUnion))
 	suite.Run(t, new(FilterKernelExtension))