WIP - trying to figure out if we can deserialize correctly

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

Hide file tree

Showing 6 changed files with 165 additions and 5 deletions.
diff --git a/core/src/main/java/org/apache/calcite/rel/externalize/RelEnumTypes.java b/core/src/main/java/org/apache/calcite/rel/externalize/RelEnumTypes.java
@@ -18,6 +18,7 @@
 
 import org.apache.calcite.avatica.util.TimeUnitRange;
 import org.apache.calcite.rel.core.TableModify;
+import org.apache.calcite.rex.RexUnknownAs;
 import org.apache.calcite.sql.JoinConditionType;
 import org.apache.calcite.sql.JoinType;
 import org.apache.calcite.sql.SqlExplain;
@@ -78,6 +79,7 @@ private RelEnumTypes() {}
     register(enumByName, SqlTrimFunction.Flag.class);
     register(enumByName, TimeUnitRange.class);
     register(enumByName, TableModify.Operation.class);
+    register(enumByName, RexUnknownAs.class);
     ENUM_BY_NAME = enumByName.build();
   }
 

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
@@ -65,9 +65,12 @@
 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;
+import com.google.common.collect.Range;
+import com.google.common.collect.RangeSet;
 
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.checkerframework.checker.nullness.qual.PolyNull;
@@ -212,7 +215,7 @@ public Constructor getConstructor(String type) {
    * Converts a type name to a class. E.g. {@code getClass("LogicalProject")}
    * returns {@link org.apache.calcite.rel.logical.LogicalProject}.class.
    */
-  public Class typeNameToClass(String type) {
+  public static Class typeNameToClass(String type) {
     if (!type.contains(".")) {
       for (String package_ : PACKAGES) {
         try {
@@ -466,6 +469,37 @@ 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));
+    map.put("nullAs", node.nullAs);
+    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 +564,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 +781,16 @@ public RexNode toRex(RelOptCluster cluster, Object o) {
         }
         return rexBuilder.makeLiteral(literal, type);
       }
+      if (map.containsKey("sargLiteral")) {
+        Object sargObject = map.get("sargLiteral");
+        if (sargObject == null) {
+          final RelDataType type = toType(typeFactory, get(map, "type"));
+          return rexBuilder.makeNullLiteral(type);
+        }
+        final RelDataType type = toType(typeFactory, get(map, "type"));
+        Sarg sarg = Sarg.fromJson((Map) sargObject);
+        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/rex/RexUnknownAs.java b/core/src/main/java/org/apache/calcite/rex/RexUnknownAs.java
@@ -120,4 +120,17 @@ public RexUnknownAs or(RexUnknownAs other) {
       return other;
     }
   }
+
+  public static RexUnknownAs toEnum(String name){
+    switch(name) {
+    case "TRUE":
+      return RexUnknownAs.TRUE;
+    case "FALSE":
+      return RexUnknownAs.FALSE;
+    case "UNKNOWN":
+      return RexUnknownAs.UNKNOWN;
+    default:
+      throw new IllegalArgumentException("unknown");
+    }
+  }
 }
diff --git a/core/src/main/java/org/apache/calcite/util/RangeSets.java b/core/src/main/java/org/apache/calcite/util/RangeSets.java
@@ -22,7 +22,12 @@
 import com.google.common.collect.RangeSet;
 import com.google.common.collect.TreeRangeSet;
 
+import java.lang.reflect.Constructor;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
@@ -271,6 +276,61 @@ public static <C extends Comparable<C>> void forEach(Range<C> range,
     }
   }
 
+  public static <C extends Comparable<C>> RangeSet<C> fromJson(ArrayList rangeSetsJson)
+  {
+    final ImmutableRangeSet.Builder<C> builder = ImmutableRangeSet.builder();
+    List<Range<C>> rangeList = Collections.emptyList();
+    try {
+      for (Object o : rangeSetsJson) {
+        // TODO: Throw if can't convert to Map
+        Map<String, String> rangeMap = (Map) o;
+        Range range = rangeFromJson(rangeMap);
+        rangeList.add(range);
+      }
+      builder.addAll(rangeList);
+      return builder.build();
+    } catch (Exception e) {
+      throw new RuntimeException("Error creating RangeSet from JSON: ", e);
+    }
+  }
+
+  public static BoundType boundTypeNameToEnum(String name){
+    switch (name){
+    case "OPEN":
+      return BoundType.OPEN;
+    case "CLOSED":
+      return BoundType.CLOSED;
+    default:
+      throw new IllegalArgumentException("Unknown BoundType enum name");
+    }
+  }
+
+  public static <C extends Comparable<C>> Range<C> rangeFromJson(Map<String, String> rangeMap) {
+    try {
+      // Can I cast a variable to a class from a name?
+      // This part is WIP (doesn't work)
+
+      Class<?> clazz = Class.forName(rangeMap.get("lowerEndpointType"));
+      Constructor<?> constructor = clazz.getConstructor();
+      C lower = (C) constructor.newInstance(rangeMap.get("lowerEndpoint"));
+      // C lower = (C) rangeMap.get("lowerEndpoint");
+
+
+      // String upperEndpointType = rangeMap.get("lowerEndpointType");
+      // Class upperEndpointTypeClass = Class.forName(upperEndpointType);
+      // Comparable<C> upper = ((Comparable<C>) rangeMap.get("upperEndPoint"));
+      C upper = lower;
+
+      BoundType lowerBoundType = boundTypeNameToEnum(rangeMap.get("lowerBoundType"));
+      BoundType upperBoundType = boundTypeNameToEnum(rangeMap.get("upperBoundType"));
+
+      Range range = Range.range(lower, lowerBoundType, upper, upperBoundType);
+      return range;
+    } catch (Exception e){
+      throw new RuntimeException("Error creating Range from JSON: ", e);
+    }
+  }
+
   /** Creates a consumer that prints values to a {@link StringBuilder}. */
   public static <C extends Comparable<C>> Consumer<C> printer(StringBuilder sb,
       BiConsumer<StringBuilder, C> valuePrinter) {

diff --git a/core/src/main/java/org/apache/calcite/util/Sarg.java b/core/src/main/java/org/apache/calcite/util/Sarg.java
@@ -16,6 +16,8 @@
  */
 package org.apache.calcite.util;
 
+import java.util.ArrayList;
+
 import org.apache.calcite.linq4j.Ord;
 import org.apache.calcite.rex.RexUnknownAs;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
@@ -27,6 +29,7 @@
 
 import org.checkerframework.checker.nullness.qual.Nullable;
 
+import java.util.Map;
 import java.util.function.BiConsumer;
 
 import static java.util.Objects.requireNonNull;
@@ -153,6 +156,16 @@ public static <C extends Comparable<C>> Sarg<C> of(RexUnknownAs nullAs,
     return new Sarg<>(ImmutableRangeSet.copyOf(rangeSet), nullAs);
   }
 
+  public static <C extends Comparable<C>> Sarg<C> fromJson(Map o){
+    if (Map.class.isAssignableFrom(o.getClass())){
+      RexUnknownAs unknownAs = RexUnknownAs.toEnum(o.get("nullAs").toString());
+      // o.get("rangeSet") is stored as a list of ranges
+      RangeSet rangeSet = RangeSets.fromJson((ArrayList) o.get("rangeSet"));
+      return new Sarg<>(ImmutableRangeSet.copyOf(rangeSet), unknownAs);
+    }
+    throw new RuntimeException("Failed to create Sarg from JSON");
+  }
+
   /**
    * {@inheritDoc}
    *

diff --git a/core/src/test/java/org/apache/calcite/plan/RelWriterTest.java b/core/src/test/java/org/apache/calcite/plan/RelWriterTest.java
@@ -852,13 +852,37 @@ private static RexNode translateInput(RelJson relJson, int input,
     // final String expected = "<TODO>";
     // assertThat(result, isLinux(expected));
 
-    RexNode between = b.getRexBuilder().makeBetween(
-        b.literal(45),
+    RexNode between =
+        b.getRexBuilder().makeBetween(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()));
   }