content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to you under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License.  You may obtain a copy of the License at
   *
   * http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.calcite.rel.externalize;
  
  import org.apache.calcite.plan.Convention;
  import org.apache.calcite.plan.RelOptCluster;
  import org.apache.calcite.plan.RelOptSchema;
  import org.apache.calcite.plan.RelOptTable;
  import org.apache.calcite.plan.RelTraitSet;
  import org.apache.calcite.rel.RelCollation;
  import org.apache.calcite.rel.RelCollations;
  import org.apache.calcite.rel.RelDistribution;
  import org.apache.calcite.rel.RelInput;
  import org.apache.calcite.rel.RelNode;
  import org.apache.calcite.rel.core.AggregateCall;
  import org.apache.calcite.rel.type.RelDataType;
  import org.apache.calcite.rel.type.RelDataTypeFactory;
  import org.apache.calcite.rex.RexLiteral;
  import org.apache.calcite.rex.RexNode;
  import org.apache.calcite.runtime.SqlFunctions;
  import org.apache.calcite.schema.Schema;
  import org.apache.calcite.sql.SqlAggFunction;
  import org.apache.calcite.util.ImmutableBitSet;
  import org.apache.calcite.util.Pair;
  import org.apache.calcite.util.Util;
  
  import com.fasterxml.jackson.core.type.TypeReference;
  import com.fasterxml.jackson.databind.DeserializationFeature;
  import com.fasterxml.jackson.databind.ObjectMapper;
  import com.google.common.collect.ImmutableList;
  
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  import java.io.IOException;
  import java.lang.reflect.Constructor;
  import java.lang.reflect.InvocationTargetException;
  import java.math.BigDecimal;
  import java.util.AbstractList;
  import java.util.ArrayList;
  import java.util.LinkedHashMap;
  import java.util.List;
  import java.util.Locale;
  import java.util.Map;
  import java.util.function.UnaryOperator;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Reads a JSON plan and converts it back to a tree of relational expressions.
   *
   * @see org.apache.calcite.rel.RelInput
   */
  public class RelJsonReader {
    private static final TypeReference<LinkedHashMap<String, Object>> TYPE_REF =
        new TypeReference<LinkedHashMap<String, Object>>() {
        };
  
    private final RelOptCluster cluster;
    private final RelOptSchema relOptSchema;
    private final RelJson relJson;
    private final Map<String, RelNode> relMap = new LinkedHashMap<>();
    private @Nullable RelNode lastRel;
  
    public RelJsonReader(RelOptCluster cluster, RelOptSchema relOptSchema,
        Schema schema) {
      this(cluster, relOptSchema, schema, UnaryOperator.identity());
    }
  
    public RelJsonReader(RelOptCluster cluster, RelOptSchema relOptSchema,
        Schema schema, UnaryOperator<RelJson> relJsonTransform) {
      this.cluster = cluster;
      this.relOptSchema = relOptSchema;
      Util.discard(schema);
      relJson = relJsonTransform.apply(RelJson.create());
    }
  
    public RelNode read(String s) throws IOException {
      lastRel = null;
      final ObjectMapper mapper = new ObjectMapper();
      Map<String, Object> o = mapper
          .configure(DeserializationFeature.USE_BIG_DECIMAL_FOR_FLOATS, true)
          .readValue(s, TYPE_REF);
      @SuppressWarnings("unchecked")
      final List<Map<String, Object>> rels = (List) requireNonNull(o.get("rels"), "rels");
      readRels(rels);
      return requireNonNull(lastRel, "lastRel");
    }
  
    /** Converts a JSON string (such as that produced by
     * {@link RelJson#toJson(Object)}) into a Calcite type. */
    public static RelDataType readType(RelDataTypeFactory typeFactory, String s)
        throws IOException {
      final ObjectMapper mapper = new ObjectMapper();
      Map<String, Object> o = mapper
          .configure(DeserializationFeature.USE_BIG_DECIMAL_FOR_FLOATS, true)
          .readValue(s, TYPE_REF);
      return RelJson.create().toType(typeFactory, o);
    }
  
    /** Converts a JSON string (such as that produced by
     * {@link RelJson#toJson(RexNode)}) into a Calcite expression. */
    public static RexNode readRex(RelOptCluster typeFactory, String s)
        throws IOException {
      final ObjectMapper mapper = new ObjectMapper();
      Map<String, Object> o = mapper
          .configure(DeserializationFeature.USE_BIG_DECIMAL_FOR_FLOATS, true)
          .readValue(s, TYPE_REF);
      return RelJson.create().toRex(typeFactory, o);
    }
  
    private void readRels(List<Map<String, Object>> jsonRels) {
      for (Map<String, Object> jsonRel : jsonRels) {
        readRel(jsonRel);
      }
    }
  
    private void readRel(final Map<String, Object> jsonRel) {
      String id = (String) requireNonNull(jsonRel.get("id"), "jsonRel.id");
      String type = (String) requireNonNull(jsonRel.get("relOp"), "jsonRel.relOp");
      Constructor constructor = relJson.getConstructor(type);
      RelInput input = new RelInput() {
        @Override public RelOptCluster getCluster() {
          return cluster;
        }
  
        @Override public RelTraitSet getTraitSet() {
          return cluster.traitSetOf(Convention.NONE);
        }
  
        @Override public RelOptTable getTable(String table) {
          final List<String> list =
              requireNonNull(getStringList(table),
                  () -> "getStringList for " + table);
          return requireNonNull(relOptSchema.getTableForMember(list),
              () -> "table " + table + " is not found in schema " + relOptSchema);
        }
  
        @Override public RelNode getInput() {
          final List<RelNode> inputs = getInputs();
          assert inputs.size() == 1;
          return inputs.get(0);
        }
  
        @Override public List<RelNode> getInputs() {
          final List<String> jsonInputs = getStringList("inputs");
          if (jsonInputs == null) {
            return ImmutableList.of(requireNonNull(lastRel, "lastRel"));
          }
          final ImmutableList.Builder<RelNode> inputs = new ImmutableList.Builder<>();
          for (String jsonInput : jsonInputs) {
            inputs.add(lookupInput(jsonInput));
          }
          return inputs.build();
        }
  
        @Override public @Nullable RexNode getExpression(String tag) {
          return relJson.toRex(this, jsonRel.get(tag));
        }
  
        @Override public ImmutableBitSet getBitSet(String tag) {
          return ImmutableBitSet.of(requireNonNull(getIntegerList(tag), tag));
        }
  
        @Override public @Nullable List<ImmutableBitSet> getBitSetList(String tag) {
          List<List<Integer>> list = getIntegerListList(tag);
          if (list == null) {
            return null;
          }
          final ImmutableList.Builder<ImmutableBitSet> builder =
              ImmutableList.builder();
          for (List<Integer> integers : list) {
            builder.add(ImmutableBitSet.of(integers));
          }
          return builder.build();
        }
  
        @Override public @Nullable List<String> getStringList(String tag) {
          //noinspection unchecked
          return (List<String>) jsonRel.get(tag);
        }
  
        @Override public @Nullable List<Integer> getIntegerList(String tag) {
          //noinspection unchecked
          return (List<Integer>) jsonRel.get(tag);
        }
  
        @Override public @Nullable List<List<Integer>> getIntegerListList(String tag) {
          //noinspection unchecked
          return (List<List<Integer>>) jsonRel.get(tag);
        }
  
        @Override public List<AggregateCall> getAggregateCalls(String tag) {
          @SuppressWarnings("unchecked")
          final List<Map<String, Object>> jsonAggs = (List) getNonNull(tag);
          final List<AggregateCall> inputs = new ArrayList<>();
          for (Map<String, Object> jsonAggCall : jsonAggs) {
            inputs.add(toAggCall(jsonAggCall));
          }
          return inputs;
        }
  
        @Override public @Nullable Object get(String tag) {
          return jsonRel.get(tag);
        }
  
        private Object getNonNull(String tag) {
          return requireNonNull(get(tag), () -> "no entry for tag " + tag);
        }
  
        @Override public @Nullable String getString(String tag) {
          return (String) get(tag);
        }
  
        @Override public float getFloat(String tag) {
          return ((Number) getNonNull(tag)).floatValue();
        }
  
        @Override public BigDecimal getBigDecimal(String tag) {
          return SqlFunctions.toBigDecimal(getNonNull(tag));
        }
  
        @Override public boolean getBoolean(String tag, boolean default_) {
          final Boolean b = (Boolean) get(tag);
          return b != null ? b : default_;
        }
  
        @Override public <E extends Enum<E>> @Nullable E getEnum(String tag, Class<E> enumClass) {
          return Util.enumVal(enumClass,
              ((String) getNonNull(tag)).toUpperCase(Locale.ROOT));
        }
  
        @Override public @Nullable List<RexNode> getExpressionList(String tag) {
          @SuppressWarnings("unchecked")
          final List<Object> jsonNodes = (List) jsonRel.get(tag);
          if (jsonNodes == null) {
            return null;
          }
          final List<RexNode> nodes = new ArrayList<>();
          for (Object jsonNode : jsonNodes) {
            nodes.add(relJson.toRex(this, jsonNode));
          }
          return nodes;
        }
  
        @Override public RelDataType getRowType(String tag) {
          final Object o = getNonNull(tag);
          return relJson.toType(cluster.getTypeFactory(), o);
        }
  
        @Override public RelDataType getRowType(String expressionsTag, String fieldsTag) {
          final List<RexNode> expressionList = getExpressionList(expressionsTag);
          @SuppressWarnings("unchecked") final List<String> names =
              (List<String>) getNonNull(fieldsTag);
          return cluster.getTypeFactory().createStructType(
              new AbstractList<Map.Entry<String, RelDataType>>() {
                @Override public Map.Entry<String, RelDataType> get(int index) {
                  return Pair.of(names.get(index),
                      requireNonNull(expressionList, "expressionList").get(index).getType());
                }
  
                @Override public int size() {
                  return names.size();
                }
              });
        }
  
        @Override public RelCollation getCollation() {
          //noinspection unchecked
          return relJson.toCollation((List) getNonNull("collation"));
        }
  
        @Override public RelDistribution getDistribution() {
          //noinspection unchecked
          return relJson.toDistribution((Map<String, Object>) getNonNull("distribution"));
        }
  
        @Override public ImmutableList<ImmutableList<RexLiteral>> getTuples(String tag) {
          //noinspection unchecked
          final List<List> jsonTuples = (List) getNonNull(tag);
          final ImmutableList.Builder<ImmutableList<RexLiteral>> builder =
              ImmutableList.builder();
          for (List jsonTuple : jsonTuples) {
            builder.add(getTuple(jsonTuple));
          }
          return builder.build();
        }
  
        public ImmutableList<RexLiteral> getTuple(List jsonTuple) {
          final ImmutableList.Builder<RexLiteral> builder =
              ImmutableList.builder();
          for (Object jsonValue : jsonTuple) {
            builder.add((RexLiteral) relJson.toRex(this, jsonValue));
          }
          return builder.build();
        }
      };
      try {
        final RelNode rel = (RelNode) constructor.newInstance(input);
        relMap.put(id, rel);
        lastRel = rel;
      } catch (InstantiationException | IllegalAccessException e) {
        throw new RuntimeException(e);
      } catch (InvocationTargetException e) {
        final Throwable e2 = e.getCause();
        if (e2 instanceof RuntimeException) {
          throw (RuntimeException) e2;
        }
        throw new RuntimeException(e2);
      }
    }
  
    private AggregateCall toAggCall(Map<String, Object> jsonAggCall) {
      @SuppressWarnings("unchecked")
      final Map<String, Object> aggMap =
          (Map) requireNonNull(jsonAggCall.get("agg"),
              "agg key is not found");
      final SqlAggFunction aggregation =
          requireNonNull(relJson.toAggregation(aggMap),
              () -> "relJson.toAggregation output for " + aggMap);
      final boolean distinct =
          requireNonNull((Boolean) jsonAggCall.get("distinct"),
              "jsonAggCall.distinct");
      @SuppressWarnings("unchecked")
      final List<Integer> operands =
          requireNonNull((List<Integer>) jsonAggCall.get("operands"),
              "jsonAggCall.operands");
      final Integer filterOperand = (Integer) jsonAggCall.get("filter");
      final Object jsonAggType =
          requireNonNull(jsonAggCall.get("type"), "jsonAggCall.type");
      final RelDataType type =
          relJson.toType(cluster.getTypeFactory(), jsonAggType);
      final String name = (String) jsonAggCall.get("name");
      return AggregateCall.create(aggregation, distinct, false, false,
          ImmutableList.of(), operands,
          filterOperand == null ? -1 : filterOperand,
          null, RelCollations.EMPTY, type, name);
    }
  
    private RelNode lookupInput(String jsonInput) {
      RelNode node = relMap.get(jsonInput);
      if (node == null) {
        throw new RuntimeException("unknown id " + jsonInput
            + " for relational expression");
      }
      return node;
    }
  }
  