diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java index c1ab9869a01..c0d66f95901 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -10,6 +10,7 @@ import java.sql.Connection; import java.util.function.BiFunction; import lombok.Getter; +import lombok.Setter; import org.apache.calcite.rex.RexNode; import org.apache.calcite.tools.FrameworkConfig; import org.apache.calcite.tools.RelBuilder; @@ -23,7 +24,7 @@ public class CalcitePlanContext { public final RelBuilder relBuilder; public final ExtendedRexBuilder rexBuilder; - @Getter private boolean isResolvingJoinCondition = false; + @Getter @Setter private boolean isResolvingJoinCondition = false; private CalcitePlanContext(FrameworkConfig config) { this.config = config; diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java index 7516b2bd5d3..c1fefcae031 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java @@ -160,18 +160,12 @@ public RexNode visitQualifiedName(QualifiedName node, CalcitePlanContext context if (parts.size() == 1) { // Handle the case of `id = cid` try { - // TODO what if there is join clause in InSubquery in join condition - // for subquery in join condition - return context.relBuilder.field(parts.get(0)); - } catch (IllegalArgumentException e) { - try { - return context.relBuilder.field(2, 0, parts.get(0)); - } catch (IllegalArgumentException ee) { - return context.relBuilder.field(2, 1, parts.get(0)); - } + return context.relBuilder.field(2, 0, parts.getFirst()); + } catch (IllegalArgumentException ee) { + return context.relBuilder.field(2, 1, parts.getFirst()); } - } else if (parts.size() - == 2) { // Handle the case of `t1.id = t2.id` or `alias1.id = alias2.id` + } else if (parts.size() == 2) { + // Handle the case of `t1.id = t2.id` or `alias1.id = alias2.id` return context.relBuilder.field(2, parts.get(0), parts.get(1)); } else if (parts.size() == 3) { throw new UnsupportedOperationException("Unsupported qualified name: " + node); @@ -183,7 +177,7 @@ public RexNode visitQualifiedName(QualifiedName node, CalcitePlanContext context return context.relBuilder.field(qualifiedName); } else if (node.getParts().size() == 2) { List parts = node.getParts(); - return context.relBuilder.field(1, parts.get(0), parts.get(1)); + return context.relBuilder.field(parts.get(0), parts.get(1)); } else if (currentFields.stream().noneMatch(f -> f.startsWith(qualifiedName))) { return context.relBuilder.field(qualifiedName); } @@ -260,9 +254,20 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) { @Override public RexNode visitInSubquery(InSubquery node, CalcitePlanContext context) { List nodes = node.getChild().stream().map(child -> analyze(child, context)).toList(); + // clear and store the outer state + boolean isResolvingJoinConditionOuter = context.isResolvingJoinCondition(); + if (isResolvingJoinConditionOuter) { + context.setResolvingJoinCondition(false); + } UnresolvedPlan subquery = node.getQuery(); + RelNode subqueryRel = subquery.accept(planVisitor, context); + // pop the inner plan context.relBuilder.build(); + // restore to the previous state + if (isResolvingJoinConditionOuter) { + context.setResolvingJoinCondition(true); + } try { return context.relBuilder.in(subqueryRel, nodes); // TODO