lidavidm commented on code in PR #1474:
URL: https://github.com/apache/arrow-go/pull/1474#discussion_r4488214223
##########
arrow/compute/exprs/field_refs.go:
##########
@@ -86,172 +89,271 @@ func GetRefSchema(ref expr.ReferenceSegment, schema
*arrow.Schema) (*arrow.Field
// ListElementRef can only reference against a List or LargeList scalar,
// and a MapKeyRef will only reference against a Map scalar. An error is
// returned if following the reference children ends up with an invalid
-// nested reference object.
+// nested reference object. A field of a null struct is null.
+//
+// The result may share memory with s, and it is not retained; use
Review Comment:
If it _may_ share memory, how is the caller supposed to know whether to
release or not release the result?
##########
arrow/compute/exprs/field_refs.go:
##########
@@ -86,172 +89,271 @@ func GetRefSchema(ref expr.ReferenceSegment, schema
*arrow.Schema) (*arrow.Field
// ListElementRef can only reference against a List or LargeList scalar,
// and a MapKeyRef will only reference against a Map scalar. An error is
// returned if following the reference children ends up with an invalid
-// nested reference object.
+// nested reference object. A field of a null struct is null.
+//
+// The result may share memory with s, and it is not retained; use
+// GetReferencedValue for a result that the caller owns.
func GetScalar(ref expr.ReferenceSegment, s scalar.Scalar, mem
memory.Allocator, ext ExtensionIDSet) (scalar.Scalar, error) {
if ref == nil {
return nil, compute.ErrEmpty
}
- var out scalar.Scalar
- for ref != nil {
- switch f := ref.(type) {
- case *expr.StructFieldRef:
- if s.DataType().ID() != arrow.STRUCT {
- return nil, fmt.Errorf("%w: attempting to
reference field from non-struct scalar %s",
- arrow.ErrInvalid, s)
- }
+ for ; ref != nil; ref = ref.GetChild() {
+ next, _, err := scalarStep(ref, s, mem, ext)
+ if err != nil {
+ return nil, err
+ }
+ s = next
+ }
+ return s, nil
+}
- st := s.(*scalar.Struct)
- if f.Field < 0 || f.Field >= int32(len(st.Value)) {
- return nil, fmt.Errorf("%w: indices=%s",
compute.ErrIndexRange, ref)
- }
+// scalarStep follows one reference segment from s. borrowed reports whether
+// next is part of s, rather than a new scalar the caller owns.
+func scalarStep(ref expr.ReferenceSegment, s scalar.Scalar, mem
memory.Allocator, ext ExtensionIDSet) (next scalar.Scalar, borrowed bool, err
error) {
+ switch f := ref.(type) {
+ case *expr.StructFieldRef:
+ if s.DataType().ID() != arrow.STRUCT {
+ return nil, false, fmt.Errorf("%w: attempting to
reference field from non-struct scalar %s",
+ arrow.ErrInvalid, s)
+ }
- out = st.Value[f.Field]
- case *expr.ListElementRef:
- switch v := s.(type) {
- case *scalar.List:
- sc, err := scalar.GetScalar(v.Value,
int(f.Offset))
- if err != nil {
- return nil, err
- }
- out = sc
- case *scalar.LargeList:
- sc, err := scalar.GetScalar(v.Value,
int(f.Offset))
- if err != nil {
- return nil, err
- }
- out = sc
- default:
- return nil, fmt.Errorf("%w: cannot get
ListElementRef from non-list scalar %s",
- arrow.ErrInvalid, v)
- }
- case *expr.MapKeyRef:
- v, ok := s.(*scalar.Map)
- if !ok {
- return nil, arrow.ErrInvalid
- }
+ st := s.(*scalar.Struct)
+ fields := st.Type.(*arrow.StructType).Fields()
+ if f.Field < 0 || f.Field >= int32(len(fields)) {
+ return nil, false, fmt.Errorf("%w: indices=%s",
compute.ErrIndexRange, ref)
+ }
+ if !st.IsValid() {
+ return scalar.MakeNullScalar(fields[f.Field].Type),
false, nil
+ }
+ return st.Value[f.Field], true, nil
+ case *expr.ListElementRef:
+ var values arrow.Array
+ switch v := s.(type) {
+ case *scalar.List:
+ values = v.Value
+ case *scalar.LargeList:
+ values = v.Value
+ default:
+ return nil, false, fmt.Errorf("%w: cannot get
ListElementRef from non-list scalar %s",
+ arrow.ErrInvalid, s)
+ }
+ if !s.IsValid() {
+ return
scalar.MakeNullScalar(s.DataType().(arrow.ListLikeType).Elem()), false, nil
+ }
+ sc, err := scalar.GetScalar(values, int(f.Offset))
+ return sc, false, err
+ case *expr.MapKeyRef:
+ v, ok := s.(*scalar.Map)
+ if !ok {
+ return nil, false, arrow.ErrInvalid
+ }
- dt, _, err := FromSubstraitType(f.MapKey.GetType(), ext)
- if err != nil {
- return nil, err
- }
+ dt, _, err := FromSubstraitType(f.MapKey.GetType(), ext)
+ if err != nil {
+ return nil, false, err
+ }
- if !arrow.TypeEqual(dt,
v.Type.(*arrow.MapType).KeyType()) {
- return nil, arrow.ErrInvalid
- }
+ mapType := v.Type.(*arrow.MapType)
+ if !arrow.TypeEqual(dt, mapType.KeyType()) {
+ return nil, false, arrow.ErrInvalid
+ }
+ if !v.IsValid() {
+ return scalar.MakeNullScalar(mapType.ItemType()),
false, nil
+ }
- keyvalDatum, err := literalToDatum(mem, f.MapKey, ext)
+ keyvalDatum, err := literalToDatum(mem, f.MapKey, ext)
+ if err != nil {
+ return nil, false, err
+ }
+ defer keyvalDatum.Release()
+
+ var (
+ keyval = keyvalDatum.(*compute.ScalarDatum)
+ m = v.Value.(*array.Struct)
+ keys = m.Field(0)
+ )
+ for i := 0; i < v.Value.Len(); i++ {
+ kv, err := scalar.GetScalar(keys, i)
if err != nil {
- return nil, err
+ return nil, false, err
}
-
- var (
- keyval = keyvalDatum.(*compute.ScalarDatum)
- m = v.Value.(*array.Struct)
- keys = m.Field(0)
- valueScalar scalar.Scalar
- )
- for i := 0; i < v.Value.Len(); i++ {
- kv, err := scalar.GetScalar(keys, i)
- if err != nil {
- return nil, err
- }
- if scalar.Equals(kv, keyval.Value) {
- valueScalar, err =
scalar.GetScalar(m.Field(1), i)
- if err != nil {
- return nil, err
- }
- break
- }
+ found := scalar.Equals(kv, keyval.Value)
+ releaseScalar(kv)
+ if found {
+ sc, err := scalar.GetScalar(m.Field(1), i)
+ return sc, false, err
}
-
- if valueScalar == nil {
- return nil, arrow.ErrNotFound
- }
-
- out = valueScalar
}
- s = out
- ref = ref.GetChild()
+ return nil, false, arrow.ErrNotFound
+ }
+ return nil, false, arrow.ErrNotImplemented
+}
+
+func retainScalar(s scalar.Scalar) {
+ if r, ok := s.(scalar.Releasable); ok {
+ r.Retain()
}
+}
- return out, nil
+func releaseScalar(s scalar.Scalar) {
+ if r, ok := s.(scalar.Releasable); ok {
+ r.Release()
+ }
}
// GetReferencedValue retrieves the referenced (potentially nested) value from
-// the provided datum which may be a scalar, array, or record batch.
+// the provided datum which may be a scalar, array, or record batch. The
+// result is a new reference that the caller must release.
+//
+// A field of a struct array is null wherever the struct is null, and a field
+// of a null struct scalar is null.
Review Comment:
I think this follow naturally from Arrow semantics, and isn't on the level
of needing to be called out in the API documentation. After all, we don't
explain this for any other type.
--
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]
