1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h"
17 #include "google/cloud/bigtable/internal/table.h"
18 #include "tensorflow/core/platform/test.h"
19
20 namespace tensorflow {
21 namespace {
22
WriteCell(const string & row,const string & family,const string & column,const string & value,::google::cloud::bigtable::noex::Table * table)23 void WriteCell(const string& row, const string& family, const string& column,
24 const string& value,
25 ::google::cloud::bigtable::noex::Table* table) {
26 ::google::cloud::bigtable::SingleRowMutation mut(row);
27 mut.emplace_back(::google::cloud::bigtable::SetCell(family, column, value));
28 table->Apply(std::move(mut));
29 }
30
TEST(BigtableTestClientTest,EmptyRowRead)31 TEST(BigtableTestClientTest, EmptyRowRead) {
32 std::shared_ptr<::google::cloud::bigtable::DataClient> client_ptr =
33 std::make_shared<BigtableTestClient>();
34 ::google::cloud::bigtable::noex::Table table(client_ptr, "test_table");
35
36 ::google::cloud::bigtable::RowSet rowset;
37 rowset.Append("r1");
38 auto filter = ::google::cloud::bigtable::Filter::Chain(
39 ::google::cloud::bigtable::Filter::Latest(1));
40 auto rows = table.ReadRows(std::move(rowset), filter);
41 EXPECT_EQ(rows.begin(), rows.end()) << "Some rows were returned in response!";
42 }
43
TEST(BigtableTestClientTest,SingleRowWriteAndRead)44 TEST(BigtableTestClientTest, SingleRowWriteAndRead) {
45 std::shared_ptr<::google::cloud::bigtable::DataClient> client_ptr =
46 std::make_shared<BigtableTestClient>();
47 ::google::cloud::bigtable::noex::Table table(client_ptr, "test_table");
48
49 WriteCell("r1", "f1", "c1", "v1", &table);
50
51 ::google::cloud::bigtable::RowSet rowset("r1");
52 auto filter = ::google::cloud::bigtable::Filter::Chain(
53 ::google::cloud::bigtable::Filter::Latest(1));
54 auto rows = table.ReadRows(std::move(rowset), filter);
55 auto itr = rows.begin();
56 EXPECT_NE(itr, rows.end()) << "No rows were returned in response!";
57 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
58 EXPECT_EQ((*itr)->row_key(), "r1");
59 EXPECT_EQ((*itr)->cells().size(), 1);
60 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
61 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
62 EXPECT_EQ((*itr)->cells()[0].value(), "v1");
63
64 ++itr;
65 EXPECT_EQ(itr, rows.end());
66 }
67
TEST(BigtableTestClientTest,MultiRowWriteAndSingleRowRead)68 TEST(BigtableTestClientTest, MultiRowWriteAndSingleRowRead) {
69 std::shared_ptr<::google::cloud::bigtable::DataClient> client_ptr =
70 std::make_shared<BigtableTestClient>();
71 ::google::cloud::bigtable::noex::Table table(client_ptr, "test_table");
72
73 WriteCell("r1", "f1", "c1", "v1", &table);
74 WriteCell("r2", "f1", "c1", "v2", &table);
75 WriteCell("r3", "f1", "c1", "v3", &table);
76
77 ::google::cloud::bigtable::RowSet rowset("r1");
78 auto filter = ::google::cloud::bigtable::Filter::Chain(
79 ::google::cloud::bigtable::Filter::Latest(1));
80 auto rows = table.ReadRows(std::move(rowset), filter);
81 auto itr = rows.begin();
82
83 EXPECT_NE(itr, rows.end()) << "Missing rows";
84 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
85 EXPECT_EQ((*itr)->row_key(), "r1");
86 EXPECT_EQ((*itr)->cells().size(), 1);
87 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
88 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
89 EXPECT_EQ((*itr)->cells()[0].value(), "v1");
90
91 ++itr;
92 EXPECT_EQ(itr, rows.end()) << "Extra rows in the response.";
93 }
94
TEST(BigtableTestClientTest,MultiRowWriteAndRead)95 TEST(BigtableTestClientTest, MultiRowWriteAndRead) {
96 std::shared_ptr<::google::cloud::bigtable::DataClient> client_ptr =
97 std::make_shared<BigtableTestClient>();
98 ::google::cloud::bigtable::noex::Table table(client_ptr, "test_table");
99
100 WriteCell("r1", "f1", "c1", "v1", &table);
101 WriteCell("r2", "f1", "c1", "v2", &table);
102 WriteCell("r3", "f1", "c1", "v3", &table);
103
104 ::google::cloud::bigtable::RowSet rowset("r1", "r2", "r3");
105 auto filter = ::google::cloud::bigtable::Filter::Chain(
106 ::google::cloud::bigtable::Filter::Latest(1));
107 auto rows = table.ReadRows(std::move(rowset), filter);
108 auto itr = rows.begin();
109
110 EXPECT_NE(itr, rows.end()) << "Missing rows";
111 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
112 EXPECT_EQ((*itr)->row_key(), "r1");
113 EXPECT_EQ((*itr)->cells().size(), 1);
114 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
115 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
116 EXPECT_EQ((*itr)->cells()[0].value(), "v1");
117
118 ++itr;
119
120 EXPECT_NE(itr, rows.end()) << "Missing rows";
121 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
122 EXPECT_EQ((*itr)->row_key(), "r2");
123 EXPECT_EQ((*itr)->cells().size(), 1);
124 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
125 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
126 EXPECT_EQ((*itr)->cells()[0].value(), "v2");
127
128 ++itr;
129
130 EXPECT_NE(itr, rows.end()) << "Missing rows";
131 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
132 EXPECT_EQ((*itr)->row_key(), "r3");
133 EXPECT_EQ((*itr)->cells().size(), 1);
134 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
135 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
136 EXPECT_EQ((*itr)->cells()[0].value(), "v3");
137
138 ++itr;
139 EXPECT_EQ(itr, rows.end()) << "Extra rows in the response.";
140 }
141
TEST(BigtableTestClientTest,MultiRowWriteAndPrefixRead)142 TEST(BigtableTestClientTest, MultiRowWriteAndPrefixRead) {
143 std::shared_ptr<::google::cloud::bigtable::DataClient> client_ptr =
144 std::make_shared<BigtableTestClient>();
145 ::google::cloud::bigtable::noex::Table table(client_ptr, "test_table");
146
147 WriteCell("r1", "f1", "c1", "v1", &table);
148 WriteCell("r2", "f1", "c1", "v2", &table);
149 WriteCell("r3", "f1", "c1", "v3", &table);
150
151 auto filter = ::google::cloud::bigtable::Filter::Chain(
152 ::google::cloud::bigtable::Filter::Latest(1));
153 auto rows =
154 table.ReadRows(::google::cloud::bigtable::RowRange::Prefix("r"), filter);
155 auto itr = rows.begin();
156
157 EXPECT_NE(itr, rows.end()) << "Missing rows";
158 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
159 EXPECT_EQ((*itr)->row_key(), "r1");
160 EXPECT_EQ((*itr)->cells().size(), 1);
161 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
162 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
163 EXPECT_EQ((*itr)->cells()[0].value(), "v1");
164
165 ++itr;
166
167 EXPECT_NE(itr, rows.end()) << "Missing rows";
168 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
169 EXPECT_EQ((*itr)->row_key(), "r2");
170 EXPECT_EQ((*itr)->cells().size(), 1);
171 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
172 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
173 EXPECT_EQ((*itr)->cells()[0].value(), "v2");
174
175 ++itr;
176
177 EXPECT_NE(itr, rows.end()) << "Missing rows";
178 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
179 EXPECT_EQ((*itr)->row_key(), "r3");
180 EXPECT_EQ((*itr)->cells().size(), 1);
181 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
182 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
183 EXPECT_EQ((*itr)->cells()[0].value(), "v3");
184
185 ++itr;
186 EXPECT_EQ(itr, rows.end()) << "Extra rows in the response.";
187 }
188
TEST(BigtableTestClientTest,ColumnFiltering)189 TEST(BigtableTestClientTest, ColumnFiltering) {
190 std::shared_ptr<::google::cloud::bigtable::DataClient> client_ptr =
191 std::make_shared<BigtableTestClient>();
192 ::google::cloud::bigtable::noex::Table table(client_ptr, "test_table");
193
194 WriteCell("r1", "f1", "c1", "v1", &table);
195 WriteCell("r2", "f1", "c1", "v2", &table);
196 WriteCell("r3", "f1", "c1", "v3", &table);
197
198 // Extra cells
199 WriteCell("r1", "f2", "c1", "v1", &table);
200 WriteCell("r2", "f2", "c1", "v2", &table);
201 WriteCell("r3", "f1", "c2", "v3", &table);
202
203 auto filter = ::google::cloud::bigtable::Filter::Chain(
204 ::google::cloud::bigtable::Filter::Latest(1),
205 ::google::cloud::bigtable::Filter::FamilyRegex("f1"),
206 ::google::cloud::bigtable::Filter::ColumnRegex("c1"));
207 auto rows =
208 table.ReadRows(::google::cloud::bigtable::RowRange::Prefix("r"), filter);
209 auto itr = rows.begin();
210
211 EXPECT_NE(itr, rows.end()) << "Missing rows";
212 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
213 EXPECT_EQ((*itr)->row_key(), "r1");
214 EXPECT_EQ((*itr)->cells().size(), 1);
215 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
216 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
217 EXPECT_EQ((*itr)->cells()[0].value(), "v1");
218
219 ++itr;
220
221 EXPECT_NE(itr, rows.end()) << "Missing rows";
222 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
223 EXPECT_EQ((*itr)->row_key(), "r2");
224 EXPECT_EQ((*itr)->cells().size(), 1);
225 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
226 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
227 EXPECT_EQ((*itr)->cells()[0].value(), "v2");
228
229 ++itr;
230
231 EXPECT_NE(itr, rows.end()) << "Missing rows";
232 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
233 EXPECT_EQ((*itr)->row_key(), "r3");
234 EXPECT_EQ((*itr)->cells().size(), 1);
235 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
236 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
237 EXPECT_EQ((*itr)->cells()[0].value(), "v3");
238
239 ++itr;
240 EXPECT_EQ(itr, rows.end()) << "Extra rows in the response.";
241 }
242
TEST(BigtableTestClientTest,RowKeys)243 TEST(BigtableTestClientTest, RowKeys) {
244 std::shared_ptr<::google::cloud::bigtable::DataClient> client_ptr =
245 std::make_shared<BigtableTestClient>();
246 ::google::cloud::bigtable::noex::Table table(client_ptr, "test_table");
247
248 WriteCell("r1", "f1", "c1", "v1", &table);
249 WriteCell("r2", "f1", "c1", "v2", &table);
250 WriteCell("r3", "f1", "c1", "v3", &table);
251
252 // Extra cells
253 WriteCell("r1", "f2", "c1", "v1", &table);
254 WriteCell("r2", "f2", "c1", "v2", &table);
255 WriteCell("r3", "f1", "c2", "v3", &table);
256
257 auto filter = ::google::cloud::bigtable::Filter::Chain(
258 ::google::cloud::bigtable::Filter::Latest(1),
259 ::google::cloud::bigtable::Filter::CellsRowLimit(1),
260 ::google::cloud::bigtable::Filter::StripValueTransformer());
261 auto rows =
262 table.ReadRows(::google::cloud::bigtable::RowRange::Prefix("r"), filter);
263 auto itr = rows.begin();
264 EXPECT_NE(itr, rows.end()) << "Missing rows";
265 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
266 EXPECT_EQ((*itr)->row_key(), "r1");
267 EXPECT_EQ((*itr)->cells().size(), 1);
268 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
269 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
270 EXPECT_EQ((*itr)->cells()[0].value(), "");
271
272 ++itr;
273
274 EXPECT_NE(itr, rows.end()) << "Missing rows";
275 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
276 EXPECT_EQ((*itr)->row_key(), "r2");
277 EXPECT_EQ((*itr)->cells().size(), 1);
278 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
279 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
280 EXPECT_EQ((*itr)->cells()[0].value(), "");
281
282 ++itr;
283
284 EXPECT_NE(itr, rows.end()) << "Missing rows";
285 EXPECT_TRUE(*itr) << "Error reading row: " << itr->status().message();
286 EXPECT_EQ((*itr)->row_key(), "r3");
287 EXPECT_EQ((*itr)->cells().size(), 1);
288 EXPECT_EQ((*itr)->cells()[0].family_name(), "f1");
289 EXPECT_EQ((*itr)->cells()[0].column_qualifier(), "c1");
290 EXPECT_EQ((*itr)->cells()[0].value(), "");
291
292 ++itr;
293 EXPECT_EQ(itr, rows.end()) << "Extra rows in the response.";
294 }
295
TEST(BigtableTestClientTest,SampleKeys)296 TEST(BigtableTestClientTest, SampleKeys) {
297 std::shared_ptr<::google::cloud::bigtable::DataClient> client_ptr =
298 std::make_shared<BigtableTestClient>();
299 ::google::cloud::bigtable::noex::Table table(client_ptr, "test_table");
300
301 WriteCell("r1", "f1", "c1", "v1", &table);
302 WriteCell("r2", "f1", "c1", "v2", &table);
303 WriteCell("r3", "f1", "c1", "v3", &table);
304 WriteCell("r4", "f1", "c1", "v4", &table);
305 WriteCell("r5", "f1", "c1", "v5", &table);
306
307 grpc::Status status;
308 auto resp = table.SampleRows(status);
309 EXPECT_TRUE(status.ok());
310 EXPECT_EQ(3, resp.size());
311 EXPECT_EQ("r1", string(resp[0].row_key));
312 EXPECT_EQ(0, resp[0].offset_bytes);
313 EXPECT_EQ("r3", string(resp[1].row_key));
314 EXPECT_EQ(100, resp[1].offset_bytes);
315 EXPECT_EQ("r5", string(resp[2].row_key));
316 EXPECT_EQ(200, resp[2].offset_bytes);
317 }
318
TEST(BigtableTestClientTest,SampleKeysShort)319 TEST(BigtableTestClientTest, SampleKeysShort) {
320 std::shared_ptr<::google::cloud::bigtable::DataClient> client_ptr =
321 std::make_shared<BigtableTestClient>();
322 ::google::cloud::bigtable::noex::Table table(client_ptr, "test_table");
323
324 WriteCell("r1", "f1", "c1", "v1", &table);
325
326 grpc::Status status;
327 auto resp = table.SampleRows(status);
328 EXPECT_TRUE(status.ok());
329 EXPECT_EQ(1, resp.size());
330 EXPECT_EQ("r1", string(resp[0].row_key));
331 }
332
TEST(BigtableTestClientTest,SampleKeysEvenNumber)333 TEST(BigtableTestClientTest, SampleKeysEvenNumber) {
334 std::shared_ptr<::google::cloud::bigtable::DataClient> client_ptr =
335 std::make_shared<BigtableTestClient>();
336 ::google::cloud::bigtable::noex::Table table(client_ptr, "test_table");
337
338 WriteCell("r1", "f1", "c1", "v1", &table);
339 WriteCell("r2", "f1", "c1", "v2", &table);
340 WriteCell("r3", "f1", "c1", "v3", &table);
341 WriteCell("r4", "f1", "c1", "v4", &table);
342
343 grpc::Status status;
344 auto resp = table.SampleRows(status);
345 EXPECT_TRUE(status.ok());
346 EXPECT_EQ(2, resp.size());
347 EXPECT_EQ("r1", string(resp[0].row_key));
348 EXPECT_EQ("r3", string(resp[1].row_key));
349 }
350
351 } // namespace
352 } // namespace tensorflow
353