Skip to content

Commit 69191fc

Browse files
ddelnanoddelnano
committed
[vizier] Refactor carnot to vizierpb result conversion to dedicated
helpers Signed-off-by: Dom Del Nano <ddelnano@gmail.com>
1 parent 319668e commit 69191fc

7 files changed

Lines changed: 436 additions & 134 deletions

File tree

‎src/experimental/standalone_pem/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ pl_cc_library(
5050
"//src/vizier/funcs:cc_library",
5151
"//src/vizier/funcs/context:cc_library",
5252
"//src/vizier/services/agent/shared/base:cc_library",
53+
"//src/vizier/services/agent/shared/vizier_results:cc_library",
5354
"@com_github_grpc_grpc//:grpc++",
5455
],
5556
)

‎src/experimental/standalone_pem/sink_server.h‎

Lines changed: 14 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
#include "src/common/base/base.h"
3232
#include "src/common/base/statuspb/status.pb.h"
3333
#include "src/common/uuid/uuid.h"
34+
#include "src/vizier/services/agent/shared/vizier_results/result_conversion.h"
3435

3536
namespace px {
3637
namespace vizier {
@@ -64,38 +65,25 @@ class StandaloneResultSinkServer final : public carnotpb::ResultSinkService::Ser
6465
consumer = consumer_pair->second;
6566
}
6667

67-
::px::api::vizierpb::ExecuteScriptResponse resp;
68-
resp.set_query_id(query_id.str());
68+
const auto resp = ExecuteScriptResponseFromChunk(*rb, query_id.str());
69+
if (resp.has_value()) {
70+
consumer->Write(*resp);
71+
}
6972

7073
if (rb->has_execution_and_timing_info()) {
71-
HandleExecutionAndTimingInfo(&resp, rb.get());
72-
consumer->Write(resp);
74+
absl::base_internal::SpinLockHolder lock(&id_to_query_consumer_map_lock_);
75+
consumer_map_.erase(query_id);
76+
}
77+
78+
if (rb->has_execution_error() && rb->execution_error().err_code() == 0) {
79+
// err_code 0 is Carnot reporting success; ExecuteScriptResponseFromChunk
80+
// emits nothing for it and the stream ends here.
81+
response->set_success(true);
7382
{
74-
// Query is complete.
7583
absl::base_internal::SpinLockHolder lock(&id_to_query_consumer_map_lock_);
7684
consumer_map_.erase(query_id);
7785
}
78-
}
79-
80-
if (rb->has_query_result()) {
81-
HandleQueryResult(&resp, rb.get());
82-
83-
consumer->Write(resp);
84-
}
85-
86-
if (rb->has_execution_error()) {
87-
auto exec_error = rb->execution_error();
88-
if (exec_error.err_code() == 0) {
89-
response->set_success(true);
90-
{
91-
absl::base_internal::SpinLockHolder lock(&id_to_query_consumer_map_lock_);
92-
consumer_map_.erase(query_id);
93-
}
94-
return ::grpc::Status::OK;
95-
}
96-
auto status = resp.mutable_status();
97-
status->set_message(exec_error.msg());
98-
consumer->Write(resp);
86+
return ::grpc::Status::OK;
9987
}
10088

10189
rb = std::make_unique<carnotpb::TransferResultChunkRequest>();
@@ -113,72 +101,6 @@ class StandaloneResultSinkServer final : public carnotpb::ResultSinkService::Ser
113101
}
114102

115103
private:
116-
void HandleQueryResult(::px::api::vizierpb::ExecuteScriptResponse* resp,
117-
carnotpb::TransferResultChunkRequest* rb) {
118-
auto query_result = rb->query_result();
119-
auto row_data = query_result.row_batch();
120-
121-
auto batch = resp->mutable_data()->mutable_batch();
122-
123-
batch->set_table_id(query_result.table_name());
124-
batch->set_num_rows(row_data.num_rows());
125-
batch->set_eow(row_data.eow());
126-
batch->set_eos(row_data.eos());
127-
128-
for (auto col : row_data.cols()) {
129-
::px::api::vizierpb::Column c;
130-
if (col.has_boolean_data()) {
131-
auto batch_cols = batch->add_cols()->mutable_boolean_data();
132-
for (auto dt : col.boolean_data().data()) {
133-
batch_cols->add_data(dt);
134-
}
135-
}
136-
if (col.has_int64_data()) {
137-
auto batch_cols = batch->add_cols()->mutable_int64_data();
138-
for (auto dt : col.int64_data().data()) {
139-
batch_cols->add_data(dt);
140-
}
141-
}
142-
if (col.has_time64ns_data()) {
143-
auto batch_cols = batch->add_cols()->mutable_time64ns_data();
144-
for (auto dt : col.time64ns_data().data()) {
145-
batch_cols->add_data(dt);
146-
}
147-
}
148-
if (col.has_float64_data()) {
149-
auto batch_cols = batch->add_cols()->mutable_float64_data();
150-
for (auto dt : col.float64_data().data()) {
151-
batch_cols->add_data(dt);
152-
}
153-
}
154-
if (col.has_string_data()) {
155-
auto batch_cols = batch->add_cols()->mutable_string_data();
156-
for (auto dt : col.string_data().data()) {
157-
batch_cols->add_data(dt);
158-
}
159-
}
160-
if (col.has_uint128_data()) {
161-
auto batch_cols = batch->add_cols()->mutable_uint128_data();
162-
for (auto dt : col.uint128_data().data()) {
163-
auto n = batch_cols->add_data();
164-
n->set_low(dt.low());
165-
n->set_high(dt.high());
166-
}
167-
}
168-
}
169-
}
170-
171-
void HandleExecutionAndTimingInfo(::px::api::vizierpb::ExecuteScriptResponse* resp,
172-
carnotpb::TransferResultChunkRequest* rb) {
173-
auto timing_info = rb->execution_and_timing_info();
174-
auto stats = resp->mutable_data()->mutable_execution_stats();
175-
stats->set_bytes_processed(timing_info.execution_stats().bytes_processed());
176-
stats->set_records_processed(timing_info.execution_stats().records_processed());
177-
auto timing = stats->mutable_timing();
178-
timing->set_execution_time_ns(timing_info.execution_stats().timing().execution_time_ns());
179-
timing->set_compilation_time_ns(timing_info.execution_stats().timing().compilation_time_ns());
180-
}
181-
182104
absl::flat_hash_map<sole::uuid, ::grpc::ServerWriter<::px::api::vizierpb::ExecuteScriptResponse>*>
183105
consumer_map_ ABSL_GUARDED_BY(id_to_query_consumer_map_lock_);
184106
mutable absl::base_internal::SpinLock id_to_query_consumer_map_lock_;

‎src/experimental/standalone_pem/vizier_server.h‎

Lines changed: 2 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
#include "src/shared/tracepoint_translation/translation.h"
4242
#include "src/shared/types/typespb/wrapper/types_pb_wrapper.h"
4343
#include "src/table_store/table_store.h"
44+
#include "src/vizier/services/agent/shared/vizier_results/result_conversion.h"
4445

4546
namespace px {
4647
namespace vizier {
@@ -132,48 +133,7 @@ class VizierServer final : public api::vizierpb::VizierService::Service {
132133
return ::grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, error_msg);
133134
}
134135
auto plan = s_or_plan.ConsumeValueOrDie();
135-
for (auto f : plan.nodes()) {
136-
for (auto n : f.nodes()) {
137-
if (n.op().op_type() == carnot::planpb::OperatorType::GRPC_SINK_OPERATOR) {
138-
auto output_table_info = n.op().grpc_sink_op();
139-
if (!output_table_info.has_output_table()) {
140-
continue;
141-
}
142-
::px::api::vizierpb::ExecuteScriptResponse schema_resp;
143-
auto metadata = schema_resp.mutable_meta_data();
144-
metadata->set_name(output_table_info.output_table().table_name());
145-
metadata->set_id(output_table_info.output_table().table_name());
146-
auto rel = metadata->mutable_relation();
147-
for (int i = 0; i < output_table_info.output_table().column_names().size(); i++) {
148-
auto col = rel->add_columns();
149-
col->set_column_name(output_table_info.output_table().column_names()[i]);
150-
switch (output_table_info.output_table().column_types()[i]) {
151-
case types::BOOLEAN:
152-
col->set_column_type(px::api::vizierpb::BOOLEAN);
153-
break;
154-
case types::INT64:
155-
col->set_column_type(px::api::vizierpb::INT64);
156-
break;
157-
case types::UINT128:
158-
col->set_column_type(px::api::vizierpb::UINT128);
159-
break;
160-
case types::FLOAT64:
161-
col->set_column_type(px::api::vizierpb::FLOAT64);
162-
break;
163-
case types::STRING:
164-
col->set_column_type(px::api::vizierpb::STRING);
165-
break;
166-
case types::TIME64NS:
167-
col->set_column_type(px::api::vizierpb::TIME64NS);
168-
break;
169-
default:
170-
break;
171-
}
172-
}
173-
response->Write(schema_resp);
174-
}
175-
}
176-
}
136+
EmitSchemaResponses(plan, query_id.str(), response);
177137

178138
sink_server_->AddConsumer(query_id, response);
179139
auto s = carnot_->ExecuteQuery(reader->query_str(), query_id, px::CurrentTimeNS());

‎src/vizier/services/agent/shared/vizier_results/BUILD.bazel‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# Copyright 2018- The Pixie Authors.
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+
# SPDX-License-Identifier: Apache-2.0
16+
17+
load("//bazel:pl_build_system.bzl", "pl_cc_library", "pl_cc_test")
18+
19+
package(default_visibility = ["//src/experimental/standalone_pem:__pkg__"])
20+
21+
pl_cc_library(
22+
name = "cc_library",
23+
srcs = ["result_conversion.cc"],
24+
hdrs = ["result_conversion.h"],
25+
deps = [
26+
"//src/api/proto/vizierpb:vizier_pl_cc_proto",
27+
"//src/carnot/carnotpb:carnot_pl_cc_proto",
28+
"//src/carnot/planpb:plan_pl_cc_proto",
29+
"//src/shared/types/typespb/wrapper:cc_library",
30+
"//src/table_store/schemapb:schema_pl_cc_proto",
31+
"@com_github_grpc_grpc//:grpc++",
32+
],
33+
)
34+
35+
pl_cc_test(
36+
name = "result_conversion_test",
37+
srcs = ["result_conversion_test.cc"],
38+
deps = [":cc_library"],
39+
)

0 commit comments

Comments
 (0)