Java Code Examples for org.apache.arrow.flight.FlightInfo

The following examples show how to use org.apache.arrow.flight.FlightInfo. These examples are extracted from open source projects. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source Project: dremio-flight-connector   Source File: Producer.java    License: Apache License 2.0 6 votes vote down vote up
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();
  }
}
 
Example 2
Source Project: dremio-flight-connector   Source File: Producer.java    License: Apache License 2.0 6 votes vote down vote up
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();
  }
}
 
Example 3
@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);
  }
}
 
Example 4
@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);
  }
}
 
Example 5
Source Project: dremio-flight-connector   Source File: Producer.java    License: Apache License 2.0 5 votes vote down vote up
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);
}
 
Example 6
Source Project: dremio-flight-connector   Source File: Producer.java    License: Apache License 2.0 5 votes vote down vote up
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());
}
 
Example 7
Source Project: dremio-flight-connector   Source File: Producer.java    License: Apache License 2.0 5 votes vote down vote up
@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()));
}
 
Example 8
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);
    }
  }
 
Example 9
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;
}
 
Example 10
Source Project: dremio-flight-connector   Source File: Producer.java    License: Apache License 2.0 4 votes vote down vote up
@Override
public SchemaResult getSchema(CallContext context, FlightDescriptor descriptor) {
  FlightInfo info = getInfoImpl(context, descriptor, new String(descriptor.getCommand()));
  return new SchemaResult(info.getSchema());
}
 
Example 11
Source Project: dremio-flight-connector   Source File: Producer.java    License: Apache License 2.0 4 votes vote down vote up
@Override
public void listFlights(CallContext callContext, Criteria arg0, StreamListener<FlightInfo> list) {
  list.onCompleted();
}
 
Example 12
Source Project: dremio-flight-connector   Source File: FormationPlugin.java    License: Apache License 2.0 4 votes vote down vote up
@Override
public void listFlights(CallContext callContext, Criteria criteria, StreamListener<FlightInfo> streamListener) {
  throw Status.UNIMPLEMENTED.asRuntimeException();
}
 
Example 13
Source Project: dremio-flight-connector   Source File: FormationPlugin.java    License: Apache License 2.0 4 votes vote down vote up
@Override
public FlightInfo getFlightInfo(CallContext callContext, FlightDescriptor flightDescriptor) {
  throw Status.UNIMPLEMENTED.asRuntimeException();
}
 
Example 14
Source Project: dremio-flight-connector   Source File: DatasetBuilder.java    License: Apache License 2.0 4 votes vote down vote up
public DatasetBuilder(List<FlightClient> clients, EntityPath key, List<FlightInfo> infos) {
  super();
  this.clients = clients;
  this.key = key;
  this.infos = infos;
}
 
Example 15
Source Project: dremio-flight-connector   Source File: DatasetBuilder.java    License: Apache License 2.0 4 votes vote down vote up
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));
  }
 
Example 16
Source Project: dremio-oss   Source File: JobsFlightProducer.java    License: Apache License 2.0 4 votes vote down vote up
@Override
public void listFlights(CallContext callContext, Criteria criteria, StreamListener<FlightInfo> streamListener) {
  throw Status.UNIMPLEMENTED.asRuntimeException();
}
 
Example 17
Source Project: dremio-oss   Source File: JobsFlightProducer.java    License: Apache License 2.0 4 votes vote down vote up
@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();
}