Java Code Examples for org.apache.calcite.linq4j.tree.BlockBuilder#append()

@Test void testAssignInConditionMultipleUsageNonOptimized() {
  // int t = 2;
  // return (t = 1) != a ? 1 : c
  final BlockBuilder builder = new BlockBuilder(true);
  final ParameterExpression t = Expressions.parameter(int.class, "t");

  builder.add(Expressions.declare(0, t, TWO));

  Expression v = builder.append("v",
      Expressions.makeTernary(ExpressionType.Conditional,
          Expressions.makeBinary(ExpressionType.NotEqual,
              Expressions.assign(t, Expressions.constant(1)),
              Expressions.parameter(int.class, "a")),
          t,
          Expressions.parameter(int.class, "c")));
  builder.add(Expressions.return_(null, v));
  assertEquals(
      "{\n"
          + "  int t = 2;\n"
          + "  return (t = 1) != a ? t : c;\n"
          + "}\n",
      Expressions.toString(builder.toBlock()));
}
 

void checkAssignInConditionOptimizedOut(int modifiers, String s) {
  // int t;
  // return (t = 1) != a ? b : c
  final BlockBuilder builder = new BlockBuilder(true);
  final ParameterExpression t =
      Expressions.parameter(int.class, "t");

  builder.add(Expressions.declare(modifiers, t, null));

  Expression v = builder.append("v",
      Expressions.makeTernary(ExpressionType.Conditional,
          Expressions.makeBinary(ExpressionType.NotEqual,
              Expressions.assign(t, Expressions.constant(1)),
              Expressions.parameter(int.class, "a")),
          Expressions.parameter(int.class, "b"),
          Expressions.parameter(int.class, "c")));
  builder.add(Expressions.return_(null, v));
  assertThat(Expressions.toString(builder.toBlock()),
      CoreMatchers.equalTo(s));
}
 

@Test void testAssignInConditionMultipleUsage() {
  // int t;
  // return (t = 1) != a ? t : c
  final BlockBuilder builder = new BlockBuilder(true);
  final ParameterExpression t = Expressions.parameter(int.class, "t");

  builder.add(Expressions.declare(0, t, null));

  Expression v = builder.append("v",
      Expressions.makeTernary(ExpressionType.Conditional,
          Expressions.makeBinary(ExpressionType.NotEqual,
              Expressions.assign(t, Expressions.constant(1)),
              Expressions.parameter(int.class, "a")),
          t,
          Expressions.parameter(int.class, "c")));
  builder.add(Expressions.return_(null, v));
  assertEquals(
      "{\n"
          + "  int t;\n"
          + "  return (t = 1) != a ? t : c;\n"
          + "}\n",
      Expressions.toString(builder.toBlock()));
}
 

public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
  // Generate:
  //   RDD rdd = ...;
  //   return SparkRuntime.asEnumerable(rdd);
  final BlockBuilder list = new BlockBuilder();
  final SparkRel child = (SparkRel) getInput();
  final PhysType physType =
      PhysTypeImpl.of(implementor.getTypeFactory(),
          getRowType(),
          JavaRowFormat.CUSTOM);
  SparkRel.Implementor sparkImplementor =
      new SparkImplementorImpl(implementor);
  final SparkRel.Result result = child.implementSpark(sparkImplementor);
  final Expression rdd = list.append("rdd", result.block);
  final Expression enumerable =
      list.append(
          "enumerable",
          Expressions.call(
              SparkMethod.AS_ENUMERABLE.method,
              rdd));
  list.add(
      Expressions.return_(null, enumerable));
  return implementor.result(physType, list.toBlock());
}
 

public Result implementSpark(Implementor implementor) {
  // Generate:
  //   Enumerable source = ...;
  //   return SparkRuntime.createRdd(sparkContext, source);
  final BlockBuilder list = new BlockBuilder();
  final EnumerableRel child = (EnumerableRel) getInput();
  final PhysType physType =
      PhysTypeImpl.of(
          implementor.getTypeFactory(), getRowType(),
          JavaRowFormat.CUSTOM);
  final Expression source = null; // TODO:
  final Expression sparkContext =
      Expressions.call(
          SparkMethod.GET_SPARK_CONTEXT.method,
          implementor.getRootExpression());
  final Expression rdd =
      list.append(
          "rdd",
          Expressions.call(
              SparkMethod.CREATE_RDD.method,
              sparkContext,
              source));
  list.add(
      Expressions.return_(null, rdd));
  return implementor.result(physType, list.toBlock());
}
 

@Test void testBlockBuilder2() {
  BlockBuilder statements = new BlockBuilder();
  Expression element =
      statements.append(
          "element", Expressions.constant(null));
  Expression comparator =
      statements.append(
          "comparator", Expressions.constant(null, Comparator.class));
  Expression treeSet =
      statements.append(
          "treeSet",
          Expressions.new_(
              TreeSet.class,
              Arrays.asList(comparator)));
  statements.add(
      Expressions.return_(
          null,
          Expressions.call(
              treeSet,
              "add",
              element)));
  BlockStatement expression = statements.toBlock();
  final String expected = "{\n"
      + "  final java.util.TreeSet treeSet = new java.util.TreeSet(\n"
      + "    (java.util.Comparator) null);\n"
      + "  return treeSet.add(null);\n"
      + "}\n";
  assertThat(Expressions.toString(expression), is(expected));
  expression.accept(new Shuttle());
}
 

private boolean implementResult(List<AggImpState> aggs,
    final BlockBuilder builder,
    final Function<BlockBuilder, WinAggFrameResultContext> frame,
    final Function<AggImpState, List<RexNode>> rexArguments,
    boolean cachedBlock) {
  boolean nonEmpty = false;
  for (final AggImpState agg : aggs) {
    boolean needCache = true;
    if (agg.implementor instanceof WinAggImplementor) {
      WinAggImplementor imp = (WinAggImplementor) agg.implementor;
      needCache = imp.needCacheWhenFrameIntact();
    }
    if (needCache ^ cachedBlock) {
      // Regular aggregates do not change when the windowing frame keeps
      // the same. Ths
      continue;
    }
    nonEmpty = true;
    Expression res = agg.implementor.implementResult(agg.context,
        new WinAggResultContextImpl(builder, agg.state, frame) {
          public List<RexNode> rexArguments() {
            return rexArguments.apply(agg);
          }
        });
    // Several count(a) and count(b) might share the result
    Expression aggRes = builder.append("a" + agg.aggIdx + "res",
        EnumUtils.convert(res, agg.result.getType()));
    builder.add(
        Expressions.statement(Expressions.assign(agg.result, aggRes)));
  }
  return nonEmpty;
}
 

public Expression field(BlockBuilder list, int index, Type storageType) {
  if (index < actualInputFieldCount) {
    Expression current = list.append("current", row);
    return rowPhysType.fieldReference(current, index, storageType);
  }
  return constants.get(index - actualInputFieldCount);
}
 

private BlockBuilder appendBlockWithSameVariable(
    Expression initializer1, Expression initializer2) {
  BlockBuilder outer = new BlockBuilder();
  ParameterExpression outerX = Expressions.parameter(int.class, "x");
  outer.add(Expressions.declare(0, outerX, initializer1));
  outer.add(Expressions.statement(Expressions.assign(outerX, Expressions.constant(1))));

  BlockBuilder inner = new BlockBuilder();
  ParameterExpression innerX = Expressions.parameter(int.class, "x");
  inner.add(Expressions.declare(0, innerX, initializer2));
  inner.add(Expressions.statement(Expressions.assign(innerX, Expressions.constant(42))));
  inner.add(Expressions.return_(null, innerX));
  outer.append("x", inner.toBlock());
  return outer;
}
 

@Test void testReuseExpressionsFromUpperLevel() {
  Expression x = b.append("x", Expressions.add(ONE, TWO));
  BlockBuilder nested = new BlockBuilder(true, b);
  Expression y = nested.append("y", Expressions.add(ONE, TWO));
  nested.add(Expressions.return_(null, Expressions.add(y, y)));
  b.add(nested.toBlock());
  assertEquals(
      "{\n"
          + "  final int x = 1 + 2;\n"
          + "  {\n"
          + "    return x + x;\n"
          + "  }\n"
          + "}\n",
      b.toBlock().toString());
}
 

public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
  final BlockBuilder builder = new BlockBuilder();
  final EnumerableRel child = (EnumerableRel) getInput();
  final Result result = implementor.visitChild(this, 0, child, pref);
  final PhysType physType =
      PhysTypeImpl.of(
          implementor.getTypeFactory(),
          getRowType(),
          result.format);
  Expression childExp =
      builder.append("child", result.block);

  PhysType inputPhysType = result.physType;
  final Pair<Expression, Expression> pair =
      inputPhysType.generateCollationKey(
          collation.getFieldCollations());

  builder.add(
      Expressions.return_(null,
          Expressions.call(childExp,
              BuiltInMethod.ORDER_BY.method,
              Expressions.list(
                  builder.append("keySelector", pair.left))
                  .appendIfNotNull(
                      builder.appendIfNotNull("comparator", pair.right)))));
  return implementor.result(physType, builder.toBlock());
}
 

@Override
public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
    final BlockBuilder builder = new BlockBuilder();
    Expression unionExp = null;
    for (Ord<RelNode> ord : Ord.zip(inputs)) {
        EnumerableRel input = (EnumerableRel) ord.e;
        final Result result = implementor.visitChild(this, ord.i, input, pref);
        Expression childExp =
                builder.append(
                        "child" + ord.i,
                        result.block);

        if (unionExp == null) {
            unionExp = childExp;
        } else {
            unionExp = createUnionExpression(unionExp, childExp, result.format == JavaRowFormat.ARRAY);
        }
    }

    builder.add(unionExp);
    final PhysType physType =
            PhysTypeImpl.of(
                    implementor.getTypeFactory(),
                    getRowType(),
                    pref.prefer(JavaRowFormat.CUSTOM));
    return implementor.result(physType, builder.toBlock());
}
 

public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
  final BlockBuilder builder = new BlockBuilder();
  final Result leftResult =
      implementor.visitChild(this, 0, (EnumerableRel) left, pref);
  Expression leftExpression =
      builder.append("left", leftResult.block);
  final Result rightResult =
      implementor.visitChild(this, 1, (EnumerableRel) right, pref);
  Expression rightExpression =
      builder.append("right", rightResult.block);
  final PhysType physType =
      PhysTypeImpl.of(implementor.getTypeFactory(),
          getRowType(),
          pref.preferArray());
  final Expression predicate =
      EnumUtils.generatePredicate(implementor, getCluster().getRexBuilder(), left, right,
          leftResult.physType, rightResult.physType, condition);
  return implementor.result(
      physType,
      builder.append(
          Expressions.call(BuiltInMethod.NESTED_LOOP_JOIN.method,
              leftExpression,
              rightExpression,
              predicate,
              EnumUtils.joinSelector(joinType,
                  physType,
                  ImmutableList.of(leftResult.physType,
                      rightResult.physType)),
              Expressions.constant(EnumUtils.toLinq4jJoinType(joinType))))
          .toBlock());
}
 

public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
  final BlockBuilder builder = new BlockBuilder();
  final EnumerableRel child = (EnumerableRel) getInput();
  final Result result = implementor.visitChild(this, 0, child, pref);
  final PhysType physType =
      PhysTypeImpl.of(
          implementor.getTypeFactory(),
          getRowType(),
          result.format);

  Expression v = builder.append("child", result.block);
  if (offset != null) {
    v = builder.append(
        "offset",
        Expressions.call(
            v,
            BuiltInMethod.SKIP.method,
            getExpression(offset)));
  }
  if (fetch != null) {
    v = builder.append(
        "fetch",
        Expressions.call(
            v,
            BuiltInMethod.TAKE.method,
            getExpression(fetch)));
  }

  builder.add(
      Expressions.return_(
          null,
          v));
  return implementor.result(physType, builder.toBlock());
}
 

@Override
public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
    final BlockBuilder builder = new BlockBuilder();
    Expression unionExp = null;
    for (Ord<RelNode> ord : Ord.zip(inputs)) {
        EnumerableRel input = (EnumerableRel) ord.e;
        final Result result = implementor.visitChild(this, ord.i, input, pref);
        Expression childExp =
                builder.append(
                        "child" + ord.i,
                        result.block);

        if (unionExp == null) {
            unionExp = childExp;
        } else {
            unionExp = createUnionExpression(unionExp, childExp, result.format == JavaRowFormat.ARRAY);
        }
    }

    builder.add(unionExp);
    final PhysType physType =
            PhysTypeImpl.of(
                    implementor.getTypeFactory(),
                    getRowType(),
                    pref.prefer(JavaRowFormat.CUSTOM));
    return implementor.result(physType, builder.toBlock());
}
 

public Expression field(BlockBuilder list, int index, Type storageType) {
  int offset = 0;
  for (Pair<Expression, PhysType> input : inputs) {
    final PhysType physType = input.right;
    int fieldCount = physType.getRowType().getFieldCount();
    if (index >= offset + fieldCount) {
      offset += fieldCount;
      continue;
    }
    final Expression left = list.append("current", input.left);
    return physType.fieldReference(left, index - offset, storageType);
  }
  throw new IllegalArgumentException("Unable to find field #" + index);
}
 

@Test void testBlockBuilder3() {
/*
    int a = 1;
    int b = a + 2;
    int c = a + 3;
    int d = a + 4;
    int e = {
      int b = a + 3;
      foo(b);
    }
    bar(a, b, c, d, e);
*/
    BlockBuilder builder0 = new BlockBuilder();
    final Expression a = builder0.append("_a", Expressions.constant(1));
    final Expression b =
        builder0.append("_b", Expressions.add(a, Expressions.constant(2)));
    final Expression c =
        builder0.append("_c", Expressions.add(a, Expressions.constant(3)));
    final Expression d =
        builder0.append("_d", Expressions.add(a, Expressions.constant(4)));

    BlockBuilder builder1 = new BlockBuilder();
    final Expression b1 =
        builder1.append("_b", Expressions.add(a, Expressions.constant(3)));
    builder1.add(
        Expressions.statement(
            Expressions.call(ExpressionTest.class, "foo", b1)));
    final Expression e = builder0.append("e", builder1.toBlock());
    builder0.add(
        Expressions.statement(
            Expressions.call(ExpressionTest.class, "bar", a, b, c, d, e)));
    // With the bug in BlockBuilder.append(String, BlockExpression),
    //    bar(1, _b, _c, _d, foo(_d));
    // Correct result is
    //    bar(1, _b, _c, _d, foo(_c));
    // because _c has the same expression (a + 3) as inner b.
    BlockStatement expression = builder0.toBlock();
    assertEquals(
        "{\n"
            + "  final int _b = 1 + 2;\n"
            + "  final int _c = 1 + 3;\n"
            + "  final int _d = 1 + 4;\n"
            + "  final int _b0 = 1 + 3;\n"
            + "  org.apache.calcite.linq4j.test.ExpressionTest.bar(1, _b, _c, _d, org.apache.calcite.linq4j.test.ExpressionTest.foo(_b0));\n"
            + "}\n",
        Expressions.toString(expression));
    expression.accept(new Shuttle());
  }
 

@Override public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
  // TODO for the moment only LAZY read & write is supported
  if (readType != Type.LAZY || writeType != Type.LAZY) {
    throw new UnsupportedOperationException(
        "EnumerableTableSpool supports for the moment only LAZY read and LAZY write");
  }

  //  ModifiableTable t = (ModifiableTable) root.getRootSchema().getTable(tableName);
  //  return lazyCollectionSpool(t.getModifiableCollection(), <inputExp>);

  BlockBuilder builder = new BlockBuilder();

  RelNode input = getInput();
  Result inputResult = implementor.visitChild(this, 0, (EnumerableRel) input, pref);

  String tableName = table.getQualifiedName().get(table.getQualifiedName().size() - 1);
  Expression tableExp = Expressions.convert_(
      Expressions.call(
          Expressions.call(
              implementor.getRootExpression(),
              BuiltInMethod.DATA_CONTEXT_GET_ROOT_SCHEMA.method),
          BuiltInMethod.SCHEMA_GET_TABLE.method,
          Expressions.constant(tableName, String.class)),
      ModifiableTable.class);
  Expression collectionExp = Expressions.call(
      tableExp,
      BuiltInMethod.MODIFIABLE_TABLE_GET_MODIFIABLE_COLLECTION.method);

  Expression inputExp = builder.append("input", inputResult.block);

  Expression spoolExp = Expressions.call(
      BuiltInMethod.LAZY_COLLECTION_SPOOL.method,
      collectionExp,
      inputExp);
  builder.add(spoolExp);

  PhysType physType = PhysTypeImpl.of(
      implementor.getTypeFactory(),
      getRowType(),
      pref.prefer(inputResult.format));
  return implementor.result(physType, builder.toBlock());
}
 

private Result implementHashJoin(EnumerableRelImplementor implementor, Prefer pref) {
  BlockBuilder builder = new BlockBuilder();
  final Result leftResult =
      implementor.visitChild(this, 0, (EnumerableRel) left, pref);
  Expression leftExpression =
      builder.append(
          "left", leftResult.block);
  final Result rightResult =
      implementor.visitChild(this, 1, (EnumerableRel) right, pref);
  Expression rightExpression =
      builder.append(
          "right", rightResult.block);
  final PhysType physType =
      PhysTypeImpl.of(
          implementor.getTypeFactory(), getRowType(), pref.preferArray());
  final PhysType keyPhysType =
      leftResult.physType.project(
          joinInfo.leftKeys, JavaRowFormat.LIST);
  Expression predicate = Expressions.constant(null);
  if (!joinInfo.nonEquiConditions.isEmpty()) {
    RexNode nonEquiCondition = RexUtil.composeConjunction(
        getCluster().getRexBuilder(), joinInfo.nonEquiConditions, true);
    if (nonEquiCondition != null) {
      predicate = EnumUtils.generatePredicate(implementor, getCluster().getRexBuilder(),
          left, right, leftResult.physType, rightResult.physType, nonEquiCondition);
    }
  }
  return implementor.result(
      physType,
      builder.append(
          Expressions.call(
              leftExpression,
              BuiltInMethod.HASH_JOIN.method,
              Expressions.list(
                  rightExpression,
                  leftResult.physType.generateAccessor(joinInfo.leftKeys),
                  rightResult.physType.generateAccessor(joinInfo.rightKeys),
                  EnumUtils.joinSelector(joinType,
                      physType,
                      ImmutableList.of(
                          leftResult.physType, rightResult.physType)))
                  .append(
                      Util.first(keyPhysType.comparer(),
                          Expressions.constant(null)))
                  .append(
                      Expressions.constant(joinType.generatesNullsOnLeft()))
                  .append(
                      Expressions.constant(
                          joinType.generatesNullsOnRight()))
                  .append(predicate)))
          .toBlock());
}
 

/** Given a method that implements {@link Scalar#execute(Context, Object[])},
 * adds a bridge method that implements {@link Scalar#execute(Context)}, and
 * compiles. */
static Scalar baz(ParameterExpression context_,
    ParameterExpression outputValues_, BlockStatement block) {
  final List<MemberDeclaration> declarations = new ArrayList<>();

  // public void execute(Context, Object[] outputValues)
  declarations.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))));
  declarations.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.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);
  }
}