fix: compute fixes for extension types

apache · Oct 24, 2024 · c779c38 · c779c38
1 parent 56b794f
commit c779c38
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 4 deletions.
diff --git a/arrow/compute/exprs/exec.go b/arrow/compute/exprs/exec.go
@@ -571,6 +571,25 @@ func executeScalarBatch(ctx context.Context, input compute.ExecBatch, exp expr.E
 			return nil, err
 		}
 
+		var newArgs []compute.Datum
+		// cast arguments if necessary
+		for i, arg := range args {
+			if !arrow.TypeEqual(argTypes[i], arg.(compute.ArrayLikeDatum).Type()) {
+				if newArgs == nil {
+					newArgs = make([]compute.Datum, len(args))
+					copy(newArgs, args)
+				}
+				newArgs[i], err = compute.CastDatum(ctx, arg, compute.SafeCastOptions(argTypes[i]))
+				if err != nil {
+					return nil, err
+				}
+				defer newArgs[i].Release()
+			}
+		}
+		if newArgs != nil {
+			args = newArgs
+		}
+
 		kctx := &exec.KernelCtx{Ctx: ctx, Kernel: k}
 		init := k.GetInitFn()
 		kinitArgs := exec.KernelInitArgs{Kernel: k, Inputs: argTypes, Options: opts}
@@ -613,7 +632,7 @@ func executeScalarBatch(ctx context.Context, input compute.ExecBatch, exp expr.E
 			result.Release()
 		}
 
-		return result, nil
+		return result, err
 	}
 
 	return nil, arrow.ErrNotImplemented

diff --git a/arrow/compute/exprs/extension_types.go b/arrow/compute/exprs/extension_types.go
@@ -75,7 +75,7 @@ func (ef *simpleExtensionTypeFactory[P]) ExtensionEquals(other arrow.ExtensionTy
 	return ef.params == rhs.params
 }
 func (ef *simpleExtensionTypeFactory[P]) ArrayType() reflect.Type {
-	return reflect.TypeOf(array.ExtensionArrayBase{})
+	return reflect.TypeOf(simpleExtensionArrayFactory[P]{})
 }
 
 func (ef *simpleExtensionTypeFactory[P]) CreateType(params P) arrow.DataType {
@@ -91,10 +91,14 @@ func (ef *simpleExtensionTypeFactory[P]) CreateType(params P) arrow.DataType {
 	}
 }
 
+type simpleExtensionArrayFactory[P comparable] struct {
+	array.ExtensionArrayBase
+}
+
 type uuidExtParams struct{}
 
 var uuidType = simpleExtensionTypeFactory[uuidExtParams]{
-	name: "uuid", getStorage: func(uuidExtParams) arrow.DataType {
+	name: "arrow.uuid", getStorage: func(uuidExtParams) arrow.DataType {
 		return &arrow.FixedSizeBinaryType{ByteWidth: 16}
 	}}
 

diff --git a/arrow/compute/scalar_compare.go b/arrow/compute/scalar_compare.go
@@ -52,6 +52,7 @@ func (fn *compareFunction) DispatchBest(vals ...arrow.DataType) (exec.Kernel, er
 	}
 
 	ensureDictionaryDecoded(vals...)
+	ensureNotExtensionType(vals...)
 	replaceNullWithOtherType(vals...)
 
 	if dt := commonNumeric(vals...); dt != nil {

diff --git a/arrow/compute/utils.go b/arrow/compute/utils.go
@@ -105,6 +105,14 @@ func ensureDictionaryDecoded(vals ...arrow.DataType) {
 	}
 }
 
+func ensureNotExtensionType(vals ...arrow.DataType) {
+	for i, v := range vals {
+		if v.ID() == arrow.EXTENSION {
+			vals[i] = v.(arrow.ExtensionType).StorageType()
+		}
+	}
+}
+
 func replaceNullWithOtherType(vals ...arrow.DataType) {
 	debug.Assert(len(vals) == 2, "should be length 2")
 

diff --git a/arrow/extensions/uuid.go b/arrow/extensions/uuid.go
@@ -228,6 +228,10 @@ func (*UUIDType) ExtensionName() string {
 	return "arrow.uuid"
 }
 
+func (*UUIDType) Bytes() int { return 16 }
+
+func (*UUIDType) BitWidth() int { return 128 }
+
 func (e *UUIDType) String() string {
 	return fmt.Sprintf("extension<%s>", e.ExtensionName())
 }

diff --git a/parquet/pqarrow/schema.go b/parquet/pqarrow/schema.go
@@ -24,6 +24,7 @@ import (
 
 	"github.com/apache/arrow-go/v18/arrow"
 	"github.com/apache/arrow-go/v18/arrow/decimal128"
+	"github.com/apache/arrow-go/v18/arrow/extensions"
 	"github.com/apache/arrow-go/v18/arrow/flight"
 	"github.com/apache/arrow-go/v18/arrow/memory"
 	"github.com/apache/arrow-go/v18/parquet"
@@ -514,8 +515,10 @@ func arrowFromFLBA(logical schema.LogicalType, length int) (arrow.DataType, erro
 	switch logtype := logical.(type) {
 	case schema.DecimalLogicalType:
 		return arrowDecimal(logtype), nil
-	case schema.NoLogicalType, schema.IntervalLogicalType, schema.UUIDLogicalType:
+	case schema.NoLogicalType, schema.IntervalLogicalType:
 		return &arrow.FixedSizeBinaryType{ByteWidth: int(length)}, nil
+	case schema.UUIDLogicalType:
+		return extensions.NewUUIDType(), nil
 	case schema.Float16LogicalType:
 		return &arrow.Float16Type{}, nil
 	default: