org.apache.arrow.flight.FlightDescriptor Java Examples

private FlightInfo getInfoImpl(CallContext callContext, FlightDescriptor descriptor, String sql) {
  try {
    final CreatePreparedStatementReq req =
      CreatePreparedStatementReq.newBuilder()
        .setSqlQuery(sql)
        .build();

    UserRequest request = new UserRequest(RpcType.CREATE_PREPARED_STATEMENT, req);
    Prepare prepare = new Prepare();

    UserBitShared.ExternalId externalId = submitWork(callContext, request, prepare);
    return prepare.getInfo(descriptor, externalId);
  } catch (Exception e) {
    throw Status.ABORTED.asRuntimeException();
  }
}
 

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();
  }
}
 

@Test
public void connect() throws Exception {
  certs();
  InetAddress ip = InetAddress.getLocalHost();
  Location location = Location.forGrpcTls(ip.getHostName(), 47470);
  try (FlightClient c = flightClient(getAllocator(), location)) {
    c.authenticate(new BasicClientAuthHandler(SystemUser.SYSTEM_USERNAME, null));
    String sql = "select * from sys.options";
    FlightInfo info = c.getInfo(FlightDescriptor.command(sql.getBytes()));
    long total = info.getEndpoints().stream()
      .map(this::submit)
      .map(TestSslFlightEndpoint::get)
      .mapToLong(Long::longValue)
      .sum();

    Assert.assertTrue(total > 1);
    System.out.println(total);
  }
}
 

@Test
public void connect() throws Exception {
  InetAddress ip = InetAddress.getLocalHost();
  Location location = Location.forGrpcInsecure(ip.getHostName(), 47470);
  try (FlightClient c = flightClient(getAllocator(), location)) {
    c.authenticate(new BasicClientAuthHandler(SystemUser.SYSTEM_USERNAME, null));
    String sql = "select * from sys.options";
    FlightInfo info = c.getInfo(FlightDescriptor.command(sql.getBytes()));
    long total = info.getEndpoints().stream()
      .map(this::submit)
      .map(TestFlightEndpoint::get)
      .mapToLong(Long::longValue)
      .sum();

    Assert.assertTrue(total > 1);
    System.out.println(total);
  }
}
 

private FlightInfo getInfoParallel(CallContext callContext, FlightDescriptor descriptor, String sql) {
  String queryId = QueryIdHelper.getQueryId(ExternalIdHelper.toQueryId(ExternalIdHelper.generateExternalId()));
  sql = String.format("create table flight.\"%s\" as (%s)", queryId, sql);
  logger.debug("Submitting ctas {}", sql);
  FlightInfo ticket = getInfoImpl(callContext, descriptor, sql);
  return getCoalesce(callContext, descriptor, ticket.getEndpoints().get(0).getTicket());
}
 

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 FlightInfo getFlightInfo(CallContext callContext, FlightDescriptor descriptor) {
  logger.debug("checking if parallel query");
  boolean isParallel = validator.getSessionOptions(callContext).isParallel();
  logger.debug("checking if parallel query: result {}", isParallel);
  if (isParallel) {
    return getInfoParallel(callContext, descriptor, new String(descriptor.getCommand()));
  }
  return getInfoImpl(callContext, descriptor, new String(descriptor.getCommand()));
}
 

private List<InputPartition<ColumnarBatch>> planBatchInputPartitionsParallel() {

    try (FlightClient client = clientFactory.apply()) {
      FlightInfo info = client.getInfo(FlightDescriptor.command(sql.getBytes()));
      return planBatchInputPartitionsSerial(info);
    } catch (InterruptedException e) {
      throw new RuntimeException(e);
    }
  }
 

@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);
    }
 

public Stream.Producer putStream(FlightDescriptor descriptor, Schema schema) {
  logger.debug("putting stream for descriptor {}", descriptor);
  boolean isIn = holders.containsKey(descriptor);
  Stream h = holders.computeIfAbsent(descriptor, (t) -> new Stream(allocator, descriptor));
  logger.debug("generated holder for {}, was it created: {}. Now setting schema", descriptor, isIn);
  if (schema != null) {
    logger.debug("schema is now set for {}", descriptor);
    h.setSchema(schema);
  } else {
    logger.debug("schema was not set for {}", descriptor);
  }
  return h.getProducer();
}
 

public FormationRecordWriter(String path, FormationPlugin.FormationFlightProducer store, ExecProtos.FragmentHandle fragmentHandle) {
  super();
  this.store = store;
  this.descriptor = FlightDescriptor.path(
    QueryIdHelper.getQueryId(fragmentHandle.getQueryId()),
    String.valueOf(fragmentHandle.getMajorFragmentId()),
    String.valueOf(fragmentHandle.getMinorFragmentId())
  );
}
 

@Override
//TODO (DX-19234): Implement me. Standard Flight procedure requires clients call getFlightInfo before getting stream
public FlightInfo getFlightInfo(CallContext callContext, FlightDescriptor flightDescriptor) {
  throw Status.UNIMPLEMENTED.asRuntimeException();
}
 

private FlightDescriptor getDescriptor(String path) {
  return FlightDescriptor.command(path.getBytes());
}
 

private void buildIfNecessary() {
    if (config != null) {
      return;
    }

    if (infos == null) {
      infos = clients.stream()
        .map(c -> c.getInfo(FlightDescriptor.path(key.getName())))
        .collect(Collectors.toList());
    }

    Preconditions.checkArgument(!infos.isEmpty());
    Schema schema = null;
    long records = 0;
    List<FlightEndpoint> endpoints = new ArrayList<>();
    for (FlightInfo info : infos) {
      schema = info.getSchema();
      records += info.getRecords();
      endpoints.addAll(info.getEndpoints());
    }

    config = new DatasetConfig()
      .setFullPathList(key.getComponents())
      .setName(key.getName())
      .setType(DatasetType.PHYSICAL_DATASET)
      .setId(new EntityId().setId(UUID.randomUUID().toString()))
      .setReadDefinition(new ReadDefinition()
        .setScanStats(new ScanStats().setRecordCount(records)
          .setScanFactor(ScanCostFactor.PARQUET.getFactor())))
      .setOwner(SystemUser.SYSTEM_USERNAME)
      .setPhysicalDataset(new PhysicalDataset())
      .setRecordSchema(new BatchSchema(schema.getFields()).toByteString())
      .setSchemaVersion(DatasetHelper.CURRENT_VERSION);

    splits = new ArrayList<>();
    List<DatasetSplit> dSplits = Lists.newArrayList();
//     int i =0;
    for (FlightEndpoint ep : endpoints) {

      List<Location> locations = ep.getLocations();
      if (locations.size() > 1) {
        throw new UnsupportedOperationException("I dont know what more than one location means, not handling it");
      }
      DatasetSplitAffinity a = DatasetSplitAffinity.of(locations.get(0).getUri().getHost(), 100d);

//       split.setSplitKey(Integer.toString(i));
      Flight.Ticket ticket = Flight.Ticket.newBuilder().setTicket(ByteString.copyFrom(ep.getTicket().getBytes())).build();
      dSplits.add(DatasetSplit.of(ImmutableList.of(a), records / endpoints.size(), records, ticket::writeTo));
    }
    splits.add(PartitionChunk.of(dSplits));
  }
 

public Consumer(BufferAllocator allocator, FlightDescriptor descriptor, BlockingQueue<ArrowRecordBatch> exchanger) {
  this.allocator = allocator.newChildAllocator("consumer", 0, Long.MAX_VALUE);
  this.descriptor = descriptor;
  this.exchanger = exchanger;
}
 

public Producer(BufferAllocator allocator, FlightDescriptor descriptor, BlockingQueue<ArrowRecordBatch> exchanger) {
  this.allocator = allocator.newChildAllocator("producer", 0, Long.MAX_VALUE);
  this.descriptor = descriptor;
  this.exchanger = exchanger;
}
 

public Stream(BufferAllocator allocator, FlightDescriptor descriptor) {
  this.descriptor = descriptor;
  this.consumer = new Consumer(allocator, descriptor, exchanger);
  this.producer = new Producer(allocator, descriptor, exchanger);
}
 

@Override
public FlightInfo getFlightInfo(CallContext callContext, FlightDescriptor flightDescriptor) {
  throw Status.UNIMPLEMENTED.asRuntimeException();
}
 

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]);
}
 

@Override
public SchemaResult getSchema(CallContext context, FlightDescriptor descriptor) {
  FlightInfo info = getInfoImpl(context, descriptor, new String(descriptor.getCommand()));
  return new SchemaResult(info.getSchema());
}