Wip serializing Sarg

looker-open-source · Mar 30, 2023 · de86cf4 · de86cf4
1 parent e3e5c77
commit de86cf4
Show file tree

Hide file tree

Showing 3 changed files with 107 additions and 3 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
@@ -16,6 +16,8 @@
  */
 package org.apache.calcite.rel.externalize;
 
+import com.google.common.collect.Range;
+import com.google.common.collect.RangeSet;
 import org.apache.calcite.avatica.AvaticaUtils;
 import org.apache.calcite.avatica.util.TimeUnit;
 import org.apache.calcite.plan.RelOptCluster;
@@ -65,6 +67,7 @@
 import org.apache.calcite.util.ImmutableBitSet;
 import org.apache.calcite.util.ImmutableIntList;
 import org.apache.calcite.util.JsonBuilder;
+import org.apache.calcite.util.Sarg;
 import org.apache.calcite.util.Util;
 
 import com.google.common.collect.ImmutableList;
@@ -466,6 +469,36 @@ public Object toJson(AggregateCall node) {
     }
   }
 
+  private Object toJson(Sarg node) {
+    final Map<String, @Nullable Object> map = jsonBuilder().map();
+    map.put("isComplementedPoints", node.isComplementedPoints());
+    map.put("pointCount", node.pointCount);
+    map.put("isAll", node.isAll());
+    map.put("isNone", node.isNone());
+    map.put("isPoints", node.isPoints());
+    map.put("rangeSet", toJson(node.rangeSet));
+    return map;
+  }
+
+  private <C extends Comparable<C>> Object toJson(RangeSet<C> rangeSet){
+    final List<@Nullable Object> list = jsonBuilder().list();
+    for (Range<C> o : rangeSet.asRanges()) {
+      list.add(toJson(o));
+    }
+    return list;
+  }
+
+  private <C extends Comparable<C>> Object toJson(Range<C> range){
+    final Map<String, @Nullable Object> map = jsonBuilder().map();
+    map.put("lowerEndpoint", range.lowerEndpoint());
+    map.put("lowerEndpointType", range.lowerEndpoint().getClass());
+    map.put("lowerBoundType", range.lowerBoundType().toString());
+    map.put("upperEndpoint", range.upperEndpoint());
+    map.put("upperEndpointType", range.upperEndpoint().getClass());
+    map.put("upperBoundType", range.upperBoundType().toString());
+    return map;
+  }
+
   private Object toJson(RelDataType node) {
     final Map<String, @Nullable Object> map = jsonBuilder().map();
     if (node.isStruct()) {
@@ -530,7 +563,11 @@ private Object toJson(RexNode node) {
       final RexLiteral literal = (RexLiteral) node;
       final Object value = literal.getValue3();
       map = jsonBuilder().map();
-      map.put("literal", RelEnumTypes.fromEnum(value));
+      if (((RexLiteral) node).getTypeName().getName().equalsIgnoreCase(Sarg.class.getSimpleName())) {
+        map.put("sargLiteral", toJson((Sarg)value));
+      } else {
+        map.put("literal", RelEnumTypes.fromEnum(value));
+      }
       map.put("type", toJson(node.getType()));
       return map;
     case INPUT_REF:
@@ -743,6 +780,15 @@ public RexNode toRex(RelOptCluster cluster, Object o) {
         }
         return rexBuilder.makeLiteral(literal, type);
       }
+      if (map.containsKey("sargLiteral")) {
+        Object sarg = map.get("sargLiteral");
+        if (sarg == null) {
+          final RelDataType type = toType(typeFactory, get(map, "type"));
+          return rexBuilder.makeNullLiteral(type);
+        }
+        final RelDataType type = toType(typeFactory, get(map, "type"));
+        return rexBuilder.makeSearchArgumentLiteral((Sarg) sarg, type);
+      }
       throw new UnsupportedOperationException("cannot convert to rex " + o);
     } else if (o instanceof Boolean) {
       return rexBuilder.makeLiteral((Boolean) o);

diff --git a/core/src/main/java/org/apache/calcite/util/JsonBuilder.java b/core/src/main/java/org/apache/calcite/util/JsonBuilder.java
@@ -16,6 +16,9 @@
  */
 package org.apache.calcite.util;
 
+import com.google.common.collect.ImmutableRangeSet;
+import com.google.common.collect.Range;
+import com.google.common.collect.RangeSet;
 import org.apache.calcite.avatica.util.Spaces;
 
 import com.google.common.collect.ImmutableList;
@@ -129,7 +132,10 @@ public void append(StringBuilder buf, int indent, @Nullable Object o) {
       appendList(buf, indent, (List<?>) o);
     } else if (o instanceof String) {
       appendString(buf, (String) o);
-    } else {
+    } else if (o instanceof Sarg) {
+      appendSarg(buf, indent, (Sarg) o);
+    }
+    else {
       assert o instanceof Number || o instanceof Boolean;
       buf.append(o);
     }
@@ -153,6 +159,34 @@ private static void appendString(StringBuilder buf, String s) {
     buf.append('"');
   }
 
+  private static <C extends Comparable<C>> void appendSarg(StringBuilder buf, int indent, Sarg<C> sarg) {
+    buf.append("{");
+    newline(buf, indent + 1);
+    buf.append("\"pointCount\"");
+    buf.append(": ");
+    buf.append(sarg.pointCount);
+    buf.append(",");
+    newline(buf, indent + 1);
+
+    buf.append("\"nullAs\"");
+    buf.append(": ");
+    buf.append(sarg.nullAs);
+    buf.append(",");
+    newline(buf, indent + 1);
+
+
+    RangeSets.forEach(sarg.rangeSet, RangeSets.printer(buf, StringBuilder::append));
+
+    // if (sarg.isComplementedPoints()){
+    //   sarg.com
+    // }
+    // for (Range<C> range: sarg.rangeSet.asRanges()){
+    //   range.lo
+    // }
+    newline(buf, indent);
+    buf.append("}");
+  }
+
   private void appendMap(
       StringBuilder buf, int indent, Map<String, @Nullable Object> map) {
     if (map.isEmpty()) {

diff --git a/core/src/test/java/org/apache/calcite/plan/RelWriterTest.java b/core/src/test/java/org/apache/calcite/plan/RelWriterTest.java
@@ -856,9 +856,33 @@ private static RexNode translateInput(RelJson relJson, int input,
         b.literal(45),
         b.literal(20),
         b.literal(30));
+    RexNode inNode = b.getRexBuilder().makeIn(
+        b.literal(12),
+        ImmutableList.of(
+          b.literal(20),
+          b.literal(14)));
+
+
     RelJson relJson = RelJson.create().withJsonBuilder(new JsonBuilder());
-    Object rexified = relJson.toJson(between);
+
+    Object rexified = relJson.toJson(inNode);
     RexNode deserialize = relJson.toRex(b.getCluster(), rexified);
+
+    final ObjectMapper mapper = new ObjectMapper();
+    final TypeReference<LinkedHashMap<String, Object>> typeRef =
+        new TypeReference<LinkedHashMap<String, Object>>() {
+        };
+    try {
+      final Map<String, Object> o;
+      String test = mapper.writeValueAsString(rexified);
+      o = mapper
+          .configure(DeserializationFeature.USE_BIG_DECIMAL_FOR_FLOATS, true)
+          .readValue(test, typeRef);
+      System.out.println(o);
+    } catch (JsonProcessingException e) {
+      throw TestUtil.rethrow(e);
+    }
+
     assertThat(deserialize.hashCode(), is(between.hashCode()));
   }