diff --git a/Bigtable/src/DataClient.php b/Bigtable/src/DataClient.php index f652bd7c8294..4522ad465200 100644 --- a/Bigtable/src/DataClient.php +++ b/Bigtable/src/DataClient.php @@ -403,4 +403,37 @@ private function convertToArray(Row $row) } return $families; } + + /** + * Returns a sample of row keys in the table. The returned row keys will + * delimit contiguous sections of the table of approximately equal size, + * which can be used to break up the data for distributed tasks like + * mapreduces. + * + * Example: + * ``` + * $rowKeyStream = $dataClient->sampleRowKeys(); + * foreach ($rowKeyStream as $rowKey) { + * print_r($rowKey) . PHP_EOL; + * } + * ``` + * + * @param array $options [optional] Configuration options. + * @return \Generator A list of associative arrays, each with the keys `rowKey` and `offset`. + * @throws ApiException if the remote call fails or operation fails + */ + public function sampleRowKeys(array $options = []) + { + $stream = $this->bigtableClient->sampleRowKeys( + $this->tableName, + $options + $this->options + ); + + foreach ($stream->readAll() as $response) { + yield [ + 'rowKey' => $response->getRowKey(), + 'offset' => $response->getOffsetBytes() + ]; + } + } } diff --git a/Bigtable/tests/Snippet/DataClientTest.php b/Bigtable/tests/Snippet/DataClientTest.php index 0ae852545a83..cd96cd469ac0 100644 --- a/Bigtable/tests/Snippet/DataClientTest.php +++ b/Bigtable/tests/Snippet/DataClientTest.php @@ -36,6 +36,7 @@ use Google\Cloud\Bigtable\V2\Row; use Google\Cloud\Bigtable\V2\RowRange; use Google\Cloud\Bigtable\V2\RowSet; +use Google\Cloud\Bigtable\V2\SampleRowKeysResponse; use Google\Cloud\Core\Testing\Snippet\SnippetTestCase; use Google\Cloud\Core\Testing\TestHelpers; use Google\Protobuf\StringValue; @@ -350,6 +351,35 @@ public function testReadModifyWriteRowIncrement() ); } + public function testSampleRowKeys() + { + $sampleRowKeyResponses[] = (new SampleRowKeysResponse) + ->setRowKey('rk1') + ->setOffsetBytes(1); + + $this->serverStream->readAll() + ->shouldBeCalled() + ->willReturn( + $this->arrayAsGenerator($sampleRowKeyResponses) + ); + $this->bigtableClient->sampleRowKeys(self::TABLE_NAME, []) + ->shouldBeCalled() + ->willReturn( + $this->serverStream->reveal() + ); + $snippet = $this->snippetFromMethod(DataClient::class, 'sampleRowKeys'); + $snippet->addLocal('dataClient', $this->dataClient); + $res = $snippet->invoke('rowKeyStream'); + $expectedRowKeys = [ + 'rowKey' => 'rk1', + 'offset' => 1 + ]; + $this->assertEquals( + print_r($expectedRowKeys, true), + $res->output() + ); + } + private function setUpReadRowsResponse() { $readRowsResponse = new ReadRowsResponse; diff --git a/Bigtable/tests/System/DataClientSampleRowKeysTest.php b/Bigtable/tests/System/DataClientSampleRowKeysTest.php new file mode 100644 index 000000000000..945da3c12b84 --- /dev/null +++ b/Bigtable/tests/System/DataClientSampleRowKeysTest.php @@ -0,0 +1,88 @@ + [ + 'cf1' => [ + 'cq1' => [ + 'value' => 'value1', + 'timeStamp' => 5000 + ] + ] + ], + 'rk2' => [ + 'cf1' => [ + 'cq2' => [ + 'value' => 'value2', + 'timeStamp' => 5000 + ] + ] + ], + 'rk3' => [ + 'cf1' => [ + 'cq3' => [ + 'value' => 'value3', + 'timeStamp' => 5000 + ] + ] + ], + 'rk4' => [ + 'cf1' => [ + 'cq4' => [ + 'value' => 'value4', + 'timeStamp' => 5000 + ] + ] + ], + 'rk5' => [ + 'cf1' => [ + 'cq5' => [ + 'value' => 'value5', + 'timeStamp' => 5000 + ] + ] + ] + ]; + self::$dataClient->upsert($insertRows); + } + + public function testSampleRowKeys() + { + $rowKeysStream = self::$dataClient->sampleRowKeys(); + $rowKeys = iterator_to_array($rowKeysStream); + $expectedRowKeys = [ + [ + 'rowKey' => '', + 'offset' => 805306368 + ] + ]; + $this->assertEquals($expectedRowKeys, $rowKeys); + } +} diff --git a/Bigtable/tests/Unit/DataClientTest.php b/Bigtable/tests/Unit/DataClientTest.php index 82d21a17e453..7f09c7eaa560 100644 --- a/Bigtable/tests/Unit/DataClientTest.php +++ b/Bigtable/tests/Unit/DataClientTest.php @@ -34,6 +34,7 @@ use Google\Cloud\Bigtable\V2\Row; use Google\Cloud\Bigtable\V2\RowRange; use Google\Cloud\Bigtable\V2\RowSet; +use Google\Cloud\Bigtable\V2\SampleRowKeysResponse; use Google\Cloud\Bigtable\Filter; use Google\Rpc\Code; use Google\Rpc\Status; @@ -637,6 +638,40 @@ public function testReadModifyWriteRowIncrement() $this->assertEquals($expectedRow, $row); } + public function testSampleRowKeys() + { + $sampleRowKeyResponses[] = (new SampleRowKeysResponse) + ->setRowKey('rk1') + ->setOffsetBytes(1); + $sampleRowKeyResponses[] = (new SampleRowKeysResponse) + ->setRowKey('rk2') + ->setOffsetBytes(2); + + $this->serverStream->readAll() + ->shouldBeCalled() + ->willReturn( + $this->arrayAsGenerator($sampleRowKeyResponses) + ); + $this->bigtableClient->sampleRowKeys(self::TABLE_NAME, $this->options) + ->shouldBeCalled() + ->willReturn( + $this->serverStream->reveal() + ); + $rowKeyStream = $this->dataClient->sampleRowKeys(); + $rowKeys = iterator_to_array($rowKeyStream); + $expectedRowKeys = [ + [ + 'rowKey' => 'rk1', + 'offset' => 1 + ], + [ + 'rowKey' => 'rk2', + 'offset' => 2 + ] + ]; + $this->assertEquals($expectedRowKeys, $rowKeys); + } + private function getMutateRowsResponse(array $status) { $mutateRowsResponses = [];