zeroshade commented on code in PR #1138: URL: https://github.com/apache/arrow-go/pull/1138#discussion_r3857268576
########## arrow/compute/scalar_nested.go: ########## @@ -0,0 +1,172 @@ +// 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" + "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" + "github.com/apache/arrow-go/v18/arrow/scalar" +) + +var listElementDoc = FunctionDoc{ + Summary: "Compute elements using nested list values and an index", + Description: "For each list value, return the element at the requested index.\n" + + "Use an integral scalar for broadcast selection. A one-element array-like\n" + + "index is accepted only for a matching one-element execution and is not\n" + + "broadcast over longer input.", + ArgNames: []string{"lists", "index"}, +} + +func validateListElementIndex(index Datum) error { + switch index.Kind() { + case KindScalar: + return kernels.ValidateListElementScalarIndex(index.(*ScalarDatum).Value) + case KindArray, KindChunked: + if index.Len() == 0 { + return fmt.Errorf("%w: list_element index array is empty", arrow.ErrInvalid) + } + if index.Len() > 1 { + return fmt.Errorf("%w: list_element does not support arrays of list indices", arrow.ErrNotImplemented) + } + } + return nil +} + +type listElementFunction struct { + ScalarFunction +} + +func listElementScalarResultSupported(typ arrow.DataType) bool { + switch typ.ID() { + case arrow.BINARY_VIEW, arrow.STRING_VIEW, arrow.LIST_VIEW, arrow.LARGE_LIST_VIEW: + return false + case arrow.EXTENSION: + return listElementScalarResultSupported(typ.(arrow.ExtensionType).StorageType()) + case arrow.STRUCT: + for _, field := range typ.(*arrow.StructType).Fields() { + if !listElementScalarResultSupported(field.Type) { + return false + } + } + case arrow.SPARSE_UNION: + for _, field := range typ.(arrow.UnionType).Fields() { + if !listElementScalarResultSupported(field.Type) { + return false + } + } + case arrow.DENSE_UNION: + return true + case arrow.RUN_END_ENCODED: + return listElementScalarResultSupported(typ.(*arrow.RunEndEncodedType).Encoded()) + } + return true +} + +func listElementScalarArrayValueSupported(values arrow.Array, index int) bool { + if values.DataType().ID() != arrow.DICTIONARY && values.IsNull(index) { + return true + } + + switch values := values.(type) { + case *array.BinaryView, *array.StringView, *array.ListView, *array.LargeListView: + return false + case array.ExtensionArray: + return listElementScalarArrayValueSupported(values.Storage(), index) + case *array.Struct: + for i := 0; i < values.NumField(); i++ { + if !listElementScalarArrayValueSupported(values.Field(i), index) { + return false + } + } + case *array.SparseUnion: + for i := 0; i < values.NumFields(); i++ { + if !listElementScalarArrayValueSupported(values.Field(i), index) { + return false + } + } + case *array.DenseUnion: + child := values.Field(values.ChildID(index)) + if child == nil { + return false + } + offset := values.ValueOffset(index) + if offset < 0 || int64(offset) >= int64(child.Len()) { + return false + } + return listElementScalarArrayValueSupported(child, int(offset)) + case *array.RunEndEncoded: + return listElementScalarArrayValueSupported(values.Values(), values.GetPhysicalIndex(index)) + } + return true +} + +// Validate the index before scalar execution splits arguments into spans. The +// index contract is defined by the original Datum, not by the length of an +// execution span. +func (fn *listElementFunction) Execute(ctx context.Context, opts FunctionOptions, args ...Datum) (Datum, error) { + if err := fn.checkArity(len(args)); err != nil { + return nil, err + } + if err := checkOptions(fn, opts); err != nil { + return nil, err + } + + if err := validateListElementIndex(args[1]); err != nil { + return nil, err + } + if args[0].Kind() == KindScalar && args[1].Kind() == KindScalar { + if listType, ok := args[0].(ArrayLikeDatum).Type().(arrow.ListLikeType); ok && + !listElementScalarResultSupported(listType.Elem()) { + return nil, fmt.Errorf("%w: list_element scalar output type %s is not supported", arrow.ErrNotImplemented, listType.Elem()) + } + + listValue := args[0].(*ScalarDatum).Value.(scalar.ListScalar) Review Comment: **Blocking:** This unchecked assertion makes unsupported scalar inputs panic instead of returning the normal dispatch/type error. For example, `ListElement(ctx, ScalarDatum(Int32(7)), ScalarDatum(Int64(0)))` panics with `*scalar.Int32 is not scalar.ListScalar`; the equivalent unsupported array input returns an error. Please guard the list-like/type assertion and either fall through to normal kernel dispatch or return `arrow.ErrType`, with a regression test for a non-list scalar argument. -- 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]
