Add Iceberg system.bucket function to Lakehouse

LakehouseMetadata will delegate all calls related to functions to IcebergMetadata.
The system.bucket function will always be delegated to Iceberg function provider.

If there is no connector selected, the full query is like this:
SELECT lakehouse.system.bucket('trino', 16)

Author: codluca

plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/LakehouseConnector.java

@@ -26,10 +26,12 @@
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.connector.ConnectorSplitManager;
 import io.trino.spi.connector.ConnectorTransactionHandle;
+import io.trino.spi.function.FunctionProvider;
 import io.trino.spi.session.PropertyMetadata;
 import io.trino.spi.transaction.IsolationLevel;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 
 import static com.google.common.collect.Sets.immutableEnumSet;
@@ -51,6 +53,7 @@ public class LakehouseConnector
     private final LakehouseSessionProperties sessionProperties;
     private final LakehouseTableProperties tableProperties;
     private final IcebergMaterializedViewProperties materializedViewProperties;
+    private final FunctionProvider functionProvider;
 
     @Inject
     public LakehouseConnector(
@@ -62,7 +65,8 @@ public LakehouseConnector(
             LakehouseNodePartitioningProvider nodePartitioningProvider,
             LakehouseSessionProperties sessionProperties,
             LakehouseTableProperties tableProperties,
-            IcebergMaterializedViewProperties materializedViewProperties)
+            IcebergMaterializedViewProperties materializedViewProperties,
+            FunctionProvider functionProvider)
     {
         this.lifeCycleManager = requireNonNull(lifeCycleManager, "lifeCycleManager is null");
         this.transactionManager = requireNonNull(transactionManager, "transactionManager is null");
@@ -73,6 +77,7 @@ public LakehouseConnector(
         this.sessionProperties = requireNonNull(sessionProperties, "sessionProperties is null");
         this.tableProperties = requireNonNull(tableProperties, "tableProperties is null");
         this.materializedViewProperties = requireNonNull(materializedViewProperties, "materializedViewProperties is null");
+        this.functionProvider = requireNonNull(functionProvider, "functionProvider is null");
     }
 
     @Override
@@ -159,4 +164,10 @@ public Set<ConnectorCapabilities> getCapabilities()
     {
         return immutableEnumSet(NOT_NULL_COLUMN_CONSTRAINT, MATERIALIZED_VIEW_GRACE_PERIOD);
     }
+
+    @Override
+    public Optional<FunctionProvider> getFunctionProvider()
+    {
+        return Optional.of(functionProvider);
+    }
 }

plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/LakehouseFunctionProvider.java

@@ -0,0 +1,51 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.lakehouse;
+
+import com.google.inject.Inject;
+import io.trino.plugin.iceberg.functions.IcebergFunctionProvider;
+import io.trino.spi.function.BoundSignature;
+import io.trino.spi.function.FunctionDependencies;
+import io.trino.spi.function.FunctionId;
+import io.trino.spi.function.FunctionProvider;
+import io.trino.spi.function.InvocationConvention;
+import io.trino.spi.function.ScalarFunctionImplementation;
+
+import static java.util.Objects.requireNonNull;
+
+public class LakehouseFunctionProvider
+        implements FunctionProvider
+{
+    private final IcebergFunctionProvider icebergFunctionProvider;
+
+    @Inject
+    public LakehouseFunctionProvider(
+            IcebergFunctionProvider icebergFunctionProvider)
+    {
+        this.icebergFunctionProvider = requireNonNull(icebergFunctionProvider, "icebergFunctionProvider is null");
+    }
+
+    @Override
+    public ScalarFunctionImplementation getScalarFunctionImplementation(
+            FunctionId functionId,
+            BoundSignature boundSignature,
+            FunctionDependencies functionDependencies,
+            InvocationConvention invocationConvention)
+    {
+        if ("bucket".equals(boundSignature.getName().getFunctionName())) {
+            return icebergFunctionProvider.getScalarFunctionImplementation(functionId, boundSignature, functionDependencies, invocationConvention);
+        }
+        throw new UnsupportedOperationException("%s provides only 'bucket' scalar function".formatted(getClass().getName()));
+    }
+}

plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/LakehouseIcebergModule.java

@@ -41,6 +41,7 @@
 import io.trino.plugin.iceberg.catalog.glue.IcebergGlueCatalogModule;
 import io.trino.plugin.iceberg.catalog.hms.IcebergHiveMetastoreCatalogModule;
 import io.trino.plugin.iceberg.fileio.ForwardingFileIoFactory;
+import io.trino.plugin.iceberg.functions.IcebergFunctionProvider;
 
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 import static io.airlift.configuration.ConfigBinder.configBinder;
@@ -75,6 +76,8 @@ protected void setup(Binder binder)
 
         binder.bind(ForwardingFileIoFactory.class).in(Scopes.SINGLETON);
 
+        binder.bind(IcebergFunctionProvider.class).in(Scopes.SINGLETON);
+
         install(switch (buildConfigObject(MetastoreTypeConfig.class).getMetastoreType()) {
             case THRIFT -> new IcebergHiveMetastoreCatalogModule();
             case FILE -> new IcebergFileMetastoreCatalogModule();

plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/LakehouseMetadata.java

@@ -85,6 +85,10 @@
 import io.trino.spi.connector.WriterScalingOptions;
 import io.trino.spi.expression.ConnectorExpression;
 import io.trino.spi.expression.Constant;
+import io.trino.spi.function.BoundSignature;
+import io.trino.spi.function.FunctionDependencyDeclaration;
+import io.trino.spi.function.FunctionId;
+import io.trino.spi.function.FunctionMetadata;
 import io.trino.spi.function.LanguageFunction;
 import io.trino.spi.function.SchemaFunctionName;
 import io.trino.spi.security.GrantInfo;
@@ -173,6 +177,30 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable
         return hiveMetadata.getTableHandle(session, tableName, startVersion, endVersion);
     }
 
+    @Override
+    public Collection<FunctionMetadata> listFunctions(ConnectorSession session, String schemaName)
+    {
+        return icebergMetadata.listFunctions(session, schemaName);
+    }
+
+    @Override
+    public Collection<FunctionMetadata> getFunctions(ConnectorSession session, SchemaFunctionName name)
+    {
+        return icebergMetadata.getFunctions(session, name);
+    }
+
+    @Override
+    public FunctionMetadata getFunctionMetadata(ConnectorSession session, FunctionId functionId)
+    {
+        return icebergMetadata.getFunctionMetadata(session, functionId);
+    }
+
+    @Override
+    public FunctionDependencyDeclaration getFunctionDependencies(ConnectorSession session, FunctionId functionId, BoundSignature boundSignature)
+    {
+        return icebergMetadata.getFunctionDependencies(session, functionId, boundSignature);
+    }
+
     @Override
     public Optional<ConnectorTableExecuteHandle> getTableHandleForExecute(ConnectorSession session, ConnectorAccessControl accessControl, ConnectorTableHandle tableHandle, String procedureName, Map<String, Object> executeProperties, RetryMode retryMode)
     {

plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/LakehouseModule.java

@@ -24,6 +24,8 @@
 import io.trino.plugin.hive.orc.OrcWriterConfig;
 import io.trino.plugin.hive.parquet.ParquetReaderConfig;
 import io.trino.plugin.hive.parquet.ParquetWriterConfig;
+import io.trino.plugin.iceberg.functions.tablechanges.TableChangesFunctionProcessorProviderFactory;
+import io.trino.spi.function.FunctionProvider;
 
 import static io.airlift.configuration.ConfigBinder.configBinder;
 import static org.weakref.jmx.guice.ExportBinder.newExporter;
@@ -53,6 +55,9 @@ protected void setup(Binder binder)
         binder.bind(FileFormatDataSourceStats.class).in(Scopes.SINGLETON);
         newExporter(binder).export(FileFormatDataSourceStats.class).withGeneratedName();
 
+        binder.bind(FunctionProvider.class).to(LakehouseFunctionProvider.class).in(Scopes.SINGLETON);
+        binder.bind(TableChangesFunctionProcessorProviderFactory.class).in(Scopes.SINGLETON);
+
         binder.bind(Key.get(boolean.class, HideDeltaLakeTables.class)).toInstance(false);
     }
 }

plugin/trino-lakehouse/src/test/java/io/trino/plugin/lakehouse/TestLakehouseConnectorTest.java

@@ -365,4 +365,11 @@ public void testShowCreateTable()
                    type = 'ICEBERG'
                 )\\E""");
     }
+
+    @Test
+    void testSystemBucket()
+    {
+        assertThat(query("SELECT lakehouse.system.bucket('trino', 16)"))
+                .matches("VALUES 10");
+    }
 }