Bigtable: Sample row keys api

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

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;

Bigtable/tests/System/DataClientSampleRowKeysTest.php

@@ -0,0 +1,88 @@
+<?php
+/**
+ * Copyright 2018 Google Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+namespace Google\Cloud\Bigtable\Tests\System;
+
+use Google\Cloud\Bigtable\Tests\System\DataClientTest;
+
+/**
+ * @group bigtable
+ * @group bigtabledata
+ */
+class DataClientSampleRowKeysTest extends DataClientTest
+{
+    public static function setUpBeforeClass()
+    {
+        parent::setUpBeforeClass();
+        $insertRows = [
+            'rk1' => [
+                '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);
+    }
+}

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 = [];