org.apache.calcite.DataContext Java Examples

private static Function1<DataContext, Object[]> compile(String code,
    Object reason) {
  try {
    final ClassBodyEvaluator cbe = new ClassBodyEvaluator();
    cbe.setClassName(GENERATED_CLASS_NAME);
    cbe.setExtendedClass(Utilities.class);
    cbe.setImplementedInterfaces(new Class[] {Function1.class, Serializable.class});
    cbe.setParentClassLoader(RexExecutable.class.getClassLoader());
    cbe.cook(new Scanner(null, new StringReader(code)));
    Class c = cbe.getClazz();
    //noinspection unchecked
    final Constructor<Function1<DataContext, Object[]>> constructor =
        c.getConstructor();
    return constructor.newInstance();
  } catch (CompileException | IOException | InstantiationException
      | IllegalAccessException | InvocationTargetException
      | NoSuchMethodException e) {
    throw new RuntimeException("While compiling " + reason, e);
  }
}
 

private boolean isSatisfiable(RexNode second, DataContext dataValues) {
  if (dataValues == null) {
    return false;
  }

  ImmutableList<RexNode> constExps = ImmutableList.of(second);
  final RexExecutable exec = RexExecutorImpl.getExecutable(builder, constExps, rowType);

  Object[] result;
  exec.setDataContext(dataValues);
  try {
    result = exec.execute();
  } catch (Exception e) {
    // TODO: CheckSupport should not allow this exception to be thrown
    // Need to monitor it and handle all the cases raising them.
    LOGGER.warn("Exception thrown while checking if => {}: {}", second, e.getMessage());
    return false;
  }
  return result != null
      && result.length == 1
      && result[0] instanceof Boolean
      && (Boolean) result[0];
}
 

protected MetaResultSet createResultSet(
    Map<String, Object> internalParameters, List<ColumnMetaData> columns,
    CursorFactory cursorFactory, final Frame firstFrame) {
  try {
    final QuarkConnectionImpl connection = getConnection();
    final AvaticaStatement statement = connection.createStatement();
    final CalcitePrepare.CalciteSignature<Object> signature =
        new CalcitePrepare.CalciteSignature<Object>("",
            ImmutableList.<AvaticaParameter>of(), internalParameters, null,
            columns, cursorFactory, null, ImmutableList.<RelCollation>of(), -1,
            null, Meta.StatementType.SELECT) {
          @Override public Enumerable<Object> enumerable(
              DataContext dataContext) {
            return Linq4j.asEnumerable(firstFrame.rows);
          }
        };
    return MetaResultSet.create(connection.id, statement.getId(), true,
        signature, firstFrame);
  } catch (SQLException e) {
    throw new RuntimeException(e);
  }
}
 

public static DataContext of(RelDataType rowType,
    List<Pair<RexInputRef, RexNode>> usageList) {
  final int size = rowType.getFieldList().size();
  final Object[] values = new Object[size];
  for (Pair<RexInputRef, RexNode> elem : usageList) {
    Pair<Integer, ?> value = getValue(elem.getKey(), elem.getValue());
    if (value == null) {
      LOGGER.warn("{} is not handled for {} for checking implication",
          elem.getKey(), elem.getValue());
      return null;
    }
    int index = value.getKey();
    values[index] = value.getValue();
  }
  return new VisitorDataContext(values);
}
 

private void bindVariable(TupleFilter filter, DataContext dataContext) {
    if (filter == null) {
        return;
    }

    for (TupleFilter childFilter : filter.getChildren()) {
        bindVariable(childFilter, dataContext);
    }

    if (filter instanceof CompareTupleFilter && dataContext != null) {
        CompareTupleFilter compFilter = (CompareTupleFilter) filter;
        for (Map.Entry<String, Object> entry : compFilter.getVariables().entrySet()) {
            String variable = entry.getKey();
            Object value = dataContext.get(variable);
            if (value != null) {
                String str = value.toString();
                str = transferDateTimeColumnToMillis(compFilter, str);
                compFilter.clearPreviousVariableValues(variable);
                compFilter.bindVariable(variable, str);
            }

        }
    }
}
 

private static Function1<DataContext, Object[]> compile(String code,
    Object reason) {
  try {
    final ClassBodyEvaluator cbe = new ClassBodyEvaluator();
    cbe.setClassName(GENERATED_CLASS_NAME);
    cbe.setExtendedClass(Utilities.class);
    cbe.setImplementedInterfaces(new Class[] {Function1.class, Serializable.class});
    cbe.setParentClassLoader(RexExecutable.class.getClassLoader());
    cbe.cook(new Scanner(null, new StringReader(code)));
    Class c = cbe.getClazz();
    //noinspection unchecked
    final Constructor<Function1<DataContext, Object[]>> constructor =
        c.getConstructor();
    return constructor.newInstance();
  } catch (CompileException | IOException | InstantiationException
      | IllegalAccessException | InvocationTargetException
      | NoSuchMethodException e) {
    throw new RuntimeException("While compiling " + reason, e);
  }
}
 

public static JdbcCatalogSchema create(
    SchemaPlus parentSchema,
    String name,
    DataSource dataSource,
    SqlDialectFactory dialectFactory,
    String catalog) {
  final Expression expression =
      parentSchema != null
          ? Schemas.subSchemaExpression(parentSchema, name,
              JdbcCatalogSchema.class)
          : Expressions.call(DataContext.ROOT,
              BuiltInMethod.DATA_CONTEXT_GET_ROOT_SCHEMA.method);
  final SqlDialect dialect =
      JdbcSchema.createDialect(dialectFactory, dataSource);
  final JdbcConvention convention =
      JdbcConvention.of(dialect, expression, name);
  return new JdbcCatalogSchema(dataSource, dialect, convention, catalog);
}
 

/**
 * Ensures that for a given context operator,
 * the correct value is retrieved from the {@link DataContext}.
 *
 * @param operator The Operator to check
 * @param variable The DataContext variable this operator should be bound to
 * @param value The expected value to retrieve.
 */
private void testContextLiteral(
    final SqlOperator operator,
    final DataContext.Variable variable,
    final Object value) {
  Frameworks.withPrepare((cluster, relOptSchema, rootSchema, statement) -> {
    final RexBuilder rexBuilder = cluster.getRexBuilder();
    final RexExecutorImpl executor =
        new RexExecutorImpl(
            new SingleValueDataContext(variable.camelName, value));
    try {
      checkConstant(value, builder -> {
        final List<RexNode> output = new ArrayList<>();
        executor.reduce(rexBuilder,
            ImmutableList.of(rexBuilder.makeCall(operator)), output);
        return output.get(0);
      });
    } catch (Exception e) {
      throw TestUtil.rethrow(e);
    }
    return null;
  });
}
 

private static CallImplementor createImplementor(final Method method) {
  return RexImpTable.createImplementor(
      new ReflectiveCallNotNullImplementor(method) {
        public Expression implement(RexToLixTranslator translator,
            RexCall call, List<Expression> translatedOperands) {
          Expression expr = super.implement(translator, call,
              translatedOperands);
          final Class<?> returnType = method.getReturnType();
          if (QueryableTable.class.isAssignableFrom(returnType)) {
            Expression queryable = Expressions.call(
                Expressions.convert_(expr, QueryableTable.class),
                BuiltInMethod.QUERYABLE_TABLE_AS_QUERYABLE.method,
                Expressions.call(DataContext.ROOT,
                    BuiltInMethod.DATA_CONTEXT_GET_QUERY_PROVIDER.method),
                Expressions.constant(null, SchemaPlus.class),
                Expressions.constant(call.getOperator().getName(), String.class));
            expr = Expressions.call(queryable,
                BuiltInMethod.QUERYABLE_AS_ENUMERABLE.method);
          } else {
            expr = Expressions.call(expr,
                BuiltInMethod.SCANNABLE_TABLE_SCAN.method, DataContext.ROOT);
          }
          return expr;
        }
      }, NullPolicy.NONE, false);
}
 

public static Enumerable<Object> computeSCALA(DataContext dataContext, RelNode relNode, RelDataType resultType)
        throws IllegalAccessException, ClassNotFoundException, InstantiationException {
    try {
        String property = System.getProperty("kylin-query-engine", "org.apache.kylin.query.runtime.SparkEngine");
        QueryEngine o = (QueryEngine) Class.forName(property).newInstance();
        return o.computeSCALA(dataContext, relNode, resultType);
    } catch (InstantiationException | IllegalAccessException | ClassNotFoundException e) {
        throw e;
    }
}
 

public static DataContext of(RelDataType rowType, RexNode rex) {
  final int size = rowType.getFieldList().size();
  final Object[] values = new Object[size];
  final List<RexNode> operands = ((RexCall) rex).getOperands();
  final RexNode firstOperand = operands.get(0);
  final RexNode secondOperand = operands.get(1);
  final Pair<Integer, ?> value = getValue(firstOperand, secondOperand);
  if (value != null) {
    int index = value.getKey();
    values[index] = value.getValue();
    return new VisitorDataContext(values);
  } else {
    return null;
  }
}
 

@Override public Enumerable<Object[]> scan(DataContext root) {
  return new AbstractEnumerable<Object[]>() {
    public Enumerator<Object[]> enumerator() {
      return new GeodeSimpleEnumerator<Object[]>(clientCache, regionName) {
        @Override public Object[] convert(Object obj) {
          Object values = convertToRowValues(relDataType.getFieldList(), obj);
          if (values instanceof Object[]) {
            return (Object[]) values;
          }
          return new Object[]{values};
        }
      };
    }
  };
}
 

public static Enumerable<Object[]> asyncResult(DataContext dataContext) {
    if (BackdoorToggles.getPrepareOnly()) {
        return Linq4j.emptyEnumerable();
    }
    OLAPRel olapRel = (OLAPRel) QueryContextFacade.current().getOlapRel();
    RelDataType rowType = (RelDataType) QueryContextFacade.current().getResultType();
    return QueryEngineFactory.computeAsync(dataContext, olapRel, rowType);
}
 

public static Enumerable<Object> collectToScalarEnumerable(DataContext dataContext) {
    if (BackdoorToggles.getPrepareOnly()) {
        return Linq4j.emptyEnumerable();
    }

    OLAPRel olapRel = (OLAPRel) QueryContextFacade.current().getOlapRel();
    RelDataType rowType = (RelDataType) QueryContextFacade.current().getResultType();
    try {
        Enumerable<Object> objects = QueryEngineFactory.computeSCALA(dataContext, olapRel, rowType);
        return objects;

    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}
 

public static Enumerable<Object[]> collectToEnumerable(DataContext dataContext) {
    if (BackdoorToggles.getPrepareOnly()) {
        return Linq4j.emptyEnumerable();
    }

    OLAPRel olapRel = (OLAPRel) QueryContextFacade.current().getOlapRel();
    RelDataType rowType = (RelDataType) QueryContextFacade.current().getResultType();
    try {
        Enumerable<Object[]> computer = QueryEngineFactory.compute(dataContext, olapRel, rowType);
        return computer;
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}
 

public Enumerable<Object[]> scan(DataContext root) {
  final List<CsvFieldType> fieldTypes = getFieldTypes(root.getTypeFactory());
  final int[] fields = CsvEnumerator.identityList(fieldTypes.size());
  final AtomicBoolean cancelFlag = DataContext.Variable.CANCEL_FLAG.get(root);
  return new AbstractEnumerable<Object[]>() {
    public Enumerator<Object[]> enumerator() {
      return new CsvEnumerator<>(source, cancelFlag, true, null,
          new CsvEnumerator.ArrayRowConverter(fieldTypes, fields, true));
    }
  };
}
 

public static DataContext of(RelDataType rowType, RexNode rex) {
  final int size = rowType.getFieldList().size();
  final Object[] values = new Object[size];
  final List<RexNode> operands = ((RexCall) rex).getOperands();
  final RexNode firstOperand = operands.get(0);
  final RexNode secondOperand = operands.get(1);
  final Pair<Integer, ?> value = getValue(firstOperand, secondOperand);
  if (value != null) {
    int index = value.getKey();
    values[index] = value.getValue();
    return new VisitorDataContext(values);
  } else {
    return null;
  }
}
 

public OLAPQuery(DataContext optiqContext, EnumeratorTypeEnum type, int ctxId) {
    this.optiqContext = optiqContext;
    this.type = type;
    this.contextId = ctxId;

    QueryContextFacade.current().addContext(ctxId, type.toString(),
            type == EnumeratorTypeEnum.OLAP);
}
 

public static Enumerable<Object[]> compute(DataContext dataContext, RelNode relNode, RelDataType resultType)
        throws IllegalAccessException, ClassNotFoundException, InstantiationException {
    try {
        String property = System.getProperty("kylin-query-engine", "org.apache.kylin.query.runtime.SparkEngine");
        QueryEngine o = (QueryEngine) Class.forName(property).newInstance();
        return o.compute(dataContext, relNode, resultType);
    } catch (InstantiationException | IllegalAccessException | ClassNotFoundException e) {
        throw e;
    }
}
 

@Override
public Expression getExpression(SchemaPlus parentSchema,
                                String name) {
  return Expressions.call(
      DataContext.ROOT,
      BuiltInMethod.DATA_CONTEXT_GET_ROOT_SCHEMA.method);
}
 

public InterpreterImplementor(Compiler compiler,
    CalcitePrepare.SparkHandler spark,
    DataContext dataContext) {
  this.compiler = compiler;
  this.spark = spark;
  this.dataContext = dataContext;
}
 

@Test void testTimestampFromContext() throws Exception {
  // CURRENT_TIMESTAMP actually rounds the value to nearest second
  // and that's why we do currentTimeInMillis / 1000 * 1000
  long val = System.currentTimeMillis() / 1000 * 1000;
  testContextLiteral(SqlStdOperatorTable.CURRENT_TIMESTAMP,
      DataContext.Variable.CURRENT_TIMESTAMP, val);
}
 

public Enumerable<Object[]> scan(DataContext root, List<RexNode> filters,
    final int[] projects) {
  final Pair<Integer, Object> filter = getFilter(cooperative, filters);
  return new AbstractEnumerable<Object[]>() {
    public Enumerator<Object[]> enumerator() {
      return beatles(buf, filter, projects);
    }
  };
}
 

public Enumerable<Object[]> scan(DataContext root) {
  return new AbstractEnumerable<Object[]>() {
    public Enumerator<Object[]> enumerator() {
      return tens();
    }
  };
}
 

public DataContext createDataContext(Map<String, Object> parameterValues,
    CalciteSchema rootSchema) {
  if (config().spark()) {
    return new SlimDataContext();
  }
  return new DataContextImpl(this, parameterValues, rootSchema);
}
 

public Enumerable<Object[]> scan(DataContext root, List<RexNode> filters) {
  final List<CsvFieldType> fieldTypes = getFieldTypes(root.getTypeFactory());
  final String[] filterValues = new String[fieldTypes.size()];
  filters.removeIf(filter -> addFilter(filter, filterValues));
  final int[] fields = CsvEnumerator.identityList(fieldTypes.size());
  final AtomicBoolean cancelFlag = DataContext.Variable.CANCEL_FLAG.get(root);
  return new AbstractEnumerable<Object[]>() {
    public Enumerator<Object[]> enumerator() {
      return new CsvEnumerator<>(source, cancelFlag, false, filterValues,
          new CsvEnumerator.ArrayRowConverter(fieldTypes, fields));
    }
  };
}
 

public static AtomicBoolean getCancelFlag(DataContext root) {
    AtomicBoolean tempFlag = DataContext.Variable.CANCEL_FLAG.get(root);
    if (tempFlag == null) {
        return new AtomicBoolean(false);
    } else {
        return tempFlag;
    }
}
 

public Enumerable<Object[]> scan(DataContext root) {
  return new AbstractEnumerable<Object[]>() {
    public Enumerator<Object[]> enumerator() {
      return new JsonEnumerator(getDataList(root.getTypeFactory()));
    }
  };
}
 

public Enumerable<Object[]> scan(DataContext root) {
  if (elementType == Object[].class) {
    //noinspection unchecked
    return enumerable;
  } else {
    //noinspection unchecked
    return enumerable.select(new FieldSelector((Class) elementType));
  }
}
 

public Enumerable<Object[]> bind(DataContext dataContext) {
  final ImmutableMap<String, Object> map = ImmutableMap.of();
  final Bindable bindable = EnumerableInterpretable.toBindable(map, null,
      (EnumerableRel) getInput(), EnumerableRel.Prefer.ARRAY);
  final ArrayBindable arrayBindable = EnumerableInterpretable.box(bindable);
  return arrayBindable.bind(dataContext);
}