github.com/apache/arrow/go/v14@v14.0.1/arrow/compute/vector_run_ends.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one
     2  // or more contributor license agreements.  See the NOTICE file
     3  // distributed with this work for additional information
     4  // regarding copyright ownership.  The ASF licenses this file
     5  // to you under the Apache License, Version 2.0 (the
     6  // "License"); you may not use this file except in compliance
     7  // with the License.  You may obtain a copy of the License at
     8  //
     9  // http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing, software
    12  // distributed under the License is distributed on an "AS IS" BASIS,
    13  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  // See the License for the specific language governing permissions and
    15  // limitations under the License.
    16  
    17  //go:build go1.18
    18  
    19  package compute
    20  
    21  import (
    22  	"context"
    23  
    24  	"github.com/apache/arrow/go/v14/arrow"
    25  	"github.com/apache/arrow/go/v14/arrow/compute/internal/kernels"
    26  )
    27  
    28  var (
    29  	runEndEncodeDoc = FunctionDoc{
    30  		Summary:         "Run-end encode array",
    31  		Description:     "Return a run-end encoded version of the input array",
    32  		ArgNames:        []string{"array"},
    33  		OptionsType:     "RunEndEncodeOptions",
    34  		OptionsRequired: true,
    35  	}
    36  	runEndDecodeDoc = FunctionDoc{
    37  		Summary:     "Decode run-end encoded array",
    38  		Description: "Return a decoded version of a run-end encoded input array",
    39  		ArgNames:    []string{"array"},
    40  	}
    41  )
    42  
    43  type RunEndEncodeOptions = kernels.RunEndEncodeState
    44  
    45  func RegisterVectorRunEndFuncs(reg FunctionRegistry) {
    46  	encKns, decKns := kernels.GetRunEndEncodeKernels()
    47  	encFn := NewVectorFunction("run_end_encode", Unary(), runEndEncodeDoc)
    48  	for _, k := range encKns {
    49  		if err := encFn.AddKernel(k); err != nil {
    50  			panic(err)
    51  		}
    52  	}
    53  	reg.AddFunction(encFn, false)
    54  
    55  	decFn := NewVectorFunction("run_end_decode", Unary(), runEndDecodeDoc)
    56  	for _, k := range decKns {
    57  		if err := decFn.AddKernel(k); err != nil {
    58  			panic(err)
    59  		}
    60  	}
    61  	reg.AddFunction(decFn, false)
    62  }
    63  
    64  func RunEndEncode(ctx context.Context, opts RunEndEncodeOptions, arg Datum) (Datum, error) {
    65  	return CallFunction(ctx, "run_end_encode", &opts, arg)
    66  }
    67  
    68  func RunEndEncodeArray(ctx context.Context, opts RunEndEncodeOptions, input arrow.Array) (arrow.Array, error) {
    69  	out, err := RunEndEncode(ctx, opts, &ArrayDatum{Value: input.Data()})
    70  	if err != nil {
    71  		return nil, err
    72  	}
    73  	defer out.Release()
    74  
    75  	return out.(*ArrayDatum).MakeArray(), nil
    76  }
    77  
    78  func RunEndDecode(ctx context.Context, arg Datum) (Datum, error) {
    79  	return CallFunction(ctx, "run_end_decode", nil, arg)
    80  }
    81  
    82  func RunEndDecodeArray(ctx context.Context, input arrow.Array) (arrow.Array, error) {
    83  	out, err := RunEndDecode(ctx, &ArrayDatum{Value: input.Data()})
    84  	if err != nil {
    85  		return nil, err
    86  	}
    87  	defer out.Release()
    88  
    89  	return out.(*ArrayDatum).MakeArray(), nil
    90  }