support lookup

Signed-off-by: Lantao Jin <[email protected]>
LantaoJin · Dec 25, 2024 · 863be2f · 863be2f
1 parent 497bf2c
commit 863be2f
Show file tree

Hide file tree

Showing 12 changed files with 663 additions and 245 deletions.
diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/calcite/CalciteAggCallVisitor.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/calcite/CalciteAggCallVisitor.java
@@ -11,8 +11,7 @@
 import org.opensearch.sql.ast.expression.AggregateFunction;
 import org.opensearch.sql.ast.expression.Alias;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
-
-import static org.opensearch.sql.calcite.CalciteHelper.translate;
+import org.opensearch.sql.calcite.utils.AggregateUtils;
 
 public class CalciteAggCallVisitor extends AbstractNodeVisitor<AggCall, CalcitePlanContext> {
     private final CalciteRexNodeVisitor rexNodeVisitor;
@@ -34,6 +33,6 @@ public AggCall visitAlias(Alias node, CalcitePlanContext context) {
     @Override
     public AggCall visitAggregateFunction(AggregateFunction node, CalcitePlanContext context) {
         RexNode field = rexNodeVisitor.analyze(node.getField(), context);
-        return translate(node, field, context);
+        return AggregateUtils.translate(node, field, context);
     }
 }
diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/calcite/CalciteHelper.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/calcite/CalciteHelper.java
diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java
@@ -6,25 +6,23 @@
 package org.opensearch.sql.calcite;
 
 import lombok.Getter;
-import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.tools.FrameworkConfig;
 import org.apache.calcite.tools.RelBuilder;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 
 import java.util.function.BiFunction;
 
-@Getter
 public class CalcitePlanContext {
 
-    private RelBuilder relBuilder;
-    private RexBuilder rexBuilder;
+    public final RelBuilder relBuilder;
+    public final ExtendedRexBuilder rexBuilder;
 
-    private boolean isResolvingJoinCondition = false;
+    @Getter private boolean isResolvingJoinCondition = false;
 
     public CalcitePlanContext(RelBuilder relBuilder) {
         this.relBuilder = relBuilder;
-        this.rexBuilder = relBuilder.getRexBuilder();
+        this.rexBuilder = new ExtendedRexBuilder(relBuilder.getRexBuilder());
     }
 
     public RexNode resolveJoinCondition(