@Test
  public void testCollectUnknownReference() throws Throwable {
    expectedException.expect(UnhandledServerException.class);
    expectedException.expectMessage("Unknown Reference some.table.some_column");

    Reference unknownReference =
        new Reference(
            new ReferenceInfo(
                new ReferenceIdent(new TableIdent("some", "table"), "some_column"),
                RowGranularity.NODE,
                DataTypes.BOOLEAN));
    CollectPhase collectNode =
        new CollectPhase(
            UUID.randomUUID(),
            0,
            "unknown",
            testRouting,
            Collections.<Symbol>singletonList(unknownReference),
            EMPTY_PROJECTIONS);
    collectNode.maxRowGranularity(RowGranularity.NODE);
    try {
      getBucket(collectNode);
    } catch (ExecutionException e) {
      throw e.getCause();
    }
  }
  @Test
  public void testWrongRouting() throws Exception {

    expectedException.expect(UnhandledServerException.class);
    expectedException.expectMessage("unsupported routing");

    CollectPhase collectNode =
        new CollectPhase(
            UUID.randomUUID(),
            0,
            "wrong",
            new Routing(
                TreeMapBuilder.<String, Map<String, List<Integer>>>newMapBuilder()
                    .put(
                        "bla",
                        TreeMapBuilder.<String, List<Integer>>newMapBuilder()
                            .put("my_index", Arrays.asList(1))
                            .put("my_index", Arrays.asList(1))
                            .map())
                    .map()),
            ImmutableList.<Symbol>of(),
            EMPTY_PROJECTIONS);
    collectNode.maxRowGranularity(RowGranularity.DOC);
    operation.collect(collectNode, new CollectingProjector(), null);
  }
  @Test
  public void testCollectShardExpressions() throws Exception {
    List<Symbol> toCollect = ImmutableList.<Symbol>of(testShardIdReference);
    CollectPhase collectNode =
        new CollectPhase(
            UUID.randomUUID(), 0, "shardCollect", shardRouting(0, 1), toCollect, EMPTY_PROJECTIONS);
    collectNode.maxRowGranularity(RowGranularity.SHARD);

    Bucket result = getBucket(collectNode);
    assertThat(result.size(), is(2));
    assertThat(result, containsInAnyOrder(isRow(0), isRow(1)));
  }
 private Bucket getBucket(CollectPhase collectNode)
     throws InterruptedException, ExecutionException {
   CollectingProjector cd = new CollectingProjector();
   JobExecutionContext.Builder builder = jobContextService.newBuilder(collectNode.jobId());
   JobCollectContext jobCollectContext =
       new JobCollectContext(
           collectNode.jobId(), collectNode, operation, RAM_ACCOUNTING_CONTEXT, cd);
   builder.addSubContext(collectNode.executionPhaseId(), jobCollectContext);
   JobExecutionContext context = jobContextService.createContext(builder);
   cd.startProjection(jobCollectContext);
   operation.collect(collectNode, cd, jobCollectContext);
   return cd.result().get();
 }
 @Test
 public void testCollectShardExpressionsLiteralsAndNodeExpressions() throws Exception {
   CollectPhase collectNode =
       new CollectPhase(
           UUID.randomUUID(),
           0,
           "shardCollect",
           shardRouting(0, 1),
           Arrays.asList(testShardIdReference, Literal.newLiteral(true), testNodeReference),
           EMPTY_PROJECTIONS);
   collectNode.maxRowGranularity(RowGranularity.SHARD);
   Bucket result = getBucket(collectNode);
   assertThat(result.size(), is(2));
   assertThat(result, containsInAnyOrder(isRow(0, true, (short) 1), isRow(1, true, (short) 1)));
 }
 @Test
 public void testCollectFunction() throws Exception {
   Function twoTimesTruthFunction = new Function(TestFunction.info, TO_COLLECT_TEST_REF);
   CollectPhase collectNode =
       new CollectPhase(
           UUID.randomUUID(),
           0,
           "unknown",
           testRouting,
           Arrays.asList(twoTimesTruthFunction, testNodeReference),
           EMPTY_PROJECTIONS);
   collectNode.maxRowGranularity(RowGranularity.NODE);
   Bucket result = getBucket(collectNode);
   assertThat(result.size(), equalTo(1));
   assertThat(result, contains(isRow(2, (short) 1)));
 }
  @Test
  public void testCollectExpressions() throws Exception {
    CollectPhase collectNode =
        new CollectPhase(
            UUID.randomUUID(),
            0,
            "collect",
            testRouting,
            Collections.<Symbol>singletonList(testNodeReference),
            EMPTY_PROJECTIONS);
    collectNode.maxRowGranularity(RowGranularity.NODE);

    Bucket result = getBucket(collectNode);

    assertThat(result.size(), equalTo(1));
    assertThat(result, contains(isRow((short) 1)));
  }
 @Test
 public void testCollectWithFalseWhereClause() throws Exception {
   CollectPhase collectNode =
       new CollectPhase(
           UUID.randomUUID(),
           0,
           "whereClause",
           testRouting,
           TO_COLLECT_TEST_REF,
           EMPTY_PROJECTIONS);
   collectNode.whereClause(
       new WhereClause(
           new Function(
               AndOperator.INFO,
               Arrays.<Symbol>asList(Literal.newLiteral(false), Literal.newLiteral(false)))));
   Bucket result = getBucket(collectNode);
   assertThat(result.size(), is(0));
 }
  private Bucket collect(CollectPhase collectNode) throws Throwable {
    ContextPreparer contextPreparer = internalCluster().getDataNodeInstance(ContextPreparer.class);
    JobContextService contextService =
        internalCluster().getDataNodeInstance(JobContextService.class);
    SharedShardContexts sharedShardContexts =
        new SharedShardContexts(internalCluster().getDataNodeInstance(IndicesService.class));
    JobExecutionContext.Builder builder = contextService.newBuilder(collectNode.jobId());
    NodeOperation nodeOperation =
        NodeOperation.withDownstream(
            collectNode, mock(ExecutionPhase.class), (byte) 0, "remoteNode");

    List<ListenableFuture<Bucket>> results =
        contextPreparer.prepareOnRemote(
            collectNode.jobId(), ImmutableList.of(nodeOperation), builder, sharedShardContexts);
    JobExecutionContext context = contextService.createContext(builder);
    context.start();
    return results.get(0).get(2, TimeUnit.SECONDS);
  }
 @Test
 public void testCollectWithTrueWhereClause() throws Exception {
   CollectPhase collectNode =
       new CollectPhase(
           UUID.randomUUID(),
           0,
           "whereClause",
           testRouting,
           TO_COLLECT_TEST_REF,
           EMPTY_PROJECTIONS);
   collectNode.whereClause(
       new WhereClause(
           new Function(
               AndOperator.INFO,
               Arrays.<Symbol>asList(Literal.newLiteral(true), Literal.newLiteral(true)))));
   collectNode.maxRowGranularity(RowGranularity.NODE);
   Bucket result = getBucket(collectNode);
   assertThat(result, contains(isRow((short) 1)));
 }
  @Test
  public void testCollectShardExpressionsWhereShardIdIs0() throws Exception {
    EqOperator op =
        (EqOperator)
            functions.get(
                new FunctionIdent(
                    EqOperator.NAME,
                    ImmutableList.<DataType>of(DataTypes.INTEGER, DataTypes.INTEGER)));

    List<Symbol> toCollect = ImmutableList.<Symbol>of(testShardIdReference);
    CollectPhase collectNode =
        new CollectPhase(
            UUID.randomUUID(), 0, "shardCollect", shardRouting(0, 1), toCollect, EMPTY_PROJECTIONS);
    collectNode.whereClause(
        new WhereClause(
            new Function(op.info(), Arrays.asList(testShardIdReference, Literal.newLiteral(0)))));
    collectNode.maxRowGranularity(RowGranularity.SHARD);
    Bucket result = getBucket(collectNode);
    assertThat(result, contains(isRow(0)));
  }
  private CollectPhase createCollectNode(
      Planner.Context plannerContext, boolean keepContextForFetcher) {
    TableInfo tableInfo = docSchemaInfo.getTableInfo("characters");

    ReferenceInfo docIdRefInfo = tableInfo.getReferenceInfo(new ColumnIdent("_docid"));
    Symbol docIdRef = new Reference(docIdRefInfo);
    List<Symbol> toCollect = ImmutableList.of(docIdRef);

    CollectPhase collectNode =
        new CollectPhase(
            UUID.randomUUID(),
            plannerContext.nextExecutionPhaseId(),
            "collect",
            tableInfo.getRouting(WhereClause.MATCH_ALL, null),
            toCollect,
            ImmutableList.<Projection>of());
    collectNode.maxRowGranularity(RowGranularity.DOC);
    collectNode.keepContextForFetcher(keepContextForFetcher);
    plannerContext.allocateJobSearchContextIds(collectNode.routing());

    return collectNode;
  }
 @Test
 public void testCollectWithNullWhereClause() throws Exception {
   EqOperator op =
       (EqOperator)
           functions.get(
               new FunctionIdent(
                   EqOperator.NAME,
                   ImmutableList.<DataType>of(DataTypes.INTEGER, DataTypes.INTEGER)));
   CollectPhase collectNode =
       new CollectPhase(
           UUID.randomUUID(),
           0,
           "whereClause",
           testRouting,
           TO_COLLECT_TEST_REF,
           EMPTY_PROJECTIONS);
   collectNode.whereClause(
       new WhereClause(
           new Function(op.info(), Arrays.<Symbol>asList(Literal.NULL, Literal.NULL))));
   Bucket result = getBucket(collectNode);
   assertThat(result.size(), is(0));
 }
  private List<Bucket> getBuckets(CollectPhase collectNode)
      throws InterruptedException, java.util.concurrent.ExecutionException {
    List<Bucket> results = new ArrayList<>();
    for (String nodeName : internalCluster().getNodeNames()) {
      ContextPreparer contextPreparer =
          internalCluster().getInstance(ContextPreparer.class, nodeName);
      JobContextService contextService =
          internalCluster().getInstance(JobContextService.class, nodeName);

      JobExecutionContext.Builder builder = contextService.newBuilder(collectNode.jobId());
      ListenableFuture<Bucket> future =
          contextPreparer.prepare(
              collectNode.jobId(),
              NodeOperation.withDownstream(collectNode, mock(ExecutionPhase.class)),
              builder);
      assert future != null;

      JobExecutionContext context = contextService.createContext(builder);
      context.start();
      results.add(future.get());
    }
    return results;
  }
  @Test
  public void testStreaming() throws Exception {
    ImmutableList<Symbol> toCollect = ImmutableList.<Symbol>of(new Value(DataTypes.STRING));
    UUID jobId = UUID.randomUUID();
    CollectPhase cn =
        new CollectPhase(
            jobId,
            0,
            "cn",
            new Routing(),
            RowGranularity.DOC,
            toCollect,
            ImmutableList.<Projection>of(),
            WhereClause.MATCH_ALL,
            DistributionInfo.DEFAULT_MODULO);

    BytesStreamOutput out = new BytesStreamOutput();
    cn.writeTo(out);

    BytesStreamInput in = new BytesStreamInput(out.bytes());
    CollectPhase cn2 = CollectPhase.FACTORY.create();
    cn2.readFrom(in);
    assertThat(cn, equalTo(cn2));

    assertThat(cn.toCollect(), is(cn2.toCollect()));
    assertThat(cn.executionNodes(), is(cn2.executionNodes()));
    assertThat(cn.jobId(), is(cn2.jobId()));
    assertThat(cn.executionPhaseId(), is(cn2.executionPhaseId()));
    assertThat(cn.maxRowGranularity(), is(cn2.maxRowGranularity()));
    assertThat(cn.distributionInfo(), is(cn2.distributionInfo()));
  }