feat: support user-defined types in YAML extensions

extensions/simple_extension_test.go

@@ -10,6 +10,9 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/substrait-io/substrait-go/extensions"
+	"github.com/substrait-io/substrait-go/proto"
+	"github.com/substrait-io/substrait-go/types"
+	"github.com/substrait-io/substrait-go/types/parser"
 )
 
 func TestUnmarshalSimpleExtension(t *testing.T) {
@@ -41,6 +44,46 @@ types:
 	assert.Nil(t, f.Types[3].Structure)
 }
 
+func TestUnmarshalCustomScalarFunction(t *testing.T) {
+	const customDef = `
+scalar_functions:
+  - name: "scalar1"
+    impls:
+      - args:
+          - name: arg1
+            value: u!customtype1
+        return: i64
+  - name: "scalar2"
+    impls:
+      - args:
+          - name: arg1
+            value: i64
+        return: u!customtype2?
+`
+
+	var f extensions.SimpleExtensionFile
+	require.NoError(t, yaml.Unmarshal([]byte(customDef), &f))
+	assert.Len(t, f.ScalarFunctions, 2)
+
+	assert.Equal(t, "scalar1", f.ScalarFunctions[0].Name)
+	assert.IsType(t, extensions.ValueArg{}, f.ScalarFunctions[0].Impls[0].Args[0])
+	arg1 := f.ScalarFunctions[0].Impls[0].Args[0].(extensions.ValueArg)
+	assert.Equal(t, "u!customtype1", arg1.Value.String())
+	typ, err := arg1.Value.Expr.(*parser.Type).TypeDef.RetType()
+	assert.NoError(t, err)
+	assert.IsType(t, &types.UserDefinedType{}, typ)
+	assert.Equal(t, proto.Type_NULLABILITY_REQUIRED, typ.GetNullability(), "expected NULLABILITY_REQUIRED")
+
+	assert.Equal(t, "scalar2", f.ScalarFunctions[1].Name)
+	assert.IsType(t, extensions.ValueArg{}, f.ScalarFunctions[1].Impls[0].Args[0])
+	ret := f.ScalarFunctions[1].Impls[0].Return
+	assert.Equal(t, "u!customtype2?", ret.String())
+	typ, err = ret.Expr.(*parser.Type).TypeDef.RetType()
+	assert.NoError(t, err)
+	assert.IsType(t, &types.UserDefinedType{}, typ)
+	assert.Equal(t, proto.Type_NULLABILITY_NULLABLE, typ.GetNullability(), "expected NULLABILITY_NULLABLE")
+}
+
 func TestUnmarshalSimpleExtensionScalarFunction(t *testing.T) {
 	const addDef = `
 scalar_functions:

types/parser/type_parser.go

@@ -122,7 +122,7 @@ func (t *typename) Capture(values []string) error {
 }
 
 type nonParamType struct {
-	TypeName    typename `parser:"@(IntType | Boolean | FPType | Temporal | BinaryType)"`
+	TypeName    typename `parser:"@(IntType | Boolean | FPType | Temporal | BinaryType | UserDefinedType)"`
 	Nullability bool     `parser:"@'?'?"`
 	// Variation   int      `parser:"'[' @\d+ ']'?"`
 }
@@ -148,7 +148,11 @@ func (t *nonParamType) RetType() (types.Type, error) {
 	} else {
 		n = types.NullabilityRequired
 	}
-	typ, err := types.SimpleTypeNameToType(types.TypeName(t.TypeName))
+	typName := t.TypeName
+	if strings.HasPrefix(string(typName), "u!") {
+		typName = "u!"
+	}
+	typ, err := types.SimpleTypeNameToType(types.TypeName(typName))
 	if err == nil {
 		return typ.WithNullability(n), nil
 	}
@@ -598,6 +602,7 @@ var (
 		{Name: "LengthType", Pattern: `fixedchar|varchar|fixedbinary|precision_timestamp_tz|precision_timestamp|interval_day`},
 		{Name: "Int", Pattern: `[-+]?\d+`},
 		{Name: "ParamType", Pattern: `(?i)(struct|list|decimal|map)`},
+		{Name: "UserDefinedType", Pattern: `u![a-zA-Z_][a-zA-Z0-9_]*`},
 		{Name: "Identifier", Pattern: `[a-zA-Z_$][a-zA-Z_$0-9]*`},
 		{Name: "Ident", Pattern: `([a-zA-Z_]\w*)|[><,?]`},
 	})

types/parser/type_parser_test.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/substrait-io/substrait-go/proto"
 	"github.com/substrait-io/substrait-go/types"
 	"github.com/substrait-io/substrait-go/types/integer_parameters"
 	"github.com/substrait-io/substrait-go/types/parser"
@@ -100,3 +101,83 @@ func TestParserRetType(t *testing.T) {
 		})
 	}
 }
+
+func TestParserListType(t *testing.T) {
+	tests := []struct {
+		expr        string
+		expected    string
+		expectedTyp types.Type
+	}{
+		{
+			expr:     "list<i32>",
+			expected: "list<i32>",
+			expectedTyp: &types.ListType{
+				Nullability: types.NullabilityRequired,
+				Type:        &types.Int32Type{Nullability: types.NullabilityRequired},
+			},
+		},
+		{
+			expr:     "list?<i16?>",
+			expected: "list?<i16?>",
+			expectedTyp: &types.ListType{
+				Nullability: types.NullabilityNullable,
+				Type:        &types.Int16Type{Nullability: types.NullabilityNullable},
+			},
+		},
+		{
+			expr:     "list<i16>",
+			expected: "list<i16>",
+			expectedTyp: &types.ListType{
+				Nullability: types.NullabilityRequired,
+				Type:        &types.Int16Type{Nullability: types.NullabilityRequired},
+			},
+		},
+	}
+
+	p, err := parser.New()
+	require.NoError(t, err)
+
+	for _, td := range tests {
+		t.Run(td.expr, func(t *testing.T) {
+			d, err := p.ParseString(td.expr)
+			assert.NoError(t, err)
+			assert.Equal(t, td.expected, d.Expr.String())
+
+			if tExpr, ok := d.Expr.(*parser.Type); ok {
+				retType, err := tExpr.RetType()
+				assert.NoError(t, err)
+				assert.Equal(t, reflect.TypeOf(td.expectedTyp), reflect.TypeOf(retType))
+				assert.Equal(t, td.expectedTyp, retType)
+			}
+		})
+	}
+}
+
+func TestParseUDT(t *testing.T) {
+	tests := []struct {
+		expr                string
+		expected            string
+		expectedTyp         types.Type
+		expectedNullability proto.Type_Nullability
+		expectedOptional    bool
+	}{
+		{"u!customtype1", "u!customtype1", &types.UserDefinedType{}, proto.Type_NULLABILITY_REQUIRED, false},
+		{"u!customtype2?", "u!customtype2?", &types.UserDefinedType{}, proto.Type_NULLABILITY_NULLABLE, true},
+	}
+
+	p, err := parser.New()
+	require.NoError(t, err)
+
+	for _, td := range tests {
+		t.Run(td.expr, func(t *testing.T) {
+			d, err := p.ParseString(td.expr)
+			assert.NoError(t, err)
+			assert.Equal(t, td.expected, d.Expr.String())
+			retType, err := d.Expr.(*parser.Type).RetType()
+			assert.NoError(t, err)
+			assert.Equal(t, td.expectedNullability, retType.GetNullability())
+			assert.Equal(t, td.expectedOptional, d.Expr.(*parser.Type).Optional())
+			assert.Equal(t, reflect.TypeOf(td.expectedTyp), reflect.TypeOf(retType))
+		})
+	}
+}

types/types.go

@@ -43,6 +43,7 @@ const (
 	TypeNameIntervalYear TypeName = "interval_year"
 	TypeNameIntervalDay  TypeName = "interval_day"
 	TypeNameUUID         TypeName = "uuid"
+	TypeNameUDT          TypeName = "u!"
 
 	TypeNameFixedBinary          TypeName = "fixedbinary"
 	TypeNameFixedChar            TypeName = "fixedchar"
@@ -68,6 +69,7 @@ var simpleTypeNameMap = map[TypeName]Type{
 	TypeNameTimestampTz:  &TimestampTzType{},
 	TypeNameIntervalYear: &IntervalYearType{},
 	TypeNameUUID:         &UUIDType{},
+	TypeNameUDT:          &UserDefinedType{},
 }
 
 var fixedTypeNameMap = map[TypeName]FixedType{