refactor: store function expr in flatmap step (#3707)

Minor refactor to have the flat map exec step store the function call
rather than the applier.

ksql-engine/src/main/java/io/confluent/ksql/structured/SchemaKStream.java

@@ -23,16 +23,12 @@
 import com.google.common.collect.ImmutableList;
 import io.confluent.ksql.engine.rewrite.StatementRewriteForRowtime;
 import io.confluent.ksql.execution.builder.KsqlQueryBuilder;
-import io.confluent.ksql.execution.codegen.CodeGenRunner;
-import io.confluent.ksql.execution.codegen.ExpressionMetadata;
 import io.confluent.ksql.execution.context.QueryContext;
 import io.confluent.ksql.execution.context.QueryLoggerUtil;
 import io.confluent.ksql.execution.ddl.commands.KsqlTopic;
 import io.confluent.ksql.execution.expression.tree.ColumnReferenceExp;
 import io.confluent.ksql.execution.expression.tree.Expression;
 import io.confluent.ksql.execution.expression.tree.FunctionCall;
-import io.confluent.ksql.execution.function.UdtfUtil;
-import io.confluent.ksql.execution.function.udtf.TableFunctionApplier;
 import io.confluent.ksql.execution.plan.AbstractStreamSource;
 import io.confluent.ksql.execution.plan.ExecutionStep;
 import io.confluent.ksql.execution.plan.ExecutionStepProperties;
@@ -55,7 +51,6 @@
 import io.confluent.ksql.execution.plan.WindowedStreamSource;
 import io.confluent.ksql.execution.streams.ExecutionStepFactory;
 import io.confluent.ksql.function.FunctionRegistry;
-import io.confluent.ksql.function.KsqlTableFunction;
 import io.confluent.ksql.metastore.model.DataSource;
 import io.confluent.ksql.metastore.model.KeyField;
 import io.confluent.ksql.metastore.model.KeyField.LegacyField;
@@ -72,7 +67,6 @@
 import io.confluent.ksql.util.IdentifierUtil;
 import io.confluent.ksql.util.KsqlConfig;
 import io.confluent.ksql.util.SchemaUtil;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
@@ -691,33 +685,13 @@ public SchemaKStream<K> flatMap(
       final List<FunctionCall> tableFunctions,
       final QueryContext.Stacker contextStacker
   ) {
-    final List<TableFunctionApplier> tableFunctionAppliers = new ArrayList<>(tableFunctions.size());
-    final CodeGenRunner codeGenRunner =
-        new CodeGenRunner(getSchema(), ksqlConfig, functionRegistry);
-    for (FunctionCall functionCall: tableFunctions) {
-      final List<ExpressionMetadata> expressionMetadataList = new ArrayList<>(
-          functionCall.getArguments().size());
-      for (Expression expression : functionCall.getArguments()) {
-        final ExpressionMetadata expressionMetadata =
-            codeGenRunner.buildCodeGenFromParseTree(expression, "Table function");
-        expressionMetadataList.add(expressionMetadata);
-      }
-      final KsqlTableFunction tableFunction = UdtfUtil.resolveTableFunction(
-          functionRegistry,
-          functionCall,
-          getSchema()
-      );
-      final TableFunctionApplier tableFunctionApplier =
-          new TableFunctionApplier(tableFunction, expressionMetadataList);
-      tableFunctionAppliers.add(tableFunctionApplier);
-    }
     final StreamFlatMap<K> step = ExecutionStepFactory.streamFlatMap(
         contextStacker,
         sourceStep,
         outputSchema,
-        tableFunctionAppliers
+        tableFunctions
     );
-    return new SchemaKStream<K>(
+    return new SchemaKStream<>(
         step,
         keyFormat,
         keyField,

ksql-execution/src/main/java/io/confluent/ksql/execution/plan/StreamFlatMap.java

@@ -15,7 +15,7 @@
 package io.confluent.ksql.execution.plan;
 
 import com.google.errorprone.annotations.Immutable;
-import io.confluent.ksql.execution.function.udtf.TableFunctionApplier;
+import io.confluent.ksql.execution.expression.tree.FunctionCall;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
@@ -25,16 +25,16 @@ public class StreamFlatMap<K> implements ExecutionStep<KStreamHolder<K>> {
 
   private final ExecutionStepProperties properties;
   private final ExecutionStep<KStreamHolder<K>> source;
-  private final List<TableFunctionApplier> tableFunctionAppliers;
+  private final List<FunctionCall> tableFunctions;
 
   public StreamFlatMap(
       final ExecutionStepProperties properties,
       final ExecutionStep<KStreamHolder<K>> source,
-      final List<TableFunctionApplier> tableFunctionAppliers
+      final List<FunctionCall> tableFunctionAppliers
   ) {
     this.properties = Objects.requireNonNull(properties, "properties");
     this.source = Objects.requireNonNull(source, "source");
-    this.tableFunctionAppliers = Objects.requireNonNull(tableFunctionAppliers);
+    this.tableFunctions = Objects.requireNonNull(tableFunctionAppliers);
   }
 
   @Override
@@ -52,8 +52,8 @@ public KStreamHolder<K> build(final PlanBuilder builder) {
     return builder.visitFlatMap(this);
   }
 
-  public List<TableFunctionApplier> getTableFunctionAppliers() {
-    return tableFunctionAppliers;
+  public List<FunctionCall> getTableFunctions() {
+    return tableFunctions;
   }
 
   public ExecutionStep<KStreamHolder<K>> getSource() {

ksql-streams/src/main/java/io/confluent/ksql/execution/streams/ExecutionStepFactory.java

@@ -19,7 +19,6 @@
 import io.confluent.ksql.execution.context.QueryContext;
 import io.confluent.ksql.execution.expression.tree.Expression;
 import io.confluent.ksql.execution.expression.tree.FunctionCall;
-import io.confluent.ksql.execution.function.udtf.TableFunctionApplier;
 import io.confluent.ksql.execution.plan.DefaultExecutionStepProperties;
 import io.confluent.ksql.execution.plan.ExecutionStep;
 import io.confluent.ksql.execution.plan.Formats;
@@ -147,13 +146,13 @@ public static <K> StreamFlatMap<K> streamFlatMap(
       final QueryContext.Stacker stacker,
       final ExecutionStep<KStreamHolder<K>> source,
       final LogicalSchema resultSchema,
-      final List<TableFunctionApplier> tableFunctionAppliers
+      final List<FunctionCall> tableFunctions
   ) {
     final QueryContext queryContext = stacker.getQueryContext();
     return new StreamFlatMap<>(
         new DefaultExecutionStepProperties(resultSchema, queryContext),
         source,
-        tableFunctionAppliers
+        tableFunctions
     );
   }
 

ksql-streams/src/main/java/io/confluent/ksql/execution/streams/KSPlanBuilder.java

@@ -130,7 +130,7 @@ public <K> KStreamHolder<K> visitStreamMapValues(
   @Override
   public <K> KStreamHolder<K> visitFlatMap(final StreamFlatMap<K> streamFlatMap) {
     final KStreamHolder<K> source = streamFlatMap.getSource().build(this);
-    return StreamFlatMapBuilder.build(source, streamFlatMap);
+    return StreamFlatMapBuilder.build(source, streamFlatMap, queryBuilder);
   }
 
   @Override

ksql-streams/src/main/java/io/confluent/ksql/execution/streams/StreamFlatMapBuilder.java

@@ -15,9 +15,20 @@
 
 package io.confluent.ksql.execution.streams;
 
+import io.confluent.ksql.execution.builder.KsqlQueryBuilder;
+import io.confluent.ksql.execution.codegen.CodeGenRunner;
+import io.confluent.ksql.execution.codegen.ExpressionMetadata;
+import io.confluent.ksql.execution.expression.tree.Expression;
+import io.confluent.ksql.execution.expression.tree.FunctionCall;
+import io.confluent.ksql.execution.function.UdtfUtil;
 import io.confluent.ksql.execution.function.udtf.KudtfFlatMapper;
+import io.confluent.ksql.execution.function.udtf.TableFunctionApplier;
 import io.confluent.ksql.execution.plan.KStreamHolder;
 import io.confluent.ksql.execution.plan.StreamFlatMap;
+import io.confluent.ksql.function.KsqlTableFunction;
+import io.confluent.ksql.schema.ksql.LogicalSchema;
+import java.util.ArrayList;
+import java.util.List;
 
 public final class StreamFlatMapBuilder {
 
@@ -26,9 +37,32 @@ private StreamFlatMapBuilder() {
 
   public static <K> KStreamHolder<K> build(
       final KStreamHolder<K> stream,
-      final StreamFlatMap<K> step) {
+      final StreamFlatMap<K> step,
+      final KsqlQueryBuilder queryBuilder) {
+    final List<FunctionCall> tableFunctions = step.getTableFunctions();
+    final LogicalSchema schema = step.getSource().getSchema();
+    final List<TableFunctionApplier> tableFunctionAppliers = new ArrayList<>(tableFunctions.size());
+    final CodeGenRunner codeGenRunner =
+        new CodeGenRunner(schema, queryBuilder.getKsqlConfig(), queryBuilder.getFunctionRegistry());
+    for (FunctionCall functionCall: tableFunctions) {
+      final List<ExpressionMetadata> expressionMetadataList = new ArrayList<>(
+          functionCall.getArguments().size());
+      for (Expression expression : functionCall.getArguments()) {
+        final ExpressionMetadata expressionMetadata =
+            codeGenRunner.buildCodeGenFromParseTree(expression, "Table function");
+        expressionMetadataList.add(expressionMetadata);
+      }
+      final KsqlTableFunction tableFunction = UdtfUtil.resolveTableFunction(
+          queryBuilder.getFunctionRegistry(),
+          functionCall,
+          schema
+      );
+      final TableFunctionApplier tableFunctionApplier =
+          new TableFunctionApplier(tableFunction, expressionMetadataList);
+      tableFunctionAppliers.add(tableFunctionApplier);
+    }
     return stream.withStream(stream.getStream().flatMapValues(
-        new KudtfFlatMapper(step.getTableFunctionAppliers())));
+        new KudtfFlatMapper(tableFunctionAppliers))
+    );
   }
-
 }