org.apache.arrow.flight.FlightInfo 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>> planBatchInputPartitionsSerial(FlightInfo info) {
  LOGGER.warn("planning partitions for endpoints {}", Joiner.on(", ").join(info.getEndpoints().stream().map(e -> e.getLocations().get(0).getUri().toString()).collect(Collectors.toList())));
  List<InputPartition<ColumnarBatch>> batches = info.getEndpoints().stream().map(endpoint -> {
    Location location = (endpoint.getLocations().isEmpty()) ?
      Location.forGrpcInsecure(defaultLocation.getUri().getHost(), defaultLocation.getUri().getPort()) :
      endpoint.getLocations().get(0);
    FactoryOptions options = dataSourceOptions.value().copy(location, endpoint.getTicket().getBytes());
    LOGGER.warn("X1 {}", dataSourceOptions.value());
    return new FlightDataReaderFactory(lazySparkContext().broadcast(options));
  }).collect(Collectors.toList());
  LOGGER.info("Created {} batches from arrow endpoints", batches.size());
  return batches;
}
 

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

@Override
public void listFlights(CallContext callContext, Criteria criteria, StreamListener<FlightInfo> streamListener) {
  throw Status.UNIMPLEMENTED.asRuntimeException();
}
 

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 DatasetBuilder(List<FlightClient> clients, EntityPath key, List<FlightInfo> infos) {
  super();
  this.clients = clients;
  this.key = key;
  this.infos = infos;
}
 

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

@Override
public void listFlights(CallContext callContext, Criteria criteria, StreamListener<FlightInfo> streamListener) {
  throw Status.UNIMPLEMENTED.asRuntimeException();
}
 

@Override
public void listFlights(CallContext callContext, Criteria arg0, StreamListener<FlightInfo> list) {
  list.onCompleted();
}
 

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