feat: add GetFunctionRegistry in LocalFunctionRegistry (substrait-io#99)

srikrishnak · Jan 6, 2025 · 18b7daf · 18b7daf
1 parent 298eedc
commit 18b7daf
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 0 deletions.
diff --git a/functions/dialect.go b/functions/dialect.go
@@ -84,6 +84,7 @@ func (d *dialectImpl) LocalizeFunctionRegistry(registry FunctionRegistry) (Local
 		allFunctions:         allVariants,
 		idToLocalFunctionMap: makeLocalFunctionVariantsMap(allVariants),
 		localTypeRegistry:    localTypeRegistry,
+		funcRegistry:         registry,
 	}, nil
 }
 

diff --git a/functions/dialect_test.go b/functions/dialect_test.go
@@ -68,6 +68,7 @@ window_functions:
 	localRegistry, err := dialect.LocalizeFunctionRegistry(gFunctionRegistry)
 	assert.NoError(t, err)
 	assert.Equal(t, t.Name(), localRegistry.GetDialect().Name())
+	assert.Equal(t, gFunctionRegistry, localRegistry.GetFunctionRegistry())
 }
 
 func TestBadDialects(t *testing.T) {

diff --git a/functions/local_functions.go b/functions/local_functions.go
@@ -19,6 +19,7 @@ type localFunctionRegistryImpl struct {
 
 	idToLocalFunctionMap map[extensions.ID]localFunctionVariant
 	localTypeRegistry    LocalTypeRegistry
+	funcRegistry         FunctionRegistry
 }
 
 func makeLocalFunctionVariantsMap(functions []extensions.FunctionVariant) map[extensions.ID]localFunctionVariant {
@@ -44,6 +45,10 @@ func (l *localFunctionRegistryImpl) GetDialect() Dialect {
 	return l.dialect
 }
 
+func (l *localFunctionRegistryImpl) GetFunctionRegistry() FunctionRegistry {
+	return l.funcRegistry
+}
+
 func (l *localFunctionRegistryImpl) GetScalarFunctions(name FunctionName, numArgs int) []*LocalScalarFunctionVariant {
 	return getFunctionVariantsByCount(getOrEmpty(name, l.scalarFunctions), numArgs)
 }

diff --git a/functions/registries.go b/functions/registries.go
@@ -86,6 +86,7 @@ type FunctionRegistry interface {
 type LocalFunctionRegistry interface {
 	functionRegistryBase[FunctionName, *LocalScalarFunctionVariant, *LocalAggregateFunctionVariant, *LocalWindowFunctionVariant]
 	GetDialect() Dialect
+	GetFunctionRegistry() FunctionRegistry
 	GetScalarFunctionByInvocation(scalarFuncInvocation *expr.ScalarFunction) (*LocalScalarFunctionVariant, error)
 	GetAggregateFunctionByInvocation(aggregateFuncInvocation *expr.AggregateFunction) (*LocalAggregateFunctionVariant, error)
 	GetWindowFunctionByInvocation(windowFuncInvocation *expr.WindowFunction) (*LocalWindowFunctionVariant, error)