feat(extensions): resolve UDT URNs while parsing

expr/expressions_test.go

@@ -380,7 +380,7 @@ func TestRoundTripUsingTestData(t *testing.T) {
 			assert.True(t, e.Equals(e))
 
 			if typTest, ok := test["type"].(string); ok {
-				exp, err := parser.ParseType(typTest)
+				exp, err := parser.ParseType(typTest, nil)
 				require.NoError(t, err)
 
 				assert.Equal(t, exp.String(), e.GetType().String())

extensions/simple_extension.go

@@ -4,9 +4,9 @@ package extensions
 
 import (
 	"fmt"
-	"reflect"
 	"strings"
 
+	"github.com/goccy/go-yaml"
 	substraitgo "github.com/substrait-io/substrait-go/v8"
 	"github.com/substrait-io/substrait-go/v8/types"
 	"github.com/substrait-io/substrait-go/v8/types/parser"
@@ -114,13 +114,17 @@ func (v TypeArg) GetTypeExpression() types.FuncDefArgType {
 
 type FuncParameterList []FuncParameter
 
-func (a *FuncParameterList) UnmarshalYAML(fn func(interface{}) error) error {
+func (a *FuncParameterList) UnmarshalYAML(data []byte) error {
+	return parseFuncParameterList(data, parser.TypeParser{}, a)
+}
+
+func parseFuncParameterList(data []byte, typeParser parser.TypeParser, out *FuncParameterList) error {
 	var args []map[string]any
-	if err := fn(&args); err != nil {
+	if err := yaml.Unmarshal(data, &args); err != nil {
 		return err
 	}
 
-	*a = make(FuncParameterList, len(args))
+	*out = make(FuncParameterList, len(args))
 	for i, arg := range args {
 		var (
 			name, desc string
@@ -138,7 +142,7 @@ func (a *FuncParameterList) UnmarshalYAML(fn func(interface{}) error) error {
 			for j, v := range vals {
 				values[j] = v.(string)
 			}
-			(*a)[i] = EnumArg{
+			(*out)[i] = EnumArg{
 				Name:        name,
 				Description: desc,
 				Options:     values,
@@ -149,45 +153,29 @@ func (a *FuncParameterList) UnmarshalYAML(fn func(interface{}) error) error {
 				constant = c.(bool)
 			}
 
-			arg := ValueArg{
-				Name:        name,
-				Description: desc,
-				Value:       new(parser.TypeExpression),
-				Constant:    constant,
-			}
-			err := arg.Value.UnmarshalYAML(func(v any) error {
-				rv := reflect.ValueOf(v)
-				if rv.Type().Kind() != reflect.Ptr {
-					return substraitgo.ErrInvalidType
-				}
-				rv.Elem().Set(reflect.ValueOf(val))
-				return nil
-			})
+			valueType, err := parseTypeExpressionFromYAMLValue(typeParser, val)
 			if err != nil {
 				return fmt.Errorf("failure reading YAML %v", err)
 			}
 
-			(*a)[i] = arg
-
-		} else if typ, ok := arg["type"]; ok {
-			arg := TypeArg{
+			(*out)[i] = ValueArg{
 				Name:        name,
 				Description: desc,
-				Type:        new(parser.TypeExpression),
+				Value:       &valueType,
+				Constant:    constant,
 			}
-			err := arg.Type.UnmarshalYAML(func(v any) error {
-				rv := reflect.ValueOf(v)
-				if rv.Type().Kind() != reflect.Ptr {
-					return substraitgo.ErrInvalidType
-				}
-				rv.Elem().Set(reflect.ValueOf(typ))
-				return nil
-			})
+
+		} else if typ, ok := arg["type"]; ok {
+			typeExpr, err := parseTypeExpressionFromYAMLValue(typeParser, typ)
 			if err != nil {
 				return fmt.Errorf("failure reading YAML %v", err)
 			}
 
-			(*a)[i] = arg
+			(*out)[i] = TypeArg{
+				Name:        name,
+				Description: desc,
+				Type:        &typeExpr,
+			}
 		}
 	}
 
@@ -501,3 +489,77 @@ type SimpleExtensionFile struct {
 	AggregateFunctions []AggregateFunction `yaml:"aggregate_functions,omitempty"`
 	WindowFunctions    []WindowFunction    `yaml:"window_functions,omitempty"`
 }
+
+func (s *SimpleExtensionFile) UnmarshalYAML(data []byte) error {
+	type typeDeclarations struct {
+		Urn   string `yaml:"urn"`
+		Types []Type `yaml:"types,omitempty"`
+	}
+	var declarations typeDeclarations
+	if err := yaml.Unmarshal(data, &declarations); err != nil {
+		return err
+	}
+
+	declaredTypes := make(map[string]struct{}, len(declarations.Types))
+	for _, typ := range declarations.Types {
+		declaredTypes[typ.Name] = struct{}{}
+	}
+
+	typeParser := parser.TypeParser{
+		ResolveUserDefinedType: func(name string, nullability types.Nullability, parameters []types.UDTParameter) (*types.ParameterizedUserDefinedType, error) {
+			if _, ok := declaredTypes[name]; !ok {
+				return nil, fmt.Errorf("%w: user-defined type %q is not declared", substraitgo.ErrInvalidSimpleExtention, name)
+			}
+			return &types.ParameterizedUserDefinedType{
+				Name:           name,
+				URN:            declarations.Urn,
+				Nullability:    nullability,
+				TypeParameters: parameters,
+			}, nil
+		},
+	}
+
+	type rawFile SimpleExtensionFile
+	var raw rawFile
+	if err := yaml.UnmarshalWithOptions(
+		data,
+		&raw,
+		yaml.CustomUnmarshaler[parser.TypeExpression](func(out *parser.TypeExpression, data []byte) error {
+			parsed, err := parseTypeExpression(data, typeParser)
+			if err != nil {
+				return err
+			}
+			*out = parsed
+			return nil
+		}),
+		yaml.CustomUnmarshaler[FuncParameterList](func(out *FuncParameterList, data []byte) error {
+			return parseFuncParameterList(data, typeParser, out)
+		}),
+	); err != nil {
+		return err
+	}
+
+	*s = SimpleExtensionFile(raw)
+	return nil
+}
+
+func parseTypeExpressionFromYAMLValue(typeParser parser.TypeParser, value any) (parser.TypeExpression, error) {
+	typeString, ok := value.(string)
+	if !ok {
+		return parser.TypeExpression{}, substraitgo.ErrInvalidType
+	}
+	return parseTypeExpression([]byte(typeString), typeParser)
+}
+
+func parseTypeExpression(data []byte, typeParser parser.TypeParser) (parser.TypeExpression, error) {
+	var typeString string
+	if err := yaml.Unmarshal(data, &typeString); err != nil {
+		return parser.TypeExpression{}, err
+	}
+
+	typ, err := typeParser.Parse(typeString)
+	if err != nil {
+		return parser.TypeExpression{}, err
+	}
+	return parser.TypeExpression{ValueType: typ}, nil
+}

extensions/simple_extension_test.go

@@ -47,6 +47,9 @@ types:
 func TestUnmarshalCustomScalarFunction(t *testing.T) {
 	const customDef = `
 urn: extension:test:test_ext
+types:
+  - name: customtype1
+  - name: customtype2
 scalar_functions:
   - name: "scalar1"
     impls:
@@ -85,6 +88,47 @@ scalar_functions:
 	assert.Equal(t, proto.Type_NULLABILITY_NULLABLE, typ.GetNullability(), "expected NULLABILITY_NULLABLE")
 }
 
+func TestUnmarshalSimpleExtensionRejectsUndeclaredUserDefinedType(t *testing.T) {
+	const customDef = `
+urn: extension:test:test_ext
+scalar_functions:
+  - name: "scalar1"
+    impls:
+      - args:
+          - name: arg1
+            value: u!missing_type
+        return: i64
+`
+
+	var f extensions.SimpleExtensionFile
+	err := yaml.Unmarshal([]byte(customDef), &f)
+
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), `user-defined type "missing_type" is not declared`)
+}
+
+func TestUnmarshalSimpleExtensionResolvesUserDefinedTypeURN(t *testing.T) {
+	const customDef = `
+urn: extension:test:test_ext
+types:
+  - name: customtype
+scalar_functions:
+  - name: "scalar1"
+    impls:
+      - args:
+          - name: arg1
+            value: i64
+        return: u!customtype
+`
+
+	var f extensions.SimpleExtensionFile
+	require.NoError(t, yaml.Unmarshal([]byte(customDef), &f))
+
+	udt, ok := f.ScalarFunctions[0].Impls[0].Return.ValueType.(*types.ParameterizedUserDefinedType)
+	require.True(t, ok)
+	assert.Equal(t, "extension:test:test_ext", udt.URN)
+}
+
 func TestUnmarshalSimpleExtensionScalarFunction(t *testing.T) {
 	const addDef = `
 urn: extension:test:test_ext

extensions/variants.go

@@ -148,7 +148,11 @@ func EvaluateTypeExpression(urn string, nullHandling NullabilityHandling, return
 	// For other types like AnyType, the TypeReference is already correctly set.
 	if udt, ok := outType.(*types.UserDefinedType); ok {
 		if paramUDT, ok := returnTypeExpr.(*types.ParameterizedUserDefinedType); ok {
-			udt.TypeReference = registry.GetTypeAnchor(ID{Name: paramUDT.Name, URN: urn})
+			udtURN := paramUDT.URN
+			if udtURN == "" {
+				udtURN = urn
+			}
+			udt.TypeReference = registry.GetTypeAnchor(ID{Name: paramUDT.Name, URN: udtURN})
 		}
 	}
 
@@ -277,7 +281,7 @@ func parseFuncName(compoundName string) (name string, args FuncParameterList) {
 	}
 	splitArgs := strings.Split(argsStr, "_")
 	for _, argStr := range splitArgs {
-		parsed, err := parser.ParseType(argStr)
+		parsed, err := parser.ParseType(argStr, nil)
 		if err != nil {
 			panic(err)
 		}
@@ -422,7 +426,7 @@ func NewAggFuncVariantOpts(id ID, opts AggVariantOptions) *AggregateFunctionVari
 				substraitgo.ErrInvalidExpr, id))
 		}
 
-		intermediate, err := parser.ParseType(opts.IntermediateOutputType)
+		intermediate, err := parser.ParseType(opts.IntermediateOutputType, nil)
 		if err != nil {
 			panic(err)
 		}
@@ -549,7 +553,7 @@ func NewWindowFuncVariantOpts(id ID, opts WindowVariantOpts) *WindowFunctionVari
 				substraitgo.ErrInvalidExpr, id))
 		}
 
-		intermediate, err := parser.ParseType(opts.IntermediateOutputType)
+		intermediate, err := parser.ParseType(opts.IntermediateOutputType, nil)
 		if err != nil {
 			panic(err)
 		}

extensions/variants_test.go

@@ -20,12 +20,12 @@ import (
 func TestEvaluateTypeExpression(t *testing.T) {
 	var (
 		// Function definition argument type shortcuts.
-		i64Null, _      = parser.ParseType("i64?")
-		i64NonNull, _   = parser.ParseType("i64")
-		strNull, _      = parser.ParseType("string?")
-		strNonNull, _   = parser.ParseType("string")
-		any1NonNull, _  = parser.ParseType("any1")
-		any1Nullable, _ = parser.ParseType("any1?")
+		i64Null, _      = parser.ParseType("i64?", nil)
+		i64NonNull, _   = parser.ParseType("i64", nil)
+		strNull, _      = parser.ParseType("string?", nil)
+		strNonNull, _   = parser.ParseType("string", nil)
+		any1NonNull, _  = parser.ParseType("any1", nil)
+		any1Nullable, _ = parser.ParseType("any1?", nil)
 		any1listNonNull = mkFuncArgList(any1NonNull)
 
 		// Few shortcut type definitions.
@@ -206,9 +206,9 @@ func TestEvaluateTypeExpression(t *testing.T) {
 
 func TestVariantWithVariadic(t *testing.T) {
 	var (
-		i64Null, _     = parser.ParseType("i64?")
-		i64NonNull, _  = parser.ParseType("i64")
-		varcharNull, _ = parser.ParseType("varchar?<20>")
+		i64Null, _     = parser.ParseType("i64?", nil)
+		i64NonNull, _  = parser.ParseType("i64", nil)
+		varcharNull, _ = parser.ParseType("varchar?<20>", nil)
 	)
 
 	tests := []struct {
@@ -473,7 +473,9 @@ func TestResolveType(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			// set up type registry
-			returnType, _ := parser.ParseType(tt.returnType)
+			returnType, _ := parser.ParseType(tt.returnType, func(name string, nullability types.Nullability, parameters []types.UDTParameter) (*types.ParameterizedUserDefinedType, error) {
+				return &types.ParameterizedUserDefinedType{Name: name, URN: "extension:org:item", Nullability: nullability, TypeParameters: parameters}, nil
+			})
 			registry := extensions.NewSet()
 			var expectedRef uint32
 			if tt.expectedUDT {
@@ -512,8 +514,10 @@ func TestResolveType(t *testing.T) {
 
 func TestResolveTypeErrorHandling(t *testing.T) {
 	// Test error propagation from EvaluateTypeExpression
-	returnType, _ := parser.ParseType("u!custom_type")
-	argType, _ := parser.ParseType("i64")
+	returnType, _ := parser.ParseType("u!custom_type", func(name string, nullability types.Nullability, parameters []types.UDTParameter) (*types.ParameterizedUserDefinedType, error) {
+		return &types.ParameterizedUserDefinedType{Name: name, URN: "extension:org:item", Nullability: nullability, TypeParameters: parameters}, nil
+	})
+	argType, _ := parser.ParseType("i64", nil)
 
 	// Create function parameter list that expects one argument
 	funcParams := extensions.FuncParameterList{valArg(argType)}

types/parameterized_user_defined_type.go

@@ -110,6 +110,7 @@ type ParameterizedUserDefinedType struct {
 	TypeVariationRef uint32
 	TypeParameters   []UDTParameter
 	Name             string
+	URN              string
 }
 
 func (m *ParameterizedUserDefinedType) SetNullability(n Nullability) FuncDefArgType {