Examples with DataContext org.apache.calcite.DataContext used on opensource projects

Example 31 with DataContext

use of org.apache.calcite.DataContext in project drill by apache.

the class EnumerableRecordReader method setup.

@SuppressWarnings("unchecked")
private void setup(OperatorContext context) {
    SchemaPlus rootSchema = context.getFragmentContext().getFullRootSchema();
    DataContext root = new DrillDataContext(schemaPath != null ? SchemaUtilites.searchSchemaTree(rootSchema, SchemaUtilites.getSchemaPathAsList(schemaPath)) : rootSchema, new JavaTypeFactoryImpl(), Collections.emptyMap());
    try {
        Class<?> implementationClass = ClassBuilder.getCompiledClass(code, CLASS_NAME, context.getFragmentContext().getConfig(), context.getFragmentContext().getOptions());
        Iterable<?> iterable = (Iterable<Map<String, Object>>) implementationClass.getMethod(BuiltInMethod.BINDABLE_BIND.method.getName(), DataContext.class).invoke(implementationClass.newInstance(), root);
        if (fieldsMap.keySet().size() == 1) {
            // for the case of projecting single column, its value is returned
            records = StreamSupport.stream(iterable.spliterator(), false).map(this::wrap).iterator();
        } else {
            // for the case when all columns were projected, array is returned
            records = StreamSupport.stream(iterable.spliterator(), false).map(row -> wrap((Object[]) row)).iterator();
        }
    } catch (CompileException | IOException | ClassTransformationException | ReflectiveOperationException e) {
        logger.error("Exception happened when executing generated code", e);
        Throwable rootCause = Throwables.getRootCause(e);
        throw new DrillRuntimeException(rootCause.getMessage(), rootCause);
    }
}

Example 32 with DataContext

use of org.apache.calcite.DataContext in project calcite by apache.

the class Prepare method optimize.

/**
 * Optimizes a query plan.
 *
 * @param root Root of relational expression tree
 * @param materializations Tables known to be populated with a given query
 * @param lattices Lattices
 * @return an equivalent optimized relational expression
 */
protected RelRoot optimize(RelRoot root, final List<Materialization> materializations, final List<CalciteSchema.LatticeEntry> lattices) {
    final RelOptPlanner planner = root.rel.getCluster().getPlanner();
    final DataContext dataContext = context.getDataContext();
    planner.setExecutor(new RexExecutorImpl(dataContext));
    final List<RelOptMaterialization> materializationList = new ArrayList<>(materializations.size());
    for (Materialization materialization : materializations) {
        List<String> qualifiedTableName = materialization.materializedTable.path();
        materializationList.add(new RelOptMaterialization(castNonNull(materialization.tableRel), castNonNull(materialization.queryRel), materialization.starRelOptTable, qualifiedTableName));
    }
    final List<RelOptLattice> latticeList = new ArrayList<>(lattices.size());
    for (CalciteSchema.LatticeEntry lattice : lattices) {
        final CalciteSchema.TableEntry starTable = lattice.getStarTable();
        final JavaTypeFactory typeFactory = context.getTypeFactory();
        final RelOptTableImpl starRelOptTable = RelOptTableImpl.create(catalogReader, starTable.getTable().getRowType(typeFactory), starTable, null);
        latticeList.add(new RelOptLattice(lattice.getLattice(), starRelOptTable));
    }
    final RelTraitSet desiredTraits = getDesiredRootTraitSet(root);
    final Program program = getProgram();
    final RelNode rootRel4 = program.run(planner, root.rel, desiredTraits, materializationList, latticeList);
    if (LOGGER.isDebugEnabled()) {
        LOGGER.debug("Plan after physical tweaks:\n{}", RelOptUtil.toString(rootRel4, SqlExplainLevel.ALL_ATTRIBUTES));
    }
    return root.withRel(rootRel4);
}

Example 33 with DataContext

use of org.apache.calcite.DataContext in project calcite by apache.

the class RexExecutorTest method check.

protected void check(final Action action) {
    Frameworks.withPrepare((cluster, relOptSchema, rootSchema, statement) -> {
        final RexBuilder rexBuilder = cluster.getRexBuilder();
        DataContext dataContext = DataContexts.of(statement.getConnection(), rootSchema);
        final RexExecutorImpl executor = new RexExecutorImpl(dataContext);
        action.check(rexBuilder, executor);
        return null;
    });
}

Example 34 with DataContext

use of org.apache.calcite.DataContext in project calcite by apache.

the class RexExecutorTest method testVariableExecution.

/**
 * Tests an executor that uses variables stored in a {@link DataContext}.
 * Can change the value of the variable and execute again.
 */
@Test
void testVariableExecution() {
    check((rexBuilder, executor) -> {
        Object[] values = new Object[1];
        final DataContext testContext = DataContexts.of(name -> name.equals("inputRecord") ? values : fail("unknown: " + name));
        final RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
        final RelDataType varchar = typeFactory.createSqlType(SqlTypeName.VARCHAR);
        final RelDataType integer = typeFactory.createSqlType(SqlTypeName.INTEGER);
        // Calcite is internally creating the input ref via a RexRangeRef
        // which eventually leads to a RexInputRef. So we are good.
        final RexInputRef input = rexBuilder.makeInputRef(varchar, 0);
        final RexNode lengthArg = rexBuilder.makeLiteral(3, integer, true);
        final RexNode substr = rexBuilder.makeCall(SqlStdOperatorTable.SUBSTRING, input, lengthArg);
        ImmutableList<RexNode> constExps = ImmutableList.of(substr);
        final RelDataType rowType = typeFactory.builder().add("someStr", varchar).build();
        final RexExecutable exec = executor.getExecutable(rexBuilder, constExps, rowType);
        exec.setDataContext(testContext);
        values[0] = "Hello World";
        Object[] result = exec.execute();
        assertTrue(result[0] instanceof String);
        assertThat((String) result[0], equalTo("llo World"));
        values[0] = "Calcite";
        result = exec.execute();
        assertTrue(result[0] instanceof String);
        assertThat((String) result[0], equalTo("lcite"));
    });
}

Example 35 with DataContext

use of org.apache.calcite.DataContext in project calcite by apache.

the class JaninoRexCompiler method baz.

/**
 * Given a method that implements {@link Scalar#execute(Context, Object[])},
 * adds a bridge method that implements {@link Scalar#execute(Context)}, and
 * compiles.
 */
static Scalar.Producer baz(ParameterExpression context_, ParameterExpression outputValues_, BlockStatement block, List<Statement> declList) {
    final List<MemberDeclaration> declarations = new ArrayList<>();
    final List<MemberDeclaration> innerDeclarations = new ArrayList<>();
    // public Scalar apply(DataContext root) {
    // <<staticList>>
    // return new Scalar() {
    // <<inner declarations>>
    // };
    // }
    final List<Statement> statements = new ArrayList<>(declList);
    statements.add(Expressions.return_(null, Expressions.new_(Scalar.class, ImmutableList.of(), innerDeclarations)));
    declarations.add(Expressions.methodDecl(Modifier.PUBLIC, Scalar.class, BuiltInMethod.FUNCTION_APPLY.method.getName(), ImmutableList.of(DataContext.ROOT), Expressions.block(statements)));
    // (bridge method)
    // public Object apply(Object root) {
    // return this.apply((DataContext) root);
    // }
    final ParameterExpression objectRoot = Expressions.parameter(Object.class, "root");
    declarations.add(Expressions.methodDecl(Modifier.PUBLIC, Object.class, BuiltInMethod.FUNCTION_APPLY.method.getName(), ImmutableList.of(objectRoot), Expressions.block(Expressions.return_(null, Expressions.call(Expressions.parameter(Scalar.Producer.class, "this"), BuiltInMethod.FUNCTION_APPLY.method, Expressions.convert_(objectRoot, DataContext.class))))));
    // public void execute(Context, Object[] outputValues)
    innerDeclarations.add(Expressions.methodDecl(Modifier.PUBLIC, void.class, BuiltInMethod.SCALAR_EXECUTE2.method.getName(), ImmutableList.of(context_, outputValues_), block));
    // public Object execute(Context)
    final BlockBuilder builder = new BlockBuilder();
    final Expression values_ = builder.append("values", Expressions.newArrayBounds(Object.class, 1, Expressions.constant(1)));
    builder.add(Expressions.statement(Expressions.call(Expressions.parameter(Scalar.class, "this"), BuiltInMethod.SCALAR_EXECUTE2.method, context_, values_)));
    builder.add(Expressions.return_(null, Expressions.arrayIndex(values_, Expressions.constant(0))));
    innerDeclarations.add(Expressions.methodDecl(Modifier.PUBLIC, Object.class, BuiltInMethod.SCALAR_EXECUTE1.method.getName(), ImmutableList.of(context_), builder.toBlock()));
    final ClassDeclaration classDeclaration = Expressions.classDecl(Modifier.PUBLIC, "Buzz", null, ImmutableList.of(Scalar.Producer.class), declarations);
    String s = Expressions.toString(declarations, "\n", false);
    if (CalciteSystemProperty.DEBUG.value()) {
        Util.debugCode(System.out, s);
    }
    try {
        return getScalar(classDeclaration, s);
    } catch (CompileException | IOException e) {
        throw new RuntimeException(e);
    }
}