Skip to content

Commit

Permalink
GH-36936: [Go] Make it possible to register custom functions. (#36959)
Browse files Browse the repository at this point in the history
### Rationale for this change

As discussed in the issue #36936 , this change makes it possible for the consumers of the module to register custom functions, while taking advantage of the existing infrastructure for matching types and choosing kernels.

### What changes are included in this PR?

Just moving the package one level higher, so that it's module-exported.

### Are these changes tested?

There's no actual code changes other than moving files around and updating imports - existing tests still pass.

### Are there any user-facing changes?

The `compute/exec` is now exported.

* Closes: #36936

Authored-by: Jakub Martin <[email protected]>
Signed-off-by: Matt Topol <[email protected]>
  • Loading branch information
cube2222 authored Aug 7, 2023
1 parent 7cbbd3e commit 3c00b08
Show file tree
Hide file tree
Showing 46 changed files with 132 additions and 41 deletions.
2 changes: 1 addition & 1 deletion go/arrow/compute/arithmetic.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import (
"fmt"

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/compute/internal/kernels"
"github.com/apache/arrow/go/v13/arrow/decimal128"
"github.com/apache/arrow/go/v13/arrow/decimal256"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/arithmetic_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ import (
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/compute"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/compute/internal/kernels"
"github.com/apache/arrow/go/v13/arrow/decimal128"
"github.com/apache/arrow/go/v13/arrow/decimal256"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/cast.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/compute/internal/kernels"
)

Expand Down
91 changes: 91 additions & 0 deletions go/arrow/compute/example_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
// 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"
"log"

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/compute"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/memory"
)

// This example demonstrates how to register a custom scalar function.
func Example_customFunction() {
pool := memory.NewGoAllocator()

ctx := context.Background()
execCtx := compute.DefaultExecCtx()
ctx = compute.SetExecCtx(ctx, execCtx)

add42 := compute.NewScalarFunction("add_42", compute.Arity{
NArgs: 1,
}, compute.FunctionDoc{
Summary: "Returns the input values plus 42",
ArgNames: []string{"input"},
})

if err := add42.AddNewKernel(
[]exec.InputType{
// We accept a single argument (array) of Int8 type.
{
Kind: exec.InputExact,
Type: arrow.PrimitiveTypes.Int8,
},
},
// We'll return a single Int8 array.
exec.NewOutputType(arrow.PrimitiveTypes.Int8),
func(ctx *exec.KernelCtx, span *exec.ExecSpan, result *exec.ExecResult) error {
// The second buffer contains the values. Both for the input and the output arrays.
for i, x := range span.Values[0].Array.Buffers[1].Buf {
result.Buffers[1].Buf[i] = x + 42
}
return nil
},
nil,
); err != nil {
log.Fatal(err)
}
execCtx.Registry.AddFunction(add42, true)

inputArrayBuilder := array.NewInt8Builder(pool)
for i := 0; i < 16; i++ {
inputArrayBuilder.Append(int8(i))
}
inputArray := inputArrayBuilder.NewArray()

outputArrayDatum, err := compute.CallFunction(
compute.SetExecCtx(context.Background(), execCtx),
"add_42",
nil,
&compute.ArrayDatum{Value: inputArray.Data()},
)
if err != nil {
log.Fatal(err)
}

fmt.Println(array.NewInt8Data(outputArrayDatum.(*compute.ArrayDatum).Value).Int8Values())

// Output:
// [42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57]
}
2 changes: 1 addition & 1 deletion go/arrow/compute/exec.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import (
"fmt"

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
)

Expand Down
File renamed without changes.
File renamed without changes.
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/compute"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/memory"
"github.com/apache/arrow/go/v13/arrow/scalar"
"github.com/stretchr/testify/assert"
Expand Down
File renamed without changes.
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/decimal128"
"github.com/apache/arrow/go/v13/arrow/endian"
"github.com/apache/arrow/go/v13/arrow/memory"
Expand Down
File renamed without changes.
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/memory"
"github.com/stretchr/testify/assert"
)
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/exec_internals_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ import (
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/internal/testing/gen"
"github.com/apache/arrow/go/v13/arrow/memory"
"github.com/apache/arrow/go/v13/arrow/scalar"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/exec_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
"github.com/apache/arrow/go/v13/arrow/scalar"
"github.com/stretchr/testify/suite"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/executor.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ import (
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/internal"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
"github.com/apache/arrow/go/v13/arrow/memory"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/expression.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/compute/internal/kernels"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
"github.com/apache/arrow/go/v13/arrow/ipc"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/exprs/exec.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/compute"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/decimal128"
"github.com/apache/arrow/go/v13/arrow/endian"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/functions.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (
"strings"

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
)

type Function interface {
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/base_arithmetic.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (

"github.com/JohnCGriffin/overflow"
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/decimal128"
"github.com/apache/arrow/go/v13/arrow/decimal256"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/base_arithmetic_amd64.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ package kernels
import (
"unsafe"

"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
"golang.org/x/exp/constraints"
"golang.org/x/sys/cpu"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
package kernels

import (
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"golang.org/x/exp/constraints"
)

Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/boolean_cast.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
)

func isNonZero[T exec.FixedWidthTypes](ctx *exec.KernelCtx, in []T, out []byte) error {
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/cast.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ package kernels
import (
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
)

type CastOptions struct {
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/cast_temporal.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
)

Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/helpers.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
"github.com/apache/arrow/go/v13/arrow/memory"
"github.com/apache/arrow/go/v13/arrow/scalar"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/numeric_cast.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/decimal128"
"github.com/apache/arrow/go/v13/arrow/decimal256"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/rounding.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import (
"math"

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/decimal128"
"github.com/apache/arrow/go/v13/arrow/decimal256"
"github.com/apache/arrow/go/v13/arrow/scalar"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/scalar_arithmetic.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/decimal128"
"github.com/apache/arrow/go/v13/arrow/decimal256"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/scalar_boolean.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ package kernels

import (
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/scalar"
)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ import (
"unsafe"

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"golang.org/x/sys/cpu"
)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@

package kernels

import "github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
import "github.com/apache/arrow/go/v13/arrow/compute/exec"

func genCompareKernel[T exec.NumericTypes](op CompareOperator) *CompareData {
return genGoCompareKernel(getCmpOp[T](op))
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/scalar_comparisons.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/decimal128"
"github.com/apache/arrow/go/v13/arrow/decimal256"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/string_casts.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/float16"
"github.com/apache/arrow/go/v13/internal/bitutils"
)
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ import (
"fmt"

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
"github.com/apache/arrow/go/v13/arrow/scalar"
)
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/vector_hash.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
"github.com/apache/arrow/go/v13/arrow/memory"
"github.com/apache/arrow/go/v13/internal/bitutils"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/vector_run_end_encode.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/decimal128"
"github.com/apache/arrow/go/v13/arrow/decimal256"
"github.com/apache/arrow/go/v13/arrow/float16"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/internal/kernels/vector_selection.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/bitutil"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/internal/debug"
"github.com/apache/arrow/go/v13/arrow/memory"
"github.com/apache/arrow/go/v13/internal/bitutils"
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/registry_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/compute"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/stretchr/testify/assert"
"golang.org/x/exp/slices"
)
Expand Down
2 changes: 1 addition & 1 deletion go/arrow/compute/scalar_bool.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ import (
"fmt"

"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/compute/internal/exec"
"github.com/apache/arrow/go/v13/arrow/compute/exec"
"github.com/apache/arrow/go/v13/arrow/compute/internal/kernels"
)

Expand Down
Loading

0 comments on commit 3c00b08

Please sign in to comment.