org.apache.arrow.flight.FlightStream Java Examples

@Override
public Long call() {
  long count = 0;
  int readIndex = 0;
  logger.debug("starting work on flight endpoint with ticket {} to {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
  try (FlightClient c = flightClient(allocator, endpoint.getLocations().get(0))) {
    c.authenticate(new BasicClientAuthHandler(SystemUser.SYSTEM_USERNAME, null));
    logger.debug("trying to get stream for flight endpoint with ticket {} to {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
    FlightStream fs = c.getStream(endpoint.getTicket());
    logger.debug("got stream for flight endpoint with ticket {} to {}. Will now try and read", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
    while (fs.next()) {
      long thisCount = fs.getRoot().getRowCount();
      count += thisCount;
      logger.debug("got results from stream for flight endpoint with ticket {} to {}. This is read {} and we got {} rows back for a total of {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri(), ++readIndex, thisCount, count);
      fs.getRoot().clear();
    }
  } catch (InterruptedException e) {

  } catch (Throwable t) {
    logger.error("Error in stream fetch", t);
  }
  logger.debug("got all results from stream for flight endpoint with ticket {} to {}. We read {} batches and we got {} rows back", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri(), ++readIndex, count);
  return count;
}
 

@Override
public Long call() {
  long count = 0;
  int readIndex = 0;
  logger.debug("starting work on flight endpoint with ticket {} to {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
  try (FlightClient c = flightClient(allocator, endpoint.getLocations().get(0))) {
    c.authenticate(new BasicClientAuthHandler(SystemUser.SYSTEM_USERNAME, null));
    logger.debug("trying to get stream for flight endpoint with ticket {} to {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
    FlightStream fs = c.getStream(endpoint.getTicket());
    logger.debug("got stream for flight endpoint with ticket {} to {}. Will now try and read", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
    while (fs.next()) {
      long thisCount = fs.getRoot().getRowCount();
      count += thisCount;
      logger.debug("got results from stream for flight endpoint with ticket {} to {}. This is read {} and we got {} rows back for a total of {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri(), ++readIndex, thisCount, count);
      fs.getRoot().clear();
    }
  } catch (InterruptedException e) {

  } catch (Throwable t) {
    logger.error("Error in stream fetch", t);
  }
  logger.debug("got all results from stream for flight endpoint with ticket {} to {}. We read {} batches and we got {} rows back", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri(), ++readIndex, count);
  return count;
}
 

/**
 * Streams JobData from server over gRPC and creates list of RecordBatchHolder to populate JobDataFragment
 * @param stream flight stream for a particular job
 * @param allocator allocator for vectors
 * @param limit max number of results to fetch
 * @return
 */
public static List<RecordBatchHolder> getData(FlightStream stream, BufferAllocator allocator, int limit) {
  final List<RecordBatchHolder> batches = new ArrayList<>();
  try (final VectorContainer container = new VectorContainer(allocator);
       final VectorSchemaRoot root = stream.getRoot()) {

    int remaining = limit;

    while (stream.next()) {
      VectorContainer.transferFromRoot(root, container, allocator);

      final int currentBatchCount = root.getRowCount();
      // Determine batchEnd so that we are not returning more records
      // than the limit even when there are extra batches are in the stream.
      int batchEnd = min(currentBatchCount, remaining);

      final RecordBatchHolder batchHolder = newRecordBatchHolder(
        new RecordBatchData(container, allocator),
        0,
        batchEnd
      );
      batches.add(batchHolder);
      remaining -= batchHolder.size();

      // break if we hit the limit
      if (remaining == 0) {
        break;
      }
    }

    // If no batches/records, return an empty result
    if (batches.isEmpty()) {
      container.setRecordCount(0);
      container.buildSchema();
      batches.add(newRecordBatchHolder(new RecordBatchData(container, allocator), 0, 0));
    }
  }
  return batches;
}
 

/**
 * Factory for returning JobData over gRPC
 * @param jobsService reference to job service
 * @param bufferAllocator allocator for vectors
 * @param jobId jobid for corresponding jobresults
 * @param offset start index of results
 * @param limit max number of results to fetch
 * @return JobDataFragment
 */
public static JobDataFragment getJobData(JobsService jobsService, BufferAllocator bufferAllocator, JobId jobId, int offset, int limit) {
  final FlightClient flightClient = jobsService.getJobsClient().getFlightClient();
  final Ticket ticket = new JobsFlightTicket(jobId.getId(), offset, limit).toTicket();
  try (FlightStream flightStream = flightClient.getStream(ticket)) {
    return new JobDataFragmentImpl(new RecordBatches(JobDataClientUtils.getData(
      flightStream, bufferAllocator, limit)), offset, jobId);
  } catch (FlightRuntimeException fre) {
    Optional<UserException> ue = JobsRpcUtils.fromFlightRuntimeException(fre);
    throw ue.isPresent() ? ue.get() : fre;
  } catch (Exception e) {
    Throwables.throwIfUnchecked(e);
    throw new RuntimeException(e);
  }
}
 

@Test
public void testException() {
  try {
    FlightStream results = client.getStream(new Ticket("test".getBytes()));
    results.getDescriptor();
    Assert.fail();
  } catch (FlightRuntimeException e) {
    Optional<UserException> ue = JobsRpcUtils.fromFlightRuntimeException(e);
    Assert.assertTrue(ue.isPresent());
    Assert.assertEquals(expectedMessage, ue.get().getOriginalMessage());
  }
}
 

public static JobDataFragmentWrapper getJobData(JobsService jobsService, BufferAllocator allocator, JobId jobId, int offset, int limit) {
  try (final FlightStream stream = jobsService.getJobsClient().getFlightClient()
    .getStream(new JobsFlightTicket(jobId.getId(), offset, limit).toTicket())) {
    List<RecordBatchHolder> batches = JobDataClientUtils.getData(stream, allocator, limit);
    return new JobDataFragmentWrapper(offset, ReleasingData.from(new RecordBatches(batches), jobId));
  } catch (FlightRuntimeException fre) {
    Optional<UserException> ue = JobsRpcUtils.fromFlightRuntimeException(fre);
    throw ue.isPresent() ? ue.get() : fre;
  } catch (Exception e) {
    Throwables.throwIfUnchecked(e);
    throw new RuntimeException(e);
  }
}
 

@Override
public Runnable acceptPut(CallContext callContext, FlightStream flightStream, StreamListener<org.apache.arrow.flight.PutResult> streamListener) {
  throw Status.UNIMPLEMENTED.asRuntimeException();
}
 

@Override
public Runnable acceptPut(CallContext callContext, FlightStream flightStream, StreamListener<PutResult> streamListener) {
  throw Status.UNIMPLEMENTED.asRuntimeException();
}
 

@Override
public Runnable acceptPut(CallContext callContext, FlightStream flightStream, StreamListener<PutResult> streamListener) {
  throw Status.UNIMPLEMENTED.asRuntimeException();
}