Skip to content

Commit 5992f28

Browse files
author
Ben Hu
committed
Migrate bigquery storage SDK from v1beta1 to v1
1 parent 010da6a commit 5992f28

7 files changed

Lines changed: 91 additions & 90 deletions

File tree

‎tensorflow_io/core/BUILD‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -169,7 +169,7 @@ cc_library(
169169
"@avro",
170170
"@com_github_grpc_grpc//:grpc++",
171171
"@com_google_absl//absl/types:any",
172-
"@com_google_googleapis//google/cloud/bigquery/storage/v1beta1:storage_cc_grpc",
172+
"@com_google_googleapis//google/cloud/bigquery/storage/v1:storage_cc_grpc",
173173
"@local_config_tf//:libtensorflow_framework",
174174
"@local_config_tf//:tf_header_lib",
175175
],
@@ -219,7 +219,7 @@ cc_library(
219219
"@com_google_absl//absl/container:flat_hash_map",
220220
"@com_google_absl//absl/types:any",
221221
"@com_google_absl//absl/types:variant",
222-
"@com_google_googleapis//google/cloud/bigquery/storage/v1beta1:storage_cc_grpc",
222+
"@com_google_googleapis//google/cloud/bigquery/storage/v1:storage_cc_grpc",
223223
"@local_config_tf//:libtensorflow_framework",
224224
"@local_config_tf//:tf_header_lib",
225225
],

‎tensorflow_io/core/kernels/bigquery/bigquery_dataset_op.cc‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ class BigQueryDatasetOp : public DatasetOpKernel {
108108
std::vector<string> default_values_;
109109
std::vector<absl::any> typed_default_values_;
110110
int64 offset_;
111-
apiv1beta1::DataFormat data_format_;
111+
apiv1::DataFormat data_format_;
112112

113113
class Dataset : public DatasetBase {
114114
public:
@@ -120,7 +120,7 @@ class BigQueryDatasetOp : public DatasetOpKernel {
120120
std::vector<string> selected_fields,
121121
std::vector<DataType> output_types,
122122
std::vector<absl::any> typed_default_values, int64 offset_,
123-
apiv1beta1::DataFormat data_format)
123+
apiv1::DataFormat data_format)
124124
: DatasetBase(DatasetContext(ctx)),
125125
client_resource_(client_resource),
126126
output_types_vector_(output_types_vector),
@@ -134,10 +134,10 @@ class BigQueryDatasetOp : public DatasetOpKernel {
134134
data_format_(data_format) {
135135
client_resource_->Ref();
136136

137-
if (data_format == apiv1beta1::DataFormat::AVRO) {
137+
if (data_format == apiv1::DataFormat::AVRO) {
138138
std::istringstream istream(schema);
139139
avro::compileJsonSchema(istream, *avro_schema_);
140-
} else if (data_format == apiv1beta1::DataFormat::ARROW) {
140+
} else if (data_format == apiv1::DataFormat::ARROW) {
141141
auto buffer_ = std::make_shared<arrow::Buffer>(
142142
reinterpret_cast<const uint8_t *>(&schema[0]), schema.length());
143143

@@ -158,11 +158,11 @@ class BigQueryDatasetOp : public DatasetOpKernel {
158158

159159
std::unique_ptr<IteratorBase> MakeIteratorInternal(
160160
const string &prefix) const override {
161-
if (data_format_ == apiv1beta1::DataFormat::AVRO) {
161+
if (data_format_ == apiv1::DataFormat::AVRO) {
162162
return std::unique_ptr<IteratorBase>(
163163
new BigQueryReaderAvroDatasetIterator<Dataset>(
164164
{this, strings::StrCat(prefix, "::BigQueryAvroDataset")}));
165-
} else if (data_format_ == apiv1beta1::DataFormat::ARROW) {
165+
} else if (data_format_ == apiv1::DataFormat::ARROW) {
166166
return std::unique_ptr<IteratorBase>(
167167
new BigQueryReaderArrowDatasetIterator<Dataset>(
168168
{this, strings::StrCat(prefix, "::BigQueryArrowDataset")}));
@@ -229,7 +229,7 @@ class BigQueryDatasetOp : public DatasetOpKernel {
229229
const std::unique_ptr<avro::ValidSchema> avro_schema_;
230230
const int64 offset_;
231231
std::shared_ptr<::arrow::Schema> arrow_schema_;
232-
const apiv1beta1::DataFormat data_format_;
232+
const apiv1::DataFormat data_format_;
233233
};
234234
};
235235

‎tensorflow_io/core/kernels/bigquery/bigquery_kernels.cc‎

Lines changed: 21 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ limitations under the License.
1919
namespace tensorflow {
2020
namespace {
2121

22-
namespace apiv1beta1 = ::google::cloud::bigquery::storage::v1beta1;
22+
namespace apiv1 = ::google::cloud::bigquery::storage::v1;
2323

2424
class BigQueryClientOp : public OpKernel {
2525
public:
@@ -105,35 +105,30 @@ class BigQueryReadSessionOp : public OpKernel {
105105
ctx, LookupResource(ctx, HandleFromInput(ctx, 0), &client_resource));
106106
core::ScopedUnref scoped_unref(client_resource);
107107

108-
apiv1beta1::CreateReadSessionRequest createReadSessionRequest;
109-
createReadSessionRequest.mutable_table_reference()->set_project_id(
110-
project_id_);
111-
createReadSessionRequest.mutable_table_reference()->set_dataset_id(
112-
dataset_id_);
113-
createReadSessionRequest.mutable_table_reference()->set_table_id(table_id_);
108+
apiv1::CreateReadSessionRequest createReadSessionRequest;
114109
createReadSessionRequest.set_parent(parent_);
115-
*createReadSessionRequest.mutable_read_options()
116-
->mutable_selected_fields() = {selected_fields_.begin(),
117-
selected_fields_.end()};
118-
createReadSessionRequest.mutable_read_options()->set_row_restriction(
119-
row_restriction_);
120-
createReadSessionRequest.set_requested_streams(requested_streams_);
121-
createReadSessionRequest.set_sharding_strategy(
122-
apiv1beta1::ShardingStrategy::BALANCED);
123-
createReadSessionRequest.set_format(data_format_);
110+
apiv1::ReadSession* read_session =
111+
createReadSessionRequest.mutable_read_session();
112+
read_session->set_table(strings::Printf(
113+
"projects/%s/datasets/%s/tables/%s", project_id_.c_str(),
114+
dataset_id_.c_str(), table_id_.c_str()));
115+
read_session->set_data_format(data_format_);
116+
*read_session->mutable_read_options()->mutable_selected_fields() = {
117+
selected_fields_.begin(), selected_fields_.end()};
118+
read_session->mutable_read_options()->set_row_restriction(row_restriction_);
119+
createReadSessionRequest.set_max_stream_count(requested_streams_);
120+
124121
VLOG(3) << "createReadSessionRequest: "
125122
<< createReadSessionRequest.DebugString();
126123
::grpc::ClientContext context;
127-
context.AddMetadata(
128-
"x-goog-request-params",
129-
strings::Printf("table_reference.dataset_id=%s&table_"
130-
"reference.project_id=%s",
131-
dataset_id_.c_str(), project_id_.c_str()));
124+
context.AddMetadata("x-goog-request-params",
125+
strings::Printf("read_session.table=%s",
126+
read_session->table().c_str()));
132127
context.set_deadline(gpr_time_add(gpr_now(GPR_CLOCK_REALTIME),
133128
gpr_time_from_seconds(60, GPR_TIMESPAN)));
134129

135-
std::shared_ptr<apiv1beta1::ReadSession> readSessionResponse =
136-
std::make_shared<apiv1beta1::ReadSession>();
130+
std::shared_ptr<apiv1::ReadSession> readSessionResponse =
131+
std::make_shared<apiv1::ReadSession>();
137132
VLOG(3) << "calling readSession";
138133
::grpc::Status status = client_resource->GetStub("")->CreateReadSession(
139134
&context, createReadSessionRequest, readSessionResponse.get());
@@ -155,13 +150,13 @@ class BigQueryReadSessionOp : public OpKernel {
155150
Tensor* schema_t = nullptr;
156151
OP_REQUIRES_OK(ctx, ctx->allocate_output("schema", {}, &schema_t));
157152

158-
if (data_format_ == apiv1beta1::DataFormat::AVRO) {
153+
if (data_format_ == apiv1::DataFormat::AVRO) {
159154
OP_REQUIRES(ctx, readSessionResponse->has_avro_schema(),
160155
errors::InvalidArgument("AVRO schema is missing"));
161156
VLOG(3) << "avro schema:" << readSessionResponse->avro_schema().schema();
162157
schema_t->scalar<tstring>()() =
163158
readSessionResponse->avro_schema().schema();
164-
} else if (data_format_ == apiv1beta1::DataFormat::ARROW) {
159+
} else if (data_format_ == apiv1::DataFormat::ARROW) {
165160
OP_REQUIRES(ctx, readSessionResponse->has_arrow_schema(),
166161
errors::InvalidArgument("ARROW schema is missing"));
167162
VLOG(3) << "arrow schema:"
@@ -183,7 +178,7 @@ class BigQueryReadSessionOp : public OpKernel {
183178
std::vector<DataType> output_types_;
184179
string row_restriction_;
185180
int requested_streams_;
186-
apiv1beta1::DataFormat data_format_;
181+
apiv1::DataFormat data_format_;
187182

188183
mutex mu_;
189184
ContainerInfo cinfo_ TF_GUARDED_BY(mu_);

‎tensorflow_io/core/kernels/bigquery/bigquery_lib.cc‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,11 +78,11 @@ string GrpcStatusToString(const ::grpc::Status& status) {
7878
}
7979

8080
Status GetDataFormat(string data_format_str,
81-
apiv1beta1::DataFormat* data_format) {
81+
apiv1::DataFormat* data_format) {
8282
if (data_format_str == "ARROW") {
83-
*data_format = apiv1beta1::DataFormat::ARROW;
83+
*data_format = apiv1::DataFormat::ARROW;
8484
} else if (data_format_str == "AVRO") {
85-
*data_format = apiv1beta1::DataFormat::AVRO;
85+
*data_format = apiv1::DataFormat::AVRO;
8686
} else {
8787
return errors::Internal("Unsupported data format: " + data_format_str);
8888
}

‎tensorflow_io/core/kernels/bigquery/bigquery_lib.h‎

Lines changed: 21 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ limitations under the License.
3838
#include "arrow/buffer.h"
3939
#include "arrow/io/memory.h"
4040
#include "arrow/ipc/api.h"
41-
#include "google/cloud/bigquery/storage/v1beta1/storage.grpc.pb.h"
41+
#include "google/cloud/bigquery/storage/v1/storage.grpc.pb.h"
4242
#include "tensorflow/core/framework/dataset.h"
4343
#include "tensorflow/core/framework/resource_mgr.h"
4444
#include "tensorflow/core/framework/tensor.h"
@@ -50,18 +50,18 @@ limitations under the License.
5050

5151
namespace tensorflow {
5252

53-
namespace apiv1beta1 = ::google::cloud::bigquery::storage::v1beta1;
53+
namespace apiv1 = ::google::cloud::bigquery::storage::v1;
5454
static constexpr int kMaxReceiveMessageSize = -1; // Disabled
5555

5656
Status GrpcStatusToTfStatus(const ::grpc::Status &status);
5757
string GrpcStatusToString(const ::grpc::Status &status);
5858
Status GetDataFormat(string data_format_str,
59-
apiv1beta1::DataFormat *data_format);
59+
apiv1::DataFormat *data_format);
6060

6161
class BigQueryClientResource : public ResourceBase {
6262
public:
6363
explicit BigQueryClientResource(
64-
std::function<std::unique_ptr<apiv1beta1::BigQueryStorage::Stub>(
64+
std::function<std::unique_ptr<apiv1::BigQueryRead::Stub>(
6565
const string &read_stream)>
6666
stub_factory)
6767
: stub_factory_(stub_factory) {}
@@ -80,10 +80,10 @@ class BigQueryClientResource : public ResourceBase {
8080
args.SetString("read_stream", read_stream);
8181
auto channel = ::grpc::CreateCustomChannel(server_name, creds, args);
8282
VLOG(3) << "Creating GRPC channel";
83-
return absl::make_unique<apiv1beta1::BigQueryStorage::Stub>(channel);
83+
return absl::make_unique<apiv1::BigQueryRead::Stub>(channel);
8484
}) {}
8585

86-
apiv1beta1::BigQueryStorage::Stub *GetStub(const string &read_stream)
86+
apiv1::BigQueryRead::Stub *GetStub(const string &read_stream)
8787
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
8888
if (stubs_.find(read_stream) == stubs_.end()) {
8989
auto stub = stub_factory_(read_stream);
@@ -95,11 +95,11 @@ class BigQueryClientResource : public ResourceBase {
9595
string DebugString() const override { return "BigQueryClientResource"; }
9696

9797
private:
98-
std::function<std::unique_ptr<apiv1beta1::BigQueryStorage::Stub>(
98+
std::function<std::unique_ptr<apiv1::BigQueryRead::Stub>(
9999
const string &)>
100100
stub_factory_;
101101
mutex mu_;
102-
std::unordered_map<string, std::unique_ptr<apiv1beta1::BigQueryStorage::Stub>>
102+
std::unordered_map<string, std::unique_ptr<apiv1::BigQueryRead::Stub>>
103103
stubs_ TF_GUARDED_BY(mu_);
104104
};
105105

@@ -156,11 +156,9 @@ class BigQueryReaderDatasetIteratorBase : public DatasetIterator<Dataset> {
156156
return OkStatus();
157157
}
158158

159-
apiv1beta1::ReadRowsRequest readRowsRequest;
160-
readRowsRequest.mutable_read_position()->mutable_stream()->set_name(
161-
this->dataset()->stream());
162-
readRowsRequest.mutable_read_position()->set_offset(
163-
this->dataset()->offset());
159+
apiv1::ReadRowsRequest readRowsRequest;
160+
readRowsRequest.set_read_stream(this->dataset()->stream());
161+
readRowsRequest.set_offset(this->dataset()->offset());
164162

165163
read_rows_context_ = absl::make_unique<::grpc::ClientContext>();
166164
// The deadline is for the entire ReadRows (not a single message receipt),
@@ -169,14 +167,14 @@ class BigQueryReaderDatasetIteratorBase : public DatasetIterator<Dataset> {
169167
std::chrono::hours(24));
170168
read_rows_context_->AddMetadata(
171169
"x-goog-request-params",
172-
absl::StrCat("read_position.stream.name=",
173-
readRowsRequest.read_position().stream().name()));
170+
absl::StrCat("read_stream=",
171+
readRowsRequest.read_stream()));
174172

175173
VLOG(3) << "getting reader, stream: "
176-
<< readRowsRequest.read_position().stream().DebugString();
174+
<< readRowsRequest.read_stream();
177175
reader_ = this->dataset()
178176
->client_resource()
179-
->GetStub(readRowsRequest.read_position().stream().name())
177+
->GetStub(readRowsRequest.read_stream())
180178
->ReadRows(read_rows_context_.get(), readRowsRequest);
181179

182180
return OkStatus();
@@ -191,9 +189,9 @@ class BigQueryReaderDatasetIteratorBase : public DatasetIterator<Dataset> {
191189
int current_row_index_ = 0;
192190
mutex mu_;
193191
std::unique_ptr<::grpc::ClientContext> read_rows_context_ TF_GUARDED_BY(mu_);
194-
std::unique_ptr<::grpc::ClientReader<apiv1beta1::ReadRowsResponse>> reader_
192+
std::unique_ptr<::grpc::ClientReader<apiv1::ReadRowsResponse>> reader_
195193
TF_GUARDED_BY(mu_);
196-
std::unique_ptr<apiv1beta1::ReadRowsResponse> response_ TF_GUARDED_BY(mu_);
194+
std::unique_ptr<apiv1::ReadRowsResponse> response_ TF_GUARDED_BY(mu_);
197195
};
198196

199197
// BigQuery reader for Arrow serialized data.
@@ -213,11 +211,11 @@ class BigQueryReaderArrowDatasetIterator
213211
TF_EXCLUSIVE_LOCKS_REQUIRED(this->mu_) override {
214212
if (this->response_ && this->response_->has_arrow_record_batch() &&
215213
this->current_row_index_ <
216-
this->response_->arrow_record_batch().row_count()) {
214+
this->response_->row_count()) {
217215
return OkStatus();
218216
}
219217

220-
this->response_ = absl::make_unique<apiv1beta1::ReadRowsResponse>();
218+
this->response_ = absl::make_unique<apiv1::ReadRowsResponse>();
221219
if (!this->reader_->Read(this->response_.get())) {
222220
*end_of_sequence = true;
223221
return GrpcStatusToTfStatus(this->reader_->Finish());
@@ -315,11 +313,11 @@ class BigQueryReaderAvroDatasetIterator
315313
Status EnsureHasRow(bool *end_of_sequence)
316314
TF_EXCLUSIVE_LOCKS_REQUIRED(this->mu_) override {
317315
if (this->response_ &&
318-
this->current_row_index_ < this->response_->avro_rows().row_count()) {
316+
this->current_row_index_ < this->response_->row_count()) {
319317
return OkStatus();
320318
}
321319

322-
this->response_ = absl::make_unique<apiv1beta1::ReadRowsResponse>();
320+
this->response_ = absl::make_unique<apiv1::ReadRowsResponse>();
323321
VLOG(3) << "calling read";
324322
if (!this->reader_->Read(this->response_.get())) {
325323
VLOG(3) << "no data";

‎tensorflow_io/core/kernels/tests/bigquery_test_client_op.cc‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ limitations under the License.
1818
namespace tensorflow {
1919
namespace {
2020

21-
namespace apiv1beta1 = ::google::cloud::bigquery::storage::v1beta1;
21+
namespace apiv1 = ::google::cloud::bigquery::storage::v1;
2222

2323
class BigQueryTestClientOp : public OpKernel {
2424
public:
@@ -58,13 +58,13 @@ class BigQueryTestClientOp : public OpKernel {
5858
std::shared_ptr<grpc::Channel> channel =
5959
::grpc::CreateChannel(this->fake_server_address_,
6060
grpc::InsecureChannelCredentials());
61-
auto stub = apiv1beta1::BigQueryStorage::NewStub(channel);
61+
auto stub = apiv1::BigQueryRead::NewStub(channel);
6262
LOG(INFO) << "BigQueryTestClientOp waiting for connections";
6363
channel->WaitForConnected(
6464
gpr_time_add(gpr_now(GPR_CLOCK_REALTIME),
6565
gpr_time_from_seconds(15, GPR_TIMESPAN)));
6666
LOG(INFO) << "Done creating BigQueryTestClientOp Fake client";
67-
return absl::make_unique<apiv1beta1::BigQueryStorage::Stub>(
67+
return absl::make_unique<apiv1::BigQueryRead::Stub>(
6868
channel);
6969
});
7070
return OkStatus();

0 commit comments

Comments
 (0)