fix: UdfIT (IGinX-THU#592)

Author: aqni

core/src/main/java/cn/edu/tsinghua/iginx/engine/physical/task/memory/StreamSourceMemoryPhysicalTask.java

@@ -22,23 +22,29 @@
 import cn.edu.tsinghua.iginx.engine.physical.exception.PhysicalException;
 import cn.edu.tsinghua.iginx.engine.physical.memory.execute.executor.util.Batch;
 import cn.edu.tsinghua.iginx.engine.physical.task.TaskResult;
+import cn.edu.tsinghua.iginx.engine.physical.task.utils.PhysicalSupplier;
 import cn.edu.tsinghua.iginx.engine.shared.RequestContext;
 import cn.edu.tsinghua.iginx.engine.shared.data.read.BatchSchema;
 import cn.edu.tsinghua.iginx.engine.shared.data.read.BatchStream;
 import java.util.Objects;
 
 public class StreamSourceMemoryPhysicalTask extends SourceMemoryPhysicalTask {
 
-  private final BatchStream stream;
+  private final PhysicalSupplier<BatchStream> supplier;
 
-  public StreamSourceMemoryPhysicalTask(RequestContext context, Object info, BatchStream stream) {
+  public StreamSourceMemoryPhysicalTask(
+      RequestContext context, Object info, PhysicalSupplier<BatchStream> supplier) {
     super(context, info);
-    this.stream = Objects.requireNonNull(stream);
+    this.supplier = Objects.requireNonNull(supplier);
   }
 
   @Override
   public TaskResult<BatchStream> execute() {
-    return new TaskResult<>(new StreamWrapper(stream));
+    try {
+      return new TaskResult<>(new StreamWrapper(supplier.get()));
+    } catch (PhysicalException e) {
+      return new TaskResult<>(e);
+    }
   }
 
   private class StreamWrapper implements BatchStream {

core/src/main/java/cn/edu/tsinghua/iginx/engine/physical/task/memory/parallel/ParallelPipelineMemoryPhysicalTask.java

@@ -106,7 +106,7 @@ private BatchStream compute(BatchStream previous) {
   private GatherMemoryPhysicalTask constructPipeline(
       RequestContext context, BatchStream stream, ScatterGatherBatchStream outputStream) {
     MemoryPhysicalTask<BatchStream> source =
-        new StreamSourceMemoryPhysicalTask(context, "Parallel Pipeline Source", stream);
+        new StreamSourceMemoryPhysicalTask(context, "Parallel Pipeline Source", () -> stream);
     PipelineMemoryPhysicalTask pipeline = pipelineFactory.createPipeline(context, source);
     GatherMemoryPhysicalTask gather = new GatherMemoryPhysicalTask(context, pipeline, outputStream);
 

core/src/main/java/cn/edu/tsinghua/iginx/engine/physical/task/utils/PhysicalSupplier.java

@@ -0,0 +1,7 @@
+package cn.edu.tsinghua.iginx.engine.physical.task.utils;
+
+import cn.edu.tsinghua.iginx.engine.physical.exception.PhysicalException;
+
+public interface PhysicalSupplier<T extends PhysicalCloseable> {
+  T get() throws PhysicalException;
+}

optimizer/src/main/java/cn/edu/tsinghua/iginx/physical/optimizer/naive/NaivePhysicalPlanner.java

@@ -152,8 +152,7 @@ public PhysicalTask<?> fetch(Source source, RequestContext context) {
     }
   }
 
-  public PhysicalTask<BatchStream> construct(
-      List<Expression> constantExpressions, RequestContext context) {
+  public PhysicalTask<?> construct(List<Expression> constantExpressions, RequestContext context) {
     Preconditions.checkArgument(
         constantExpressions.stream()
             .allMatch(cn.edu.tsinghua.iginx.sql.utils.ExpressionUtils::isConstantArithmeticExpr));
@@ -162,7 +161,7 @@ public PhysicalTask<BatchStream> construct(
         new StreamSourceMemoryPhysicalTask(
             context,
             "Produce 1 emtpy row to calculate constant values",
-            BatchStreams.nonColumn(context.getBatchRowCount(), 1));
+            () -> BatchStreams.nonColumn(context.getBatchRowCount(), 1));
 
     List<FunctionCall> functionCalls = new ArrayList<>();
     for (Expression expression : constantExpressions) {
@@ -172,6 +171,12 @@ public PhysicalTask<BatchStream> construct(
 
     RowTransform rowTransform = new RowTransform(EmptySource.EMPTY_SOURCE, functionCalls);
 
+    if (rowTransform.getFunctionCallList().stream()
+        .anyMatch(UDFDetector::containNonSystemFunction)) {
+      return new UnaryRowMemoryPhysicalTask(
+          convert(source, context, RowStream.class), rowTransform, context);
+    }
+
     return new PipelineMemoryPhysicalTask(
         source,
         Collections.singletonList(rowTransform),