Пример #1
0
  @Test
  public void testCollectWithPartitionedColumns() throws Throwable {
    Routing routing =
        docSchemaInfo.getTableInfo(PARTITIONED_TABLE_NAME).getRouting(WhereClause.MATCH_ALL, null);
    TableIdent tableIdent = new TableIdent(Schemas.DEFAULT_SCHEMA_NAME, PARTITIONED_TABLE_NAME);
    CollectPhase collectNode =
        getCollectNode(
            Arrays.<Symbol>asList(
                new Reference(
                    new ReferenceInfo(
                        new ReferenceIdent(tableIdent, "id"),
                        RowGranularity.DOC,
                        DataTypes.INTEGER)),
                new Reference(
                    new ReferenceInfo(
                        new ReferenceIdent(tableIdent, "date"),
                        RowGranularity.SHARD,
                        DataTypes.TIMESTAMP))),
            routing,
            WhereClause.MATCH_ALL);

    Bucket result = collect(collectNode);
    for (Row row : result) {
      System.out.println("Row:" + Arrays.toString(row.materialize()));
    }

    assertThat(result, containsInAnyOrder(isRow(1, 0L), isRow(2, 1L)));
  }
Пример #2
0
 protected boolean synchronizedSetNextRow(Row row) {
   if (paused) {
     pauseFifo.add(row.materialize());
     return true;
   } else {
     Object[] bufferedCells;
     while ((bufferedCells = pauseFifo.poll()) != null) {
       sharedRow.cells(bufferedCells);
       boolean wantMore = delegate.setNextRow(sharedRow);
       if (!wantMore) {
         return false;
       }
       if (paused) {
         pauseFifo.add(row.materialize());
         return true;
       }
     }
     return delegate.setNextRow(row);
   }
 }
  @Test
  public void testFetchAction() throws Exception {
    setUpCharacters();

    Analysis analysis = analyze("select id, name from characters");
    QueryThenFetchConsumer queryThenFetchConsumer =
        internalCluster().getInstance(QueryThenFetchConsumer.class);
    Planner.Context plannerContext = new Planner.Context(clusterService(), UUID.randomUUID());
    ConsumerContext consumerContext = new ConsumerContext(analysis.rootRelation(), plannerContext);
    QueryThenFetch plan =
        (QueryThenFetch)
            queryThenFetchConsumer.consume(analysis.rootRelation(), consumerContext).plan();

    List<Bucket> results = getBuckets(plan.collectNode());

    TransportFetchNodeAction transportFetchNodeAction =
        internalCluster().getInstance(TransportFetchNodeAction.class);

    // extract docIds by nodeId and jobSearchContextId
    Map<String, LongArrayList> jobSearchContextDocIds = new HashMap<>();
    for (Bucket rows : results) {
      long docId = (long) rows.iterator().next().get(0);
      // unpack jobSearchContextId and reader doc id from docId
      int jobSearchContextId = (int) (docId >> 32);
      String nodeId = plannerContext.nodeId(jobSearchContextId);
      LongArrayList docIdsPerNode = jobSearchContextDocIds.get(nodeId);
      if (docIdsPerNode == null) {
        docIdsPerNode = new LongArrayList();
        jobSearchContextDocIds.put(nodeId, docIdsPerNode);
      }
      docIdsPerNode.add(docId);
    }

    Iterable<Projection> projections =
        Iterables.filter(
            plan.mergeNode().projections(), Predicates.instanceOf(FetchProjection.class));
    FetchProjection fetchProjection = (FetchProjection) Iterables.getOnlyElement(projections);
    RowInputSymbolVisitor rowInputSymbolVisitor =
        new RowInputSymbolVisitor(internalCluster().getInstance(Functions.class));
    RowInputSymbolVisitor.Context context =
        rowInputSymbolVisitor.extractImplementations(fetchProjection.outputSymbols());

    final CountDownLatch latch = new CountDownLatch(jobSearchContextDocIds.size());
    final List<Row> rows = new ArrayList<>();
    for (Map.Entry<String, LongArrayList> nodeEntry : jobSearchContextDocIds.entrySet()) {
      NodeFetchRequest nodeFetchRequest = new NodeFetchRequest();
      nodeFetchRequest.jobId(plan.collectNode().jobId());
      nodeFetchRequest.executionPhaseId(plan.collectNode().executionPhaseId());
      nodeFetchRequest.toFetchReferences(context.references());
      nodeFetchRequest.closeContext(true);
      nodeFetchRequest.jobSearchContextDocIds(nodeEntry.getValue());

      transportFetchNodeAction.execute(
          nodeEntry.getKey(),
          nodeFetchRequest,
          new ActionListener<NodeFetchResponse>() {
            @Override
            public void onResponse(NodeFetchResponse nodeFetchResponse) {
              for (Row row : nodeFetchResponse.rows()) {
                rows.add(row);
              }
              latch.countDown();
            }

            @Override
            public void onFailure(Throwable e) {
              latch.countDown();
              fail(e.getMessage());
            }
          });
    }
    latch.await();

    assertThat(rows.size(), is(2));
    for (Row row : rows) {
      assertThat((Integer) row.get(0), anyOf(is(1), is(2)));
      assertThat(
          (BytesRef) row.get(1), anyOf(is(new BytesRef("Arthur")), is(new BytesRef("Ford"))));
    }
  }