com.google.cloud.spanner.PartitionOptions Java Examples

void partitionReadUsingIndex() {
  // [START partition_read_using_index]
  final BatchReadOnlyTransaction txn =
      batchClient.batchReadOnlyTransaction(TimestampBound.strong());
  List<Partition> partitions =
      txn.partitionReadUsingIndex(
          PartitionOptions.getDefaultInstance(),
          "Singers",
          "SingerId",
          KeySet.all(),
          Arrays.asList("SingerId", "FirstName", "LastName"));

  for (Partition p : partitions) {
    try (ResultSet results = txn.execute(p)) {
      while (results.next()) {
        long singerId = results.getLong(0);
        String firstName = results.getString(1);
        String lastName = results.getString(2);
        System.out.println("[" + singerId + "] " + firstName + " " + lastName);
      }
    }
  }
  // [END partition_read_using_index]
}
 

void partitionRead() {
  // [START partition_read]
  final BatchReadOnlyTransaction txn =
      batchClient.batchReadOnlyTransaction(TimestampBound.strong());
  List<Partition> partitions =
      txn.partitionRead(
          PartitionOptions.getDefaultInstance(),
          "Singers",
          KeySet.all(),
          Arrays.asList("SingerId", "FirstName", "LastName"));
  for (final Partition p : partitions) {
    try (ResultSet results = txn.execute(p)) {
      while (results.next()) {
        long singerId = results.getLong(0);
        String firstName = results.getString(1);
        String lastName = results.getString(2);
        System.out.println("[" + singerId + "] " + firstName + " " + lastName);
      }
    }
  }
  // [END partition_read]
}
 

void partitionQuery() {
  // [START partition_query]
  final BatchReadOnlyTransaction txn =
      batchClient.batchReadOnlyTransaction(TimestampBound.strong());
  List<Partition> partitions =
      txn.partitionQuery(
          PartitionOptions.getDefaultInstance(),
          Statement.of("SELECT SingerId, FirstName, LastName FROM Singers"));

  for (final Partition p : partitions) {
    try (ResultSet results = txn.execute(p)) {
      while (results.next()) {
        long singerId = results.getLong(0);
        String firstName = results.getString(1);
        String lastName = results.getString(2);
        System.out.println("[" + singerId + "] " + firstName + " " + lastName);
      }
    }
  }
  // [END partition_query]

}
 

@Override
public List<Partition> partitionReadUsingIndex(PartitionOptions partitionOptions, String table,
    String index, KeySet keys, Iterable<String> columns, ReadOption... options)
    throws SpannerException {
  throw SpannerExceptionFactory.newSpannerException(ErrorCode.UNIMPLEMENTED,
      METHOD_NOT_IMPLEMENTED);
}
 

@Override
public List<Partition> partitionQuery(PartitionOptions partitionOptions, Statement statement,
    QueryOption... options) throws SpannerException {
  checkTransaction();
  if (batchReadOnlyTransaction != null) {
    return batchReadOnlyTransaction.partitionQuery(partitionOptions, statement, options);
  }
  throw SpannerExceptionFactory.newSpannerException(ErrorCode.FAILED_PRECONDITION,
      METHOD_ONLY_IN_BATCH_READONLY);
}
 

@Nullable
abstract PartitionOptions getPartitionOptions();
 

/**
 * This example showcases how to create a batch client, partition a query, and concurrently read
 * from multiple partitions.
 */
public static void main(String[] args) throws InterruptedException {
  if (args.length != 2) {
    System.err.println("Usage: BatchSample <instance_id> <database_id>");
    return;
  }

  /*
   * CREATE TABLE Singers (
   *   SingerId   INT64 NOT NULL,
   *   FirstName  STRING(1024),
   *   LastName   STRING(1024),
   *   SingerInfo BYTES(MAX),
   * ) PRIMARY KEY (SingerId);
   */

  String instanceId = args[0];
  String databaseId = args[1];

  SpannerOptions options = SpannerOptions.newBuilder().build();
  Spanner spanner = options.getService();

  // [START spanner_batch_client]
  int numThreads = Runtime.getRuntime().availableProcessors();
  ExecutorService executor = Executors.newFixedThreadPool(numThreads);

  // Statistics
  int totalPartitions;
  AtomicInteger totalRecords = new AtomicInteger(0);

  try {
    BatchClient batchClient =
        spanner.getBatchClient(DatabaseId.of(options.getProjectId(), instanceId, databaseId));

    final BatchReadOnlyTransaction txn =
        batchClient.batchReadOnlyTransaction(TimestampBound.strong());

    // A Partition object is serializable and can be used from a different process.
    List<Partition> partitions =
        txn.partitionQuery(
            PartitionOptions.getDefaultInstance(),
            Statement.of("SELECT SingerId, FirstName, LastName FROM Singers"));

    totalPartitions = partitions.size();

    for (final Partition p : partitions) {
      executor.execute(
          () -> {
            try (ResultSet results = txn.execute(p)) {
              while (results.next()) {
                long singerId = results.getLong(0);
                String firstName = results.getString(1);
                String lastName = results.getString(2);
                System.out.println("[" + singerId + "] " + firstName + " " + lastName);
                totalRecords.getAndIncrement();
              }
            }
          });
    }
  } finally {
    executor.shutdown();
    executor.awaitTermination(1, TimeUnit.HOURS);
    spanner.close();
  }

  double avgRecordsPerPartition = 0.0;
  if (totalPartitions != 0) {
    avgRecordsPerPartition = (double) totalRecords.get() / totalPartitions;
  }
  System.out.println("totalPartitions=" + totalPartitions);
  System.out.println("totalRecords=" + totalRecords);
  System.out.println("avgRecordsPerPartition=" + avgRecordsPerPartition);
  // [END spanner_batch_client]
}
 

@Test
public void readAllPipeline() throws Exception {
  Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
  TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);

  SpannerConfig spannerConfig =
      SpannerConfig.create()
          .withProjectId("test")
          .withInstanceId("123")
          .withDatabaseId("aaa")
          .withServiceFactory(serviceFactory);

  PCollectionView<Transaction> tx =
      pipeline.apply(
          "tx",
          SpannerIO.createTransaction()
              .withSpannerConfig(spannerConfig)
              .withTimestampBound(timestampBound));

  PCollection<ReadOperation> reads =
      pipeline.apply(
          Create.of(
              ReadOperation.create().withQuery("SELECT * FROM users"),
              ReadOperation.create().withTable("users").withColumns("id", "name")));

  PCollection<Struct> one =
      reads.apply(
          "read all", SpannerIO.readAll().withSpannerConfig(spannerConfig).withTransaction(tx));

  BatchTransactionId txId = new FakeBatchTransactionId("tx");
  when(mockBatchTx.getBatchTransactionId()).thenReturn(txId);

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound))
      .thenReturn(mockBatchTx);

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class)))
      .thenReturn(mockBatchTx);

  Partition fakePartition =
      FakePartitionFactory.createFakeReadPartition(ByteString.copyFromUtf8("partition"));
  when(mockBatchTx.partitionQuery(
          any(PartitionOptions.class), eq(Statement.of("SELECT * FROM users"))))
      .thenReturn(Arrays.asList(fakePartition, fakePartition));
  when(mockBatchTx.partitionRead(
          any(PartitionOptions.class),
          eq("users"),
          eq(KeySet.all()),
          eq(Arrays.asList("id", "name"))))
      .thenReturn(Arrays.asList(fakePartition));

  when(mockBatchTx.execute(any(Partition.class)))
      .thenReturn(
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(0, 2)),
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(2, 4)),
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(4, 6)));

  PAssert.that(one).containsInAnyOrder(FAKE_ROWS);

  pipeline.run();
}
 

@Test
public void readPipeline() throws Exception {
  Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
  TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);

  SpannerConfig spannerConfig =
      SpannerConfig.create()
          .withProjectId("test")
          .withInstanceId("123")
          .withDatabaseId("aaa")
          .withServiceFactory(serviceFactory);

  PCollection<Struct> one =
      pipeline.apply(
          "read q",
          SpannerIO.read()
              .withSpannerConfig(spannerConfig)
              .withQuery("SELECT * FROM users")
              .withTimestampBound(timestampBound));
  FakeBatchTransactionId txId = new FakeBatchTransactionId("readPipelineTest");
  when(mockBatchTx.getBatchTransactionId()).thenReturn(txId);

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound))
      .thenReturn(mockBatchTx);

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class)))
      .thenReturn(mockBatchTx);

  Partition fakePartition =
      FakePartitionFactory.createFakeQueryPartition(ByteString.copyFromUtf8("one"));
  when(mockBatchTx.partitionQuery(
          any(PartitionOptions.class), eq(Statement.of("SELECT * FROM users"))))
      .thenReturn(Arrays.asList(fakePartition, fakePartition));

  when(mockBatchTx.execute(any(Partition.class)))
      .thenReturn(
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(0, 2)),
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(2, 6)));

  PAssert.that(one).containsInAnyOrder(FAKE_ROWS);

  pipeline.run();
}
 

@Test
public void runReadUsingIndex() throws Exception {
  Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
  TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);

  SpannerConfig spannerConfig =
      SpannerConfig.create()
          .withProjectId("test")
          .withInstanceId("123")
          .withDatabaseId("aaa")
          .withServiceFactory(serviceFactory);

  PCollection<Struct> one =
      pipeline.apply(
          "read q",
          SpannerIO.read()
              .withTimestamp(Timestamp.now())
              .withSpannerConfig(spannerConfig)
              .withTable("users")
              .withColumns("id", "name")
              .withIndex("theindex")
              .withTimestampBound(timestampBound));

  FakeBatchTransactionId id = new FakeBatchTransactionId("runReadUsingIndexTest");
  when(mockBatchTx.getBatchTransactionId()).thenReturn(id);

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound))
      .thenReturn(mockBatchTx);
  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class)))
      .thenReturn(mockBatchTx);

  Partition fakePartition =
      FakePartitionFactory.createFakeReadPartition(ByteString.copyFromUtf8("one"));

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(id)).thenReturn(mockBatchTx);
  when(mockBatchTx.partitionReadUsingIndex(
          any(PartitionOptions.class),
          eq("users"),
          eq("theindex"),
          eq(KeySet.all()),
          eq(Arrays.asList("id", "name"))))
      .thenReturn(Arrays.asList(fakePartition, fakePartition, fakePartition));

  when(mockBatchTx.execute(any(Partition.class)))
      .thenReturn(
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(0, 2)),
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(2, 4)),
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(4, 6)));

  PAssert.that(one).containsInAnyOrder(FAKE_ROWS);

  pipeline.run();
}
 

@Test
public void runRead() throws Exception {
  Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
  TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);

  SpannerConfig spannerConfig =
      SpannerConfig.create()
          .withProjectId("test")
          .withInstanceId("123")
          .withDatabaseId("aaa")
          .withServiceFactory(serviceFactory);

  PCollection<Struct> one =
      pipeline.apply(
          "read q",
          SpannerIO.read()
              .withSpannerConfig(spannerConfig)
              .withTable("users")
              .withColumns("id", "name")
              .withTimestampBound(timestampBound));

  FakeBatchTransactionId id = new FakeBatchTransactionId("runReadTest");
  when(mockBatchTx.getBatchTransactionId()).thenReturn(id);

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound))
      .thenReturn(mockBatchTx);
  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class)))
      .thenReturn(mockBatchTx);

  Partition fakePartition =
      FakePartitionFactory.createFakeReadPartition(ByteString.copyFromUtf8("one"));

  when(mockBatchTx.partitionRead(
          any(PartitionOptions.class),
          eq("users"),
          eq(KeySet.all()),
          eq(Arrays.asList("id", "name"))))
      .thenReturn(Arrays.asList(fakePartition, fakePartition, fakePartition));
  when(mockBatchTx.execute(any(Partition.class)))
      .thenReturn(
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(0, 2)),
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(2, 4)),
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(4, 6)));

  PAssert.that(one).containsInAnyOrder(FAKE_ROWS);

  pipeline.run();
}
 

@Test
public void runQuery() throws Exception {
  Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
  TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);

  SpannerConfig spannerConfig =
      SpannerConfig.create()
          .withProjectId("test")
          .withInstanceId("123")
          .withDatabaseId("aaa")
          .withServiceFactory(serviceFactory);

  PCollection<Struct> one =
      pipeline.apply(
          "read q",
          SpannerIO.read()
              .withSpannerConfig(spannerConfig)
              .withQuery("SELECT * FROM users")
              .withTimestampBound(timestampBound));

  FakeBatchTransactionId id = new FakeBatchTransactionId("runQueryTest");
  when(mockBatchTx.getBatchTransactionId()).thenReturn(id);

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound))
      .thenReturn(mockBatchTx);
  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class)))
      .thenReturn(mockBatchTx);

  Partition fakePartition =
      FakePartitionFactory.createFakeQueryPartition(ByteString.copyFromUtf8("one"));

  when(mockBatchTx.partitionQuery(
          any(PartitionOptions.class), eq(Statement.of("SELECT * FROM users"))))
      .thenReturn(Arrays.asList(fakePartition, fakePartition));
  when(mockBatchTx.execute(any(Partition.class)))
      .thenReturn(
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(0, 2)),
          ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(2, 6)));

  PAssert.that(one).containsInAnyOrder(FAKE_ROWS);

  pipeline.run();
}
 

public Read withPartitionOptions(PartitionOptions partitionOptions) {
  return withReadOperation(getReadOperation().withPartitionOptions(partitionOptions));
}
 

@Nullable
abstract PartitionOptions getPartitionOptions();
 

public ReadOperation withPartitionOptions(PartitionOptions partitionOptions) {
  return toBuilder().setPartitionOptions(partitionOptions).build();
}
 

@Nullable
abstract PartitionOptions getPartitionOptions();
 

public static ReadOperation create() {
  return new AutoValue_ReadOperation.Builder()
      .setPartitionOptions(PartitionOptions.getDefaultInstance())
      .setKeySet(KeySet.all())
      .build();
}
 

protected List<Partition> partitionQuery(com.google.cloud.spanner.Statement statement) {
  PartitionOptions po = PartitionOptions.getDefaultInstance();
  return connection.getTransaction().partitionQuery(po, statement);
}
 

@Override
public List<Partition> partitionRead(PartitionOptions partitionOptions, String table, KeySet keys,
    Iterable<String> columns, ReadOption... options) throws SpannerException {
  throw SpannerExceptionFactory.newSpannerException(ErrorCode.UNIMPLEMENTED,
      METHOD_NOT_IMPLEMENTED);
}
 

public Read withPartitionOptions(PartitionOptions partitionOptions) {
  return withReadOperation(getReadOperation().withPartitionOptions(partitionOptions));
}
 

abstract Builder setPartitionOptions(PartitionOptions partitionOptions); 

abstract Builder setPartitionOptions(PartitionOptions partitionOptions); 

abstract Builder setPartitionOptions(PartitionOptions partitionOptions);