[CALCITE-6200] RelJson throw UnsupportedOperationException for RexDyn…

…amicParam
QAQ-Jun · Jan 16, 2024 · 24ea6be · 24ea6be
1 parent e71d59f
commit 24ea6be
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 0 deletions.
diff --git a/core/src/main/java/org/apache/calcite/rel/externalize/RelJson.java b/core/src/main/java/org/apache/calcite/rel/externalize/RelJson.java
@@ -40,6 +40,7 @@
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.rex.RexCorrelVariable;
+import org.apache.calcite.rex.RexDynamicParam;
 import org.apache.calcite.rex.RexFieldAccess;
 import org.apache.calcite.rex.RexFieldCollation;
 import org.apache.calcite.rex.RexLiteral;
@@ -578,6 +579,13 @@ private static Object toJson(CorrelationId node) {
   public Object toJson(RexNode node) {
     final Map<String, @Nullable Object> map;
     switch (node.getKind()) {
+    case DYNAMIC_PARAM:
+      map = jsonBuilder().map();
+      final RexDynamicParam rexDynamicParam = (RexDynamicParam) node;
+      final RelDataType rdpType = rexDynamicParam.getType();
+      map.put("dynamicParam", rexDynamicParam.getIndex());
+      map.put("type", toJson(rdpType));
+      return map;
     case FIELD_ACCESS:
       map = jsonBuilder().map();
       final RexFieldAccess fieldAccess = (RexFieldAccess) node;
@@ -830,6 +838,12 @@ public RexNode toRex(RelOptCluster cluster, Object o) {
         Sarg sarg = sargFromJson((Map) sargObject);
         return rexBuilder.makeSearchArgumentLiteral(sarg, type);
       }
+      if (map.containsKey("dynamicParam")) {
+        final Object dynamicParamObject = requireNonNull(map.get("dynamicParam"));
+        final Integer index = (Integer) dynamicParamObject;
+        final RelDataType type = toType(typeFactory, get(map, "type"));
+        return rexBuilder.makeDynamicParam(type, index);
+      }
       throw new UnsupportedOperationException("cannot convert to rex " + o);
     } else if (o instanceof Boolean) {
       return rexBuilder.makeLiteral((Boolean) o);

diff --git a/core/src/test/java/org/apache/calcite/plan/RelWriterTest.java b/core/src/test/java/org/apache/calcite/plan/RelWriterTest.java
@@ -44,6 +44,7 @@
 import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexCorrelVariable;
+import org.apache.calcite.rex.RexDynamicParam;
 import org.apache.calcite.rex.RexFieldCollation;
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexLiteral;
@@ -1426,6 +1427,54 @@ private RelNode mockCountOver(RelBuilder builder, String table,
         .assertThatPlan(isLinux(expected));
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6200">[CALCITE-6200]
+   * RelJson throw UnsupportedOperationException for RexDynamicParam</a>. */
+  @Test void testDynamicParam() {
+    final Function<RelBuilder, RelNode> relFn = relBuilder -> {
+      final RexBuilder rexBuilder = relBuilder.getRexBuilder();
+      final RelDataTypeFactory typeFactory = relBuilder.getTypeFactory();
+      final RelDataType intType = typeFactory.createSqlType(SqlTypeName.INTEGER);
+      final RexDynamicParam dynamicParam = rexBuilder.makeDynamicParam(intType, 0);
+      final RelNode relNode = relBuilder
+          .scan("EMP")
+          .sortLimit(null, dynamicParam, relBuilder.fields(RelCollations.EMPTY))
+          .build();
+      return relNode;
+    };
+
+    final String expectedJson = "{\n"
+        + "  \"rels\": [\n"
+        + "    {\n"
+        + "      \"id\": \"0\",\n"
+        + "      \"relOp\": \"LogicalTableScan\",\n"
+        + "      \"table\": [\n"
+        + "        \"scott\",\n"
+        + "        \"EMP\"\n"
+        + "      ],\n"
+        + "      \"inputs\": []\n"
+        + "    },\n"
+        + "    {\n"
+        + "      \"id\": \"1\",\n"
+        + "      \"relOp\": \"LogicalSort\",\n"
+        + "      \"collation\": [],\n"
+        + "      \"fetch\": {\n"
+        + "        \"dynamicParam\": 0,\n"
+        + "        \"type\": {\n"
+        + "          \"type\": \"INTEGER\",\n"
+        + "          \"nullable\": false\n"
+        + "        }\n"
+        + "      }\n"
+        + "    }\n"
+        + "  ]\n"
+        + "}";
+    final String expectedPlan = "LogicalSort(fetch=[?0])\n"
+        + "  LogicalTableScan(table=[[scott, EMP]])\n";
+    relFn(relFn)
+        .assertThatJson(is(expectedJson))
+        .assertThatPlan(isLinux(expectedPlan));
+  }
+
   @Test void testWriteSortExchangeWithHashDistribution() {
     final Function<RelBuilder, RelNode> relFn = b ->
         createSortPlan(b, RelDistributions.hash(Lists.newArrayList(0)));