org.apache.arrow.flight.Ticket Java Examples

public FlightInfo getInfo(FlightDescriptor descriptor, UserBitShared.ExternalId externalId) {
  try {
    logger.debug("Waiting for prepared statement handle to return for job id {}", ExternalIdHelper.toQueryId(externalId));
    CreatePreparedStatementResp handle = future.get();
    logger.debug("prepared statement handle for job id {} has returned", ExternalIdHelper.toQueryId(externalId));
    if (handle.getStatus() == RequestStatus.FAILED) {
      logger.warn("prepared statement handle for job id " + ExternalIdHelper.toQueryId(externalId) + " has failed", UserRemoteException.create(handle.getError()));
      throw Status.INTERNAL.withDescription(handle.getError().getMessage()).withCause(UserRemoteException.create(handle.getError())).asRuntimeException();
    }
    logger.debug("prepared statement handle for job id {} has succeeded", ExternalIdHelper.toQueryId(externalId));
    PreparedStatement statement = handle.getPreparedStatement();
    Ticket ticket = new Ticket(statement.getServerHandle().toByteArray());
    FlightEndpoint endpoint = new FlightEndpoint(ticket, location);
    logger.debug("flight endpoint for job id {} has been created with ticket {}", ExternalIdHelper.toQueryId(externalId), new String(ticket.getBytes()));
    Schema schema = fromMetadata(statement.getColumnsList());
    FlightInfo info = new FlightInfo(schema, descriptor, Lists.newArrayList(endpoint), -1L, -1L);
    logger.debug("flight info for job id {} has been created with schema {}", ExternalIdHelper.toQueryId(externalId), schema.toJson());
    return info;
  } catch (Exception e) {
    logger.warn("prepared statement handle for job id " + ExternalIdHelper.toQueryId(externalId) + " has failed", UserException.parseError(e).buildSilently());
    throw Status.UNKNOWN.withCause(UserException.parseError(e).buildSilently()).asRuntimeException();
  }
}
 

@Override
public void getStream(CallContext callContext, Ticket ticket, ServerStreamListener listener) {
  RetrieveData d = new RetrieveData(listener);
  RunQuery query;
  try {
    query = RunQuery.newBuilder()
      .setType(QueryType.PREPARED_STATEMENT)
      .setSource(UserProtos.SubmissionSource.UNKNOWN_SOURCE)
      .setPriority(UserProtos.QueryPriority
        .newBuilder()
        .setWorkloadClass(UserBitShared.WorkloadClass.GENERAL)
        .setWorkloadType(UserBitShared.WorkloadType.UNKNOWN)
        .build())
      .setPreparedStatementHandle(PreparedStatementHandle.parseFrom(ticket.getBytes()))
      .build();
  } catch (InvalidProtocolBufferException e) {
    throw Status.UNKNOWN.withCause(e).asRuntimeException();
  }

  UserRequest request = new UserRequest(RpcType.RUN_QUERY, query);
  submitWork(callContext, request, d);
}
 

@Override
public void doAction(CallContext callContext, Action action, StreamListener<Result> resultStreamListener) {
  try {
    logger.debug("do action is called with action type {} and body {}", action.getType(), new String(action.getBody()));
    if ("create".equals(action.getType())) {
      logger.debug("create stream called, creasting ticket with {}", new String(action.getBody()));
      putStream(getDescriptor(new Ticket(action.getBody())), null);
      logger.debug("put stream succeeded for {}, returning result 0", new String(action.getBody()));
      resultStreamListener.onNext(new Result(new byte[]{0}));
    } else if ("delete".equals(action.getType())) {
      logger.debug("delete stream called, deleting ticket with {}", new String(action.getBody()));
      deleteStream(new Ticket(action.getBody()));
      logger.debug("delete stream succeeded for {}, returning result 0", new String(action.getBody()));
      resultStreamListener.onNext(new Result(new byte[]{0}));
    }
    resultStreamListener.onCompleted();
  } catch (Throwable e) {
    resultStreamListener.onError(e);
  }
}
 

private FlightInfo getCoalesce(CallContext callContext, FlightDescriptor descriptor, Ticket cmd) {
  PrepareParallel d = new PrepareParallel();
  logger.debug("coalescing query {}", new String(cmd.getBytes()));
  RunQuery query;
  try {
    PreparedStatementHandle handle = PreparedStatementHandle.parseFrom(cmd.getBytes());

    query = RunQuery.newBuilder()
      .setType(QueryType.PREPARED_STATEMENT)
      .setSource(UserProtos.SubmissionSource.UNKNOWN_SOURCE)
      .setPreparedStatementHandle(handle)
      .setPriority(UserProtos.QueryPriority
        .newBuilder()
        .setWorkloadClass(UserBitShared.WorkloadClass.GENERAL)
        .setWorkloadType(UserBitShared.WorkloadType.UNKNOWN)
        .build())
      .build();
  } catch (InvalidProtocolBufferException e) {
    throw Status.UNKNOWN.withCause(e).asRuntimeException();
  }

  UserRequest request = new UserRequest(RpcType.RUN_QUERY, query);
  UserBitShared.ExternalId externalId = submitWork(callContext, request, d);
  String queryId = QueryIdHelper.getQueryId(ExternalIdHelper.toQueryId(externalId));
  logger.debug("submitted query for {} and got query id {}. Will now wait for parallel planner to return.", new String(cmd.getBytes()), queryId);
  return d.getInfo(descriptor, queryId);
}
 

@Override
    public void getStream(CallContext callContext, Ticket ticket, ServerStreamListener serverStreamListener) {
      String ticketId = new String(ticket.getBytes());
      logger.debug("formation getStream for ticket {}", ticketId);
      FlightDescriptor descriptor = getDescriptor(ticket);
      boolean isIn = holders.containsKey(descriptor);
      Stream holder = holders.computeIfAbsent(descriptor, (t) -> new Stream(allocator, descriptor));
      logger.debug("generated holder for {}, was it created: {}. Now submitting job", ticketId, isIn);
//      Future<?> f = executor.submit(new DrainRunnable(holder.getConsumer(), ticketId, serverStreamListener));
      new DrainRunnable(holder.getConsumer(), ticketId, serverStreamListener).run();
      logger.debug("job for {} was submitted", ticketId);
//      futures.put(ticket, f);
    }
 

private void deleteStream(Ticket ticket) {
      logger.debug("deleting stream {}", new String(ticket.getBytes()));
//      Future<?> f = futures.remove(ticket);
      Stream h = holders.get(getDescriptor(ticket));
      if (h != null) {
        logger.debug("sending cancel to stream {}", new String(ticket.getBytes()));
        if (h.cancel()) {
          futures.remove(ticket).cancel(true);
        }
      } else {
        logger.debug("stream does not exist! {}", new String(ticket.getBytes()));
      }
    }
 

/**
 * 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);
  }
}
 

/**
 * Deserializes a new instance from the protocol buffer ticket.
 */
public static JobsFlightTicket from(Ticket ticket) {
  try {
    return MAPPER.readValue(ticket.getBytes(), JobsFlightTicket.class);
  } catch (IOException e) {
    throw Throwables.propagate(e);
  }
}
 

/**
 *  Creates a new protocol buffer Ticket by serializing to JSON.
 */
public Ticket toTicket() {
  try {
    return new Ticket(MAPPER.writeValueAsBytes(this));
  } catch (JsonProcessingException e) {
    throw Throwables.propagate(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 void planParallelized(PlanningSet planningSet) {
  logger.debug("plan parallel called, collecting endpoints");
  List<FlightEndpoint> endpoints = Lists.newArrayList();
  FlightEndpoint screenEndpoint = null;
  for (Wrapper wrapper : planningSet.getFragmentWrapperMap().values()) {
    String majorId = String.valueOf(wrapper.getMajorFragmentId());
    try {
      Boolean isWriter = wrapper.getNode().getRoot().accept(new WriterVisitor(), false);
      if (!isWriter) {
        continue;
      }
    } catch (Throwable throwable) {
      logger.warn("unable to complete visitor ", throwable);
    }


    CoreOperatorType op = CoreOperatorType.valueOf(wrapper.getNode().getRoot().getOperatorType());
    boolean isScreen = CoreOperatorType.SCREEN.equals(op) && planningSet.getFragmentWrapperMap().size() > 1;

    logger.info("Creating tickets for {}. MajorId {}", wrapper.getNode().getRoot(), majorId);
    for (int i = 0; i < wrapper.getAssignedEndpoints().size(); i++) {
      CoordinationProtos.NodeEndpoint endpoint = wrapper.getAssignedEndpoint(i);
      logger.warn("Creating ticket for {} . MajorId {} and index {}", wrapper.getNode().getRoot(), majorId, i);
      Ticket ticket = new Ticket(JOINER.join(
        majorId,
        String.valueOf(i),
        endpoint.getAddress(),
        endpoint.getUserPort()
      ).getBytes());
      int port = Integer.parseInt(PropertyHelper.getFromEnvProperty("dremio.formation.port", Integer.toString(FormationPlugin.FLIGHT_PORT)));
      String host = PropertyHelper.getFromEnvProperty("dremio.formation.host", endpoint.getAddress());
      Location location = Location.forGrpcInsecure(host, port);

      if (isScreen) {
        screenEndpoint = new FlightEndpoint(ticket, location);
      } else {
        endpoints.add(new FlightEndpoint(ticket, location));
      }
    }
  }
  if (screenEndpoint != null && endpoints.isEmpty()) {
    logger.info("Adding a screen endpoint as its the only possible endpoint!");
    endpoints.add(screenEndpoint);
  } else {
    logger.warn("Skipping a screen as it lies above the writer.");
  }
  logger.debug("built {} parallel endpoints", endpoints.size());
  future.complete(endpoints);
}
 

private static FlightDescriptor getDescriptor(Ticket ticket) {
  String path = new String(ticket.getBytes());
  String[] pathParts = path.split(":");
  return FlightDescriptor.path(pathParts[0], pathParts[1], pathParts[2]);
}
 

public FlightDataReader(Broadcast<FlightDataSourceReader.FactoryOptions> options) {
  this.options = options;
  this.location = Location.forGrpcInsecure(options.value().getHost(), options.value().getPort());
  this.ticket = new Ticket(options.value().getTicket());
}
 

@Override
public void getStream(CallContext callContext, Ticket ticket, ServerStreamListener serverStreamListener) {
  /* Note that we do not trim record batches that we receive from the Job Results Store. This may result
   * in sending record that the client does not care about, or in the case of sequential requests, sending
   * duplicate records. We may want to trim the record batches if this presents a problem.
   */
  try {
    final JobsFlightTicket jobsFlightTicket = JobsFlightTicket.from(ticket);
    final JobProtobuf.JobId jobId = JobProtobuf.JobId.newBuilder().setId(jobsFlightTicket.getJobId()).build();
    final int offset = jobsFlightTicket.getOffset();
    final int limit = jobsFlightTicket.getLimit();

    try (final JobDataFragment jobDataFragment = jobsService.get().getJobData(JobsProtoUtil.toStuff(jobId), offset, limit)) {
      final Schema schema = jobDataFragment.getSchema();
      try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
        serverStreamListener.start(root);
        for (RecordBatchHolder holder : jobDataFragment.getRecordBatches()) {
          // iterate over the columns
          int numRecords = holder.size();
          for (int i = 0; i < schema.getFields().size(); i++) {
            ValueVector vector = root.getVector(schema.getFields().get(i).getName());
            ValueVector dataVector = holder.getData().getVectors().get(i);
            int k = 0; // index at which value need to written in "vector" from "dataVector"
            // iterate over values in the column to copy data
            for (int j = holder.getStart(); j < holder.getEnd(); j++, k++ ) {
              // Copy value at dataVector[j] into vector[k]
              vector.copyFromSafe(j, k, dataVector);
            }
            vector.setValueCount(numRecords);
            root.setRowCount(numRecords);
          }
          serverStreamListener.putNext();
          root.allocateNew();
        }
      }
      serverStreamListener.completed();
    }
  } catch (UserException ue) {
    serverStreamListener.error(JobsRpcUtils.toStatusRuntimeException(ue));
  } catch (Exception e) {
    serverStreamListener.error(Status.UNKNOWN.withCause(e).withDescription(e.getMessage()).asException());
  }
}
 

@Override
public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
  UserException e = UserException.unsupportedError().message(expectedMessage).buildSilently();
  listener.error(toStatusRuntimeException(e));
}