Support UDF in plan generator

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaData.java

@@ -197,6 +197,11 @@ public Map<String, RecordType> getRecordTypes() {
         return recordTypes;
     }
 
+    @Nonnull
+    public Map<String, UserDefinedFunction> getUserDefinedFunctionMap() {
+        return userDefinedFunctionMap;
+    }
+
     @Nonnull
     public RecordType getRecordType(@Nonnull String name) {
         RecordType recordType = recordTypes.get(name);

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaDataBuilder.java

@@ -228,9 +228,10 @@ private void loadProtoExceptRecords(@Nonnull RecordMetaDataProto.MetaData metaDa
                 typeBuilder.setRecordTypeKey(LiteralKeyExpression.fromProtoValue(typeProto.getExplicitKey()));
             }
         }
+        PlanSerializationContext serializationContext = new PlanSerializationContext(DefaultPlanSerializationRegistry.INSTANCE,
+                PlanHashable.CURRENT_FOR_CONTINUATION);
         for (RecordMetaDataProto.PUserDefinedFunction function: metaDataProto.getUserDefinedFunctionsList()) {
-            UserDefinedFunction func = (UserDefinedFunction)PlanSerialization.dispatchFromProtoContainer(new PlanSerializationContext(DefaultPlanSerializationRegistry.INSTANCE,
-                    PlanHashable.CURRENT_FOR_CONTINUATION), function);
+            UserDefinedFunction func = (UserDefinedFunction)PlanSerialization.dispatchFromProtoContainer(serializationContext, function);
             userDefinedFunctionMap.put(func.getFunctionName(), func);
         }
         if (metaDataProto.hasSplitLongRecords()) {

fdb-relational-core/src/main/antlr/RelationalParser.g4

@@ -88,7 +88,7 @@ utilityStatement
 
 templateClause
     :
-        CREATE ( structDefinition | tableDefinition | enumDefinition | indexDefinition )
+        CREATE ( structDefinition | tableDefinition | enumDefinition | indexDefinition | functionDefinition)
     ;
 
 createStatement
@@ -154,6 +154,10 @@ indexDefinition
     : (UNIQUE)? INDEX indexName=uid AS queryTerm indexAttributes?
     ;
 
+functionDefinition
+    : FUNCTION functionName=uid LEFT_ROUND_BRACKET paramName=uid inputTypeName=columnType RIGHT_ROUND_BRACKET RETURNS columnType AS fullId
+    ;
+
 indexAttributes
     : WITH ATTRIBUTES indexAttribute (COMMA indexAttribute)*
     ;
@@ -560,6 +564,11 @@ uid
     | DOUBLE_QUOTE_ID
     ;
 
+userDefinedFunctionName
+    : ID
+    | DOUBLE_QUOTE_ID
+    ;
+
 // done
 simpleId
     : ID
@@ -789,6 +798,7 @@ functionCall
     : aggregateWindowedFunction                                     #aggregateFunctionCall // done (supported)
     | specificFunction                                              #specificFunctionCall //
     | scalarFunctionName '(' functionArgs? ')'                      #scalarFunctionCall // done (unsupported)
+    | userDefinedFunctionName '(' functionArgs? ')'                 #userDefinedFunctionCall
     ;
 
 specificFunction
@@ -899,7 +909,7 @@ levelInWeightListElement
     ;
 
 aggregateWindowedFunction
-    : functionName=(AVG | MAX | MIN | SUM | MAX_EVER | MIN_EVER )
+    : functionName=(AVG | MAX | MIN | SUM | MAX_EVER | MIN_EVER)
       '(' aggregator=(ALL | DISTINCT)? functionArg ')' overClause?
     | functionName=BITMAP_CONSTRUCT_AGG '(' functionArg ')'
     | functionName=COUNT '(' (starArg='*' | aggregator=ALL? functionArg | aggregator=DISTINCT functionArgs) ')' overClause?

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/RecordLayerSchemaTemplate.java

@@ -26,6 +26,7 @@
 import com.apple.foundationdb.record.RecordMetaDataProto;
 import com.apple.foundationdb.record.metadata.Key;
 import com.apple.foundationdb.record.query.combinatorics.TopologicalSort;
+import com.apple.foundationdb.record.query.plan.cascades.UserDefinedFunction;
 import com.apple.foundationdb.relational.api.exceptions.ErrorCode;
 import com.apple.foundationdb.relational.api.exceptions.RelationalException;
 import com.apple.foundationdb.relational.api.metadata.DataType;
@@ -50,6 +51,7 @@
 import javax.annotation.Nonnull;
 import java.util.BitSet;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
@@ -68,6 +70,9 @@ public final class RecordLayerSchemaTemplate implements SchemaTemplate {
     @Nonnull
     private final Set<RecordLayerTable> tables;
 
+    @Nonnull
+    private final Map<String, UserDefinedFunction> userDefinedFunctionMap;
+
     private final int version;
 
     private final boolean enableLongRows;
@@ -85,11 +90,13 @@ public final class RecordLayerSchemaTemplate implements SchemaTemplate {
 
     private RecordLayerSchemaTemplate(@Nonnull final String name,
                                       @Nonnull final Set<RecordLayerTable> tables,
+                                      @Nonnull final Map<String, UserDefinedFunction> userDefinedFunctionMap,
                                       int version,
                                       boolean enableLongRows,
                                       boolean storeRowVersions) {
         this.name = name;
         this.tables = tables;
+        this.userDefinedFunctionMap = userDefinedFunctionMap;
         this.version = version;
         this.enableLongRows = enableLongRows;
         this.storeRowVersions = storeRowVersions;
@@ -100,13 +107,15 @@ private RecordLayerSchemaTemplate(@Nonnull final String name,
 
     private RecordLayerSchemaTemplate(@Nonnull final String name,
                                       @Nonnull final Set<RecordLayerTable> tables,
+                                      @Nonnull final Map<String, UserDefinedFunction> userDefinedFunctionMap,
                                       int version,
                                       boolean enableLongRows,
                                       boolean storeRowVersions,
                                       @Nonnull final RecordMetaData cachedMetadata) {
         this.name = name;
         this.version = version;
         this.tables = tables;
+        this.userDefinedFunctionMap = userDefinedFunctionMap;
         this.enableLongRows = enableLongRows;
         this.storeRowVersions = storeRowVersions;
         this.metaDataSupplier = Suppliers.memoize(() -> cachedMetadata);
@@ -147,6 +156,11 @@ public RecordLayerSchema generateSchema(@Nonnull String databaseId, @Nonnull Str
         return new RecordLayerSchema(schemaName, databaseId, this);
     }
 
+    @Nonnull
+    public Collection<UserDefinedFunction> getAllUserDefinedFunctions() {
+        return userDefinedFunctionMap.values();
+    }
+
     @Nonnull
     public Descriptors.Descriptor getDescriptor(@Nonnull final String tableName) {
         return toRecordMetadata().getRecordType(tableName).getDescriptor();
@@ -309,12 +323,15 @@ public static final class Builder {
 
         private final Map<String, RecordLayerTable> tables;
 
+        private final Map<String, UserDefinedFunction> functionMap;
+
         private final Map<String, DataType.Named> auxiliaryTypes; // for quick lookup
 
         private RecordMetaData cachedMetadata;
 
         private Builder() {
             tables = new LinkedHashMap<>();
+            functionMap = new HashMap<>();
             auxiliaryTypes = new LinkedHashMap<>();
             // enable long rows is TRUE by default
             enableLongRows = true;
@@ -402,6 +419,18 @@ public Builder addAuxiliaryTypes(@Nonnull Collection<DataType.Named> auxiliaryTy
             return this;
         }
 
+        @Nonnull
+        public Builder addUserDefinedFunction(@Nonnull UserDefinedFunction userDefinedFunction) {
+            functionMap.put(userDefinedFunction.getFunctionName(), userDefinedFunction);
+            return this;
+        }
+
+        @Nonnull
+        public Builder addUserDefinedFunctions(@Nonnull Collection<UserDefinedFunction> functions) {
+            functions.forEach(this::addUserDefinedFunction);
+            return this;
+        }
+
         @Nonnull
         Builder setCachedMetadata(@Nonnull final RecordMetaData metadata) {
             this.cachedMetadata = metadata;
@@ -459,11 +488,10 @@ public RecordLayerSchemaTemplate build() {
             if (needsResolution) {
                 resolveTypes();
             }
-
             if (cachedMetadata != null) {
-                return new RecordLayerSchemaTemplate(name, new LinkedHashSet<>(tables.values()), version, enableLongRows, storeRowVersions, cachedMetadata);
+                return new RecordLayerSchemaTemplate(name, new LinkedHashSet<>(tables.values()), functionMap, version, enableLongRows, storeRowVersions, cachedMetadata);
             } else {
-                return new RecordLayerSchemaTemplate(name, new LinkedHashSet<>(tables.values()), version, enableLongRows, storeRowVersions);
+                return new RecordLayerSchemaTemplate(name, new LinkedHashSet<>(tables.values()), functionMap, version, enableLongRows, storeRowVersions);
             }
         }
 

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/serde/RecordMetadataDeserializer.java

@@ -61,6 +61,9 @@ public RecordLayerSchemaTemplate.Builder getSchemaTemplate(@Nonnull final String
                 .setEnableLongRows(recordMetaData.isSplitLongRecords())
                 .setName(schemaTemplateName)
                 .setIntermingleTables(!recordMetaData.primaryKeyHasRecordTypePrefix());
+        for (final var u: recordMetaData.getUserDefinedFunctionMap().values()) {
+            schemaTemplateBuilder.addUserDefinedFunction(u);
+        }
         final var nameToTableBuilder = new HashMap<String, RecordLayerTable.Builder>();
         for (final var registeredType : registeredTypes) {
             switch (registeredType.getType()) {

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/serde/RecordMetadataSerializer.java

@@ -87,6 +87,7 @@ public void visit(@Nonnull com.apple.foundationdb.relational.api.metadata.Index
     @Override
     public void visit(@Nonnull SchemaTemplate schemaTemplate) {
         Assert.thatUnchecked(schemaTemplate instanceof RecordLayerSchemaTemplate);
+        getBuilder().addUserDefinedFunctions(((RecordLayerSchemaTemplate) schemaTemplate).getAllUserDefinedFunctions());
         getBuilder().setSplitLongRecords(schemaTemplate.isEnableLongRows());
         getBuilder().setStoreRecordVersions(schemaTemplate.isStoreRowVersions());
         getBuilder().setVersion(schemaTemplate.getVersion());

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/Identifier.java

@@ -128,6 +128,13 @@ public boolean prefixedWith(@Nonnull Identifier identifier) {
         return true;
     }
 
+    @Nonnull
+    public List<String> removePrefix(@Nonnull Identifier prefix) {
+        // assume the identifier has the prefix, should call prefixedWith(prefix) to check before calling this method
+        final var fullName = fullyQualifiedName();
+        return fullName.subList(prefix.fullyQualifiedName().size(), fullName.size());
+    }
+
     public boolean qualifiedWith(@Nonnull Identifier identifier) {
         final var identifierFullName = identifier.fullyQualifiedName();
         final var fullName = fullyQualifiedName();

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/SemanticAnalyzer.java

@@ -28,6 +28,7 @@
 import com.apple.foundationdb.record.query.plan.cascades.CorrelationIdentifier;
 import com.apple.foundationdb.record.query.plan.cascades.IndexAccessHint;
 import com.apple.foundationdb.record.query.plan.cascades.Quantifier;
+import com.apple.foundationdb.record.query.plan.cascades.UserDefinedFunction;
 import com.apple.foundationdb.record.query.plan.cascades.typing.Type;
 import com.apple.foundationdb.record.query.plan.cascades.typing.Typed;
 import com.apple.foundationdb.record.query.plan.cascades.values.AggregateValue;
@@ -38,6 +39,7 @@
 import com.apple.foundationdb.record.query.plan.cascades.values.IndexableAggregateValue;
 import com.apple.foundationdb.record.query.plan.cascades.values.LiteralValue;
 import com.apple.foundationdb.record.query.plan.cascades.values.NotValue;
+import com.apple.foundationdb.record.query.plan.cascades.values.QuantifiedObjectValue;
 import com.apple.foundationdb.record.query.plan.cascades.values.RecordConstructorValue;
 import com.apple.foundationdb.record.query.plan.cascades.values.RelOpValue;
 import com.apple.foundationdb.record.query.plan.cascades.values.StreamableAggregateValue;
@@ -51,6 +53,7 @@
 import com.apple.foundationdb.relational.api.metadata.Table;
 import com.apple.foundationdb.relational.generated.RelationalParser;
 import com.apple.foundationdb.relational.recordlayer.metadata.DataTypeUtils;
+import com.apple.foundationdb.relational.recordlayer.metadata.RecordLayerSchemaTemplate;
 import com.apple.foundationdb.relational.recordlayer.query.functions.SqlFunctionCatalog;
 import com.apple.foundationdb.relational.recordlayer.query.functions.SqlFunctionCatalogImpl;
 import com.apple.foundationdb.relational.recordlayer.query.visitors.QueryVisitor;
@@ -100,6 +103,10 @@ public SemanticAnalyzer(@Nonnull SchemaTemplate metadataCatalog,
                             @Nonnull SqlFunctionCatalog functionCatalog) {
         this.metadataCatalog = metadataCatalog;
         this.functionCatalog = functionCatalog;
+        // add UDFs to functionCatalog
+        for (UserDefinedFunction function: ((RecordLayerSchemaTemplate) metadataCatalog).getAllUserDefinedFunctions()) {
+            (this.functionCatalog).addUdfFunction(function);
+        }
     }
 
     /**
@@ -473,6 +480,46 @@ public Optional<Expression> lookupNestedField(@Nonnull Identifier requestedIdent
         return Optional.of(nestedAttribute);
     }
 
+    @Nonnull
+    public Optional<Value> lookupNestedField(@Nonnull Identifier requestedIdentifier,
+                                             @Nonnull Identifier paramId,
+                                             @Nonnull QuantifiedObjectValue existingValue,
+                                             @Nonnull DataType targetDataType) {
+        Assert.thatUnchecked(requestedIdentifier.prefixedWith(paramId), "Invalid function definition");
+
+        // x -> x
+        if (requestedIdentifier.fullyQualifiedName().size() == paramId.fullyQualifiedName().size()) {
+            Assert.thatUnchecked(existingValue.getResultType().equals(DataTypeUtils.toRecordLayerType(targetDataType)), ErrorCode.DATATYPE_MISMATCH, "Result data types don't match!");
+            return Optional.of(existingValue);
+        }
+        // find nested field path
+        final var remainingPath = requestedIdentifier.removePrefix(paramId);
+        final ImmutableList.Builder<FieldValue.Accessor> accessors = ImmutableList.builder();
+        DataType existingDataType = DataTypeUtils.toRelationalType(existingValue.getResultType());
+        for (String s : remainingPath) {
+            if (existingDataType.getCode() != DataType.Code.STRUCT) {
+                return Optional.empty();
+            }
+            final var fields = ((DataType.StructType) existingDataType).getFields();
+            var found = false;
+            for (int j = 0; j < fields.size(); j++) {
+                if (fields.get(j).getName().equals(s)) {
+                    accessors.add(new FieldValue.Accessor(fields.get(j).getName(), j));
+                    existingDataType = fields.get(j).getType();
+                    found = true;
+                    break;
+                }
+            }
+            if (!found) {
+                return Optional.empty();
+            }
+        }
+        final var fieldPath = FieldValue.resolveFieldPath(existingValue.getResultType(), accessors.build());
+        final var fieldValue = FieldValue.ofFieldsAndFuseIfPossible(existingValue, fieldPath);
+        Assert.thatUnchecked(fieldValue.getResultType().equals(DataTypeUtils.toRecordLayerType(targetDataType)), ErrorCode.DATATYPE_MISMATCH, "Result data types don't match!");
+        return Optional.of(fieldValue);
+    }
+
     @Nonnull
     public DataType lookupType(@Nonnull Identifier typeIdentifier, boolean isNullable, boolean isRepeated,
                                @Nonnull Function<String, Optional<DataType>> dataTypeProvider) {

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/functions/SqlFunctionCatalog.java

@@ -20,17 +20,19 @@
 
 package com.apple.foundationdb.relational.recordlayer.query.functions;
 
-import com.apple.foundationdb.record.query.plan.cascades.BuiltInFunction;
-import com.apple.foundationdb.record.query.plan.cascades.typing.Typed;
+import com.apple.foundationdb.record.query.plan.cascades.CatalogedFunction;
+import com.apple.foundationdb.record.query.plan.cascades.UserDefinedFunction;
 import com.apple.foundationdb.relational.recordlayer.query.Expression;
 
 import javax.annotation.Nonnull;
 
 public interface SqlFunctionCatalog {
     @Nonnull
-    BuiltInFunction<? extends Typed> lookUpFunction(@Nonnull String name, @Nonnull Expression... expressions);
+    CatalogedFunction lookUpFunction(@Nonnull String name, @Nonnull Expression... expressions);
 
     boolean containsFunction(@Nonnull String name);
 
     boolean isUdfFunction(@Nonnull String name);
+
+    void addUdfFunction(@Nonnull UserDefinedFunction function);
 }

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/functions/SqlFunctionCatalogImpl.java

@@ -23,6 +23,8 @@
 import com.apple.foundationdb.annotation.API;
 
 import com.apple.foundationdb.record.query.plan.cascades.BuiltInFunction;
+import com.apple.foundationdb.record.query.plan.cascades.CatalogedFunction;
+import com.apple.foundationdb.record.query.plan.cascades.UserDefinedFunction;
 import com.apple.foundationdb.record.query.plan.cascades.typing.Typed;
 import com.apple.foundationdb.record.query.plan.cascades.values.FunctionCatalog;
 import com.apple.foundationdb.record.query.plan.cascades.values.RecordConstructorValue;
@@ -33,7 +35,9 @@
 import com.google.common.collect.ImmutableMap;
 
 import javax.annotation.Nonnull;
+import java.util.HashMap;
 import java.util.Locale;
+import java.util.Map;
 import java.util.Objects;
 import java.util.function.Function;
 import java.util.stream.StreamSupport;
@@ -52,26 +56,39 @@ public final class SqlFunctionCatalogImpl implements SqlFunctionCatalog {
     @Nonnull
     private final ImmutableMap<String, Function<Integer, BuiltInFunction<? extends Typed>>> synonyms;
 
+    @Nonnull
+    private final Map<String, UserDefinedFunction> userDefinedFunctionMap = new HashMap<>();
+
+
     private SqlFunctionCatalogImpl() {
         this.synonyms = createSynonyms();
     }
 
     @Nonnull
     @Override
-    public BuiltInFunction<? extends Typed> lookUpFunction(@Nonnull final String name, @Nonnull final Expression... expressions) {
-        return Assert.notNullUnchecked(Objects.requireNonNull(synonyms.get(name.toLowerCase(Locale.ROOT))).apply(expressions.length));
+    public CatalogedFunction lookUpFunction(@Nonnull final String name, @Nonnull final Expression... expressions) {
+        if (synonyms.get(name.toLowerCase(Locale.ROOT)) != null) {
+            return Assert.notNullUnchecked(Objects.requireNonNull(synonyms.get(name.toLowerCase(Locale.ROOT))).apply(expressions.length));
+        } else {
+            return userDefinedFunctionMap.get(name.toLowerCase(Locale.ROOT));
+        }
     }
 
     @Override
     public boolean containsFunction(@Nonnull String name) {
-        return synonyms.containsKey(name.toLowerCase(Locale.ROOT));
+        return synonyms.containsKey(name.toLowerCase(Locale.ROOT)) || userDefinedFunctionMap.containsKey(name.toLowerCase(Locale.ROOT));
     }
 
     @Override
     public boolean isUdfFunction(@Nonnull final String name) {
         return "java_call".equals(name.trim().toLowerCase(Locale.ROOT));
     }
 
+    @Override
+    public void addUdfFunction(@Nonnull final UserDefinedFunction function) {
+        userDefinedFunctionMap.put(function.getFunctionName(), function);
+    }
+
     @Nonnull
     private static ImmutableMap<String, Function<Integer, BuiltInFunction<? extends Typed>>> createSynonyms() {
         return ImmutableMap.<String, Function<Integer, BuiltInFunction<? extends Typed>>>builder()