diff --git a/google/cloud/bigquery/v2/minimal/internal/job_request.cc b/google/cloud/bigquery/v2/minimal/internal/job_request.cc index d3404627870fb..5a6382dd510cf 100644 --- a/google/cloud/bigquery/v2/minimal/internal/job_request.cc +++ b/google/cloud/bigquery/v2/minimal/internal/job_request.cc @@ -281,6 +281,10 @@ void to_json(nlohmann::json& j, QueryRequest const& q) { {"formatOptions", q.format_options()}, {"labels", q.labels()}}; + if (!q.query_results_format().empty()) { + j["queryResultsFormat"] = q.query_results_format(); + } + ToIntJson(q.timeout(), j, "timeoutMs"); // timeoutMs value is a number for this request type. } @@ -309,6 +313,8 @@ void from_json(nlohmann::json const& j, QueryRequest& q) { SafeGetTo(j, "defaultDataset", &QueryRequest::set_default_dataset, q); SafeGetTo(j, "formatOptions", &QueryRequest::set_format_options, q); SafeGetTo(j, "labels", &QueryRequest::set_labels, q); + SafeGetTo(j, "queryResultsFormat", &QueryRequest::set_query_results_format, + q); std::chrono::milliseconds timeout; FromJson(timeout, j, "timeoutMs"); @@ -370,6 +376,7 @@ std::string QueryRequest::DebugString(absl::string_view name, .SubMessage("default_dataset", default_dataset()) .SubMessage("format_options", format_options()) .SubMessage("job_creation_mode", job_creation_mode()) + .StringField("query_results_format", query_results_format()) .Build(); } diff --git a/google/cloud/bigquery/v2/minimal/internal/job_request.h b/google/cloud/bigquery/v2/minimal/internal/job_request.h index 6dfd12633850c..ab832b818a959 100644 --- a/google/cloud/bigquery/v2/minimal/internal/job_request.h +++ b/google/cloud/bigquery/v2/minimal/internal/job_request.h @@ -538,6 +538,17 @@ class QueryRequest { return std::move(set_job_creation_mode(std::move(job_creation_mode))); } + std::string const& query_results_format() const { + return query_results_format_; + } + QueryRequest& set_query_results_format(std::string query_results_format) & { + query_results_format_ = std::move(query_results_format); + return *this; + } + QueryRequest&& set_query_results_format(std::string query_results_format) && { + return std::move(set_query_results_format(std::move(query_results_format))); + } + std::string DebugString(absl::string_view name, TracingOptions const& options = {}, int indent = 0) const; @@ -548,6 +559,7 @@ class QueryRequest { std::string parameter_mode_; std::string location_; std::string request_id_; + std::string query_results_format_; bool dry_run_ = false; bool preserve_nulls_ = false; diff --git a/google/cloud/bigquery/v2/minimal/internal/job_request_test.cc b/google/cloud/bigquery/v2/minimal/internal/job_request_test.cc index d1be2f64247a8..b15bd2ba1bd44 100644 --- a/google/cloud/bigquery/v2/minimal/internal/job_request_test.cc +++ b/google/cloud/bigquery/v2/minimal/internal/job_request_test.cc @@ -1029,7 +1029,8 @@ TEST(PostQueryRequestTest, DebugString) { R"( labels { key: "lk1" value: "lv1" } labels { key: "lk2" value: "lv2" })" R"( default_dataset { project_id: "2" dataset_id: "1" })" R"( format_options { use_int64_timestamp: true timestamp_output_format: TIMESTAMP_OUTPUT_FORMAT_UNSPECIFIED })" - R"( job_creation_mode { value: "JOB_CREATION_MODE_UNSPECIFIED" } } })"); + R"( job_creation_mode { value: "JOB_CREATION_MODE_UNSPECIFIED" })" + R"( query_results_format: "ARROW" } })"); EXPECT_EQ( request.DebugString( @@ -1050,7 +1051,8 @@ TEST(PostQueryRequestTest, DebugString) { R"( labels { key: "lk1" value: "lv1" } labels { key: "lk2" value: "lv2" })" R"( default_dataset { project_id: "2" dataset_id: "1" })" R"( format_options { use_int64_timestamp: true timestamp_output_format: TIMESTAMP_OUTPUT_FORMAT_UNSPECIFIED })" - R"( job_creation_mode { value: "JOB_CRE......" } } })"); + R"( job_creation_mode { value: "JOB_CRE......" })" + R"( query_results_format: "ARROW" } })"); EXPECT_EQ(request.DebugString("PostQueryRequest", TracingOptions{}.SetOptions( "single_line_mode=F")), @@ -1108,6 +1110,7 @@ TEST(PostQueryRequestTest, DebugString) { job_creation_mode { value: "JOB_CREATION_MODE_UNSPECIFIED" } + query_results_format: "ARROW" } })"); } diff --git a/google/cloud/bigquery/v2/minimal/internal/job_response.cc b/google/cloud/bigquery/v2/minimal/internal/job_response.cc index 1c00a4fa072d0..dbbab595d62ed 100644 --- a/google/cloud/bigquery/v2/minimal/internal/job_response.cc +++ b/google/cloud/bigquery/v2/minimal/internal/job_response.cc @@ -14,6 +14,7 @@ #include "google/cloud/bigquery/v2/minimal/internal/job_response.h" #include "google/cloud/bigquery/v2/minimal/internal/json_utils.h" +#include "google/cloud/internal/base64_transforms.h" #include "google/cloud/internal/debug_string.h" #include "google/cloud/internal/make_status.h" #include "absl/strings/str_cat.h" @@ -192,6 +193,60 @@ std::string CancelJobResponse::DebugString(absl::string_view name, .Build(); } +std::string ArrowSchema::DebugString(absl::string_view name, + TracingOptions const& options, + int indent) const { + return internal::DebugFormatter(name, options, indent) + .StringField("serialized_schema", serialized_schema) + .Build(); +} + +std::string ArrowRecordBatch::DebugString(absl::string_view name, + TracingOptions const& options, + int indent) const { + return internal::DebugFormatter(name, options, indent) + .StringField("serialized_record_batch", serialized_record_batch) + .Field("row_count", row_count) + .Build(); +} + +void to_json(nlohmann::json& j, ArrowSchema const& a) { + j = nlohmann::json{ + {"serializedSchema", internal::UrlsafeBase64Encode(a.serialized_schema)}}; +} + +void from_json(nlohmann::json const& j, ArrowSchema& a) { + if (j.contains("serializedSchema") && j["serializedSchema"].is_string()) { + std::string b64 = j["serializedSchema"].get(); + auto bytes = internal::UrlsafeBase64Decode(b64); + if (bytes.ok()) { + a.serialized_schema.assign(reinterpret_cast(bytes->data()), + bytes->size()); + } + } +} + +void to_json(nlohmann::json& j, ArrowRecordBatch const& a) { + j = nlohmann::json{{"serializedRecordBatch", + internal::UrlsafeBase64Encode(a.serialized_record_batch)}, + {"rowCount", std::to_string(a.row_count)}}; +} + +void from_json(nlohmann::json const& j, ArrowRecordBatch& a) { + if (j.contains("serializedRecordBatch") && + j["serializedRecordBatch"].is_string()) { + std::string b64 = j["serializedRecordBatch"].get(); + auto bytes = internal::UrlsafeBase64Decode(b64); + if (bytes.ok()) { + a.serialized_record_batch.assign( + reinterpret_cast(bytes->data()), bytes->size()); + } + } + if (j.contains("rowCount")) { + a.row_count = GetNumberFromJson(j, "rowCount"); + } +} + std::string PostQueryResults::DebugString(absl::string_view name, TracingOptions const& options, int indent) const { @@ -209,6 +264,9 @@ std::string PostQueryResults::DebugString(absl::string_view name, .SubMessage("job_reference", job_reference) .SubMessage("session_info", session_info) .SubMessage("dml_stats", dml_stats) + .SubMessage("arrow_schema", arrow_schema) + .SubMessage("arrow_record_batch", arrow_record_batch) + .Field("page_row_count", page_row_count) .Build(); } @@ -262,6 +320,11 @@ StatusOr QueryResponse::BuildFromHttpResponse( SafeGetTo(query_results.session_info, *json, "sessionInfo"); SafeGetTo(query_results.dml_stats, *json, "dmlStats"); + SafeGetTo(query_results.arrow_schema, *json, "arrowSchema"); + SafeGetTo(query_results.arrow_record_batch, *json, "arrowRecordBatch"); + if (json->contains("pageRowCount")) { + query_results.page_row_count = GetNumberFromJson(*json, "pageRowCount"); + } QueryResponse response; response.http_response = http_response; @@ -299,6 +362,15 @@ void to_json(nlohmann::json& j, PostQueryResults const& q) { {"errors", q.errors}, {"sessionInfo", q.session_info}, {"dmlStats", q.dml_stats}}; + if (!q.arrow_schema.serialized_schema.empty()) { + j["arrowSchema"] = q.arrow_schema; + } + if (!q.arrow_record_batch.serialized_record_batch.empty()) { + j["arrowRecordBatch"] = q.arrow_record_batch; + } + if (q.page_row_count > 0) { + j["pageRowCount"] = std::to_string(q.page_row_count); + } } void from_json(nlohmann::json const& j, PostQueryResults& q) { @@ -315,6 +387,11 @@ void from_json(nlohmann::json const& j, PostQueryResults& q) { SafeGetTo(q.errors, j, "errors"); SafeGetTo(q.session_info, j, "sessionInfo"); SafeGetTo(q.dml_stats, j, "dmlStats"); + SafeGetTo(q.arrow_schema, j, "arrowSchema"); + SafeGetTo(q.arrow_record_batch, j, "arrowRecordBatch"); + if (j.contains("pageRowCount")) { + q.page_row_count = GetNumberFromJson(j, "pageRowCount"); + } } void to_json(nlohmann::json& j, GetQueryResults const& q) { @@ -330,6 +407,15 @@ void to_json(nlohmann::json& j, GetQueryResults const& q) { {"jobReference", q.job_reference}, {"rows", q.rows}, {"errors", q.errors}}; + if (!q.arrow_schema.serialized_schema.empty()) { + j["arrowSchema"] = q.arrow_schema; + } + if (!q.arrow_record_batch.serialized_record_batch.empty()) { + j["arrowRecordBatch"] = q.arrow_record_batch; + } + if (q.page_row_count > 0) { + j["pageRowCount"] = std::to_string(q.page_row_count); + } } void from_json(nlohmann::json const& j, GetQueryResults& q) { SafeGetTo(q.kind, j, "kind"); @@ -344,6 +430,11 @@ void from_json(nlohmann::json const& j, GetQueryResults& q) { SafeGetTo(q.job_reference, j, "jobReference"); SafeGetTo(q.rows, j, "rows"); SafeGetTo(q.errors, j, "errors"); + SafeGetTo(q.arrow_schema, j, "arrowSchema"); + SafeGetTo(q.arrow_record_batch, j, "arrowRecordBatch"); + if (j.contains("pageRowCount")) { + q.page_row_count = GetNumberFromJson(j, "pageRowCount"); + } } std::string GetQueryResults::DebugString(absl::string_view name, @@ -362,6 +453,9 @@ std::string GetQueryResults::DebugString(absl::string_view name, .Field("errors", errors) .SubMessage("schema", schema) .SubMessage("job_reference", job_reference) + .SubMessage("arrow_schema", arrow_schema) + .SubMessage("arrow_record_batch", arrow_record_batch) + .Field("page_row_count", page_row_count) .Build(); } @@ -413,6 +507,12 @@ GetQueryResultsResponse::BuildFromHttpResponse( } } + SafeGetTo(get_query_results.arrow_schema, *json, "arrowSchema"); + SafeGetTo(get_query_results.arrow_record_batch, *json, "arrowRecordBatch"); + if (json->contains("pageRowCount")) { + get_query_results.page_row_count = GetNumberFromJson(*json, "pageRowCount"); + } + GetQueryResultsResponse response; response.http_response = http_response; response.get_query_results = get_query_results; diff --git a/google/cloud/bigquery/v2/minimal/internal/job_response.h b/google/cloud/bigquery/v2/minimal/internal/job_response.h index 4cfefd634bdaf..7c19e9328a0f0 100644 --- a/google/cloud/bigquery/v2/minimal/internal/job_response.h +++ b/google/cloud/bigquery/v2/minimal/internal/job_response.h @@ -96,6 +96,35 @@ class CancelJobResponse { BigQueryHttpResponse http_response; }; +struct ArrowSchema { + std::string serialized_schema; + + std::string DebugString(absl::string_view name, + TracingOptions const& options = {}, + int indent = 0) const; +}; +void to_json(nlohmann::json& j, ArrowSchema const& a); +void from_json(nlohmann::json const& j, ArrowSchema& a); +inline bool operator==(ArrowSchema const& lhs, ArrowSchema const& rhs) { + return lhs.serialized_schema == rhs.serialized_schema; +} + +struct ArrowRecordBatch { + std::string serialized_record_batch; + std::int64_t row_count = 0; + + std::string DebugString(absl::string_view name, + TracingOptions const& options = {}, + int indent = 0) const; +}; +void to_json(nlohmann::json& j, ArrowRecordBatch const& a); +void from_json(nlohmann::json const& j, ArrowRecordBatch& a); +inline bool operator==(ArrowRecordBatch const& lhs, + ArrowRecordBatch const& rhs) { + return lhs.serialized_record_batch == rhs.serialized_record_batch && + lhs.row_count == rhs.row_count; +} + struct PostQueryResults { std::string DebugString(absl::string_view name, TracingOptions const& options = {}, @@ -117,6 +146,9 @@ struct PostQueryResults { std::vector errors; SessionInfo session_info; DmlStats dml_stats; + ArrowSchema arrow_schema; + ArrowRecordBatch arrow_record_batch; + std::int64_t page_row_count = 0; }; void to_json(nlohmann::json& j, PostQueryResults const& q); void from_json(nlohmann::json const& j, PostQueryResults& q); @@ -153,6 +185,9 @@ struct GetQueryResults { std::vector rows; std::vector errors; + ArrowSchema arrow_schema; + ArrowRecordBatch arrow_record_batch; + std::int64_t page_row_count = 0; std::string DebugString(absl::string_view name, TracingOptions const& options = {}, diff --git a/google/cloud/bigquery/v2/minimal/internal/job_response_test.cc b/google/cloud/bigquery/v2/minimal/internal/job_response_test.cc index 4e56ee42c4365..6c0d3ee8b37ac 100644 --- a/google/cloud/bigquery/v2/minimal/internal/job_response_test.cc +++ b/google/cloud/bigquery/v2/minimal/internal/job_response_test.cc @@ -2009,7 +2009,10 @@ TEST(QueryResponseTest, DebugString) { R"( rounding_mode { value: "" } range_element_type { type: "" } } })" R"( job_reference { project_id: "p123" job_id: "j123" location: "useast" })" R"( session_info { session_id: "123" } dml_stats { inserted_row_count: 10)" - R"( deleted_row_count: 10 updated_row_count: 10 } } })"); + R"( deleted_row_count: 10 updated_row_count: 10 })" + R"( arrow_schema { serialized_schema: "testing_schema_data" })" + R"( arrow_record_batch { serialized_record_batch: "testing_batch_data" row_count: 10 })" + R"( page_row_count: 10 } })"); EXPECT_EQ( response->DebugString( @@ -2029,7 +2032,10 @@ TEST(QueryResponseTest, DebugString) { R"( rounding_mode { value: "" } range_element_type { type: "" } } })" R"( job_reference { project_id: "p123" job_id: "j123" location: "useast" })" R"( session_info { session_id: "123" } dml_stats {)" - R"( inserted_row_count: 10 deleted_row_count: 10 updated_row_count: 10 } } })"); + R"( inserted_row_count: 10 deleted_row_count: 10 updated_row_count: 10 })" + R"( arrow_schema { serialized_schema: "testing......" })" + R"( arrow_record_batch { serialized_record_batch: "testing......" row_count: 10 })" + R"( page_row_count: 10 } })"); EXPECT_EQ(response->DebugString("QueryResponse", TracingOptions{}.SetOptions( "single_line_mode=F")), @@ -2108,6 +2114,14 @@ TEST(QueryResponseTest, DebugString) { deleted_row_count: 10 updated_row_count: 10 } + arrow_schema { + serialized_schema: "testing_schema_data" + } + arrow_record_batch { + serialized_record_batch: "testing_batch_data" + row_count: 10 + } + page_row_count: 10 } })"); } @@ -2166,7 +2180,10 @@ TEST(GetQueryResultsResponseTest, DebugString) { R"( description: "" collation: "" default_value_expression: "")" R"( max_length: 0 precision: 0 scale: 0 categories { } policy_tags { })" R"( rounding_mode { value: "" } range_element_type { type: "" } } })" - R"( job_reference { project_id: "p123" job_id: "j123" location: "useast" } } })"); + R"( job_reference { project_id: "p123" job_id: "j123" location: "useast" })" + R"( arrow_schema { serialized_schema: "testing_schema_data" })" + R"( arrow_record_batch { serialized_record_batch: "testing_batch_data" row_count: 10 })" + R"( page_row_count: 10 } })"); EXPECT_EQ( response->DebugString( @@ -2185,7 +2202,10 @@ TEST(GetQueryResultsResponseTest, DebugString) { R"( description: "" collation: "" default_value_expression: "")" R"( max_length: 0 precision: 0 scale: 0 categories { } policy_tags { })" R"( rounding_mode { value: "" } range_element_type { type: "" } } })" - R"( job_reference { project_id: "p123" job_id: "j123" location: "useast" } } })"); + R"( job_reference { project_id: "p123" job_id: "j123" location: "useast" })" + R"( arrow_schema { serialized_schema: "testing......" })" + R"( arrow_record_batch { serialized_record_batch: "testing......" row_count: 10 })" + R"( page_row_count: 10 } })"); EXPECT_EQ( response->DebugString("GetQueryResultsResponse", @@ -2258,6 +2278,14 @@ TEST(GetQueryResultsResponseTest, DebugString) { job_id: "j123" location: "useast" } + arrow_schema { + serialized_schema: "testing_schema_data" + } + arrow_record_batch { + serialized_record_batch: "testing_batch_data" + row_count: 10 + } + page_row_count: 10 } })"); } diff --git a/google/cloud/bigquery/v2/minimal/testing/job_query_test_utils.cc b/google/cloud/bigquery/v2/minimal/testing/job_query_test_utils.cc index 05cd11eb88966..49daebd3c4b2b 100644 --- a/google/cloud/bigquery/v2/minimal/testing/job_query_test_utils.cc +++ b/google/cloud/bigquery/v2/minimal/testing/job_query_test_utils.cc @@ -78,6 +78,7 @@ QueryRequest MakeQueryRequest() { .set_labels(labels) .set_format_options(dfo) .set_job_creation_mode(JobCreationMode::UnSpecified()) + .set_query_results_format("ARROW") .set_default_dataset(MakeDatasetReference()); return expected; @@ -135,6 +136,7 @@ void AssertEquals(QueryRequest const& lhs, QueryRequest const& rhs) { EXPECT_EQ(lhs.default_dataset().project_id, rhs.default_dataset().project_id); EXPECT_EQ(lhs.format_options().use_int64_timestamp, rhs.format_options().use_int64_timestamp); + EXPECT_EQ(lhs.query_results_format(), rhs.query_results_format()); } void AssertEquals(bigquery_v2_minimal_internal::PostQueryRequest const& lhs, @@ -165,6 +167,10 @@ PostQueryResults MakePostQueryResults() { expected.schema = MakeTable().schema; expected.total_bytes_processed = 1000; expected.total_rows = 1000; + expected.arrow_schema.serialized_schema = "testing_schema_data"; + expected.arrow_record_batch.serialized_record_batch = "testing_batch_data"; + expected.arrow_record_batch.row_count = 10; + expected.page_row_count = 10; return expected; } @@ -188,6 +194,10 @@ GetQueryResults MakeGetQueryResults() { expected.schema = MakeTable().schema; expected.total_bytes_processed = 1000; expected.total_rows = 1000; + expected.arrow_schema.serialized_schema = "testing_schema_data"; + expected.arrow_record_batch.serialized_record_batch = "testing_batch_data"; + expected.arrow_record_batch.row_count = 10; + expected.page_row_count = 10; return expected; } @@ -230,6 +240,9 @@ void AssertEquals(bigquery_v2_minimal_internal::PostQueryResults const& lhs, std::equal(lhs.errors.begin(), lhs.errors.end(), rhs.errors.begin())); EXPECT_TRUE(std::equal(lhs.rows.begin(), lhs.rows.end(), rhs.rows.begin())); + EXPECT_EQ(lhs.arrow_schema, rhs.arrow_schema); + EXPECT_EQ(lhs.arrow_record_batch, rhs.arrow_record_batch); + EXPECT_EQ(lhs.page_row_count, rhs.page_row_count); } void AssertEquals(bigquery_v2_minimal_internal::GetQueryResults const& lhs, @@ -255,6 +268,9 @@ void AssertEquals(bigquery_v2_minimal_internal::GetQueryResults const& lhs, std::equal(lhs.errors.begin(), lhs.errors.end(), rhs.errors.begin())); EXPECT_TRUE(std::equal(lhs.rows.begin(), lhs.rows.end(), rhs.rows.begin())); + EXPECT_EQ(lhs.arrow_schema, rhs.arrow_schema); + EXPECT_EQ(lhs.arrow_record_batch, rhs.arrow_record_batch); + EXPECT_EQ(lhs.page_row_count, rhs.page_row_count); } GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END