From 9105e34c0f21c4be9f57b08698ad4f0bd397bf8f Mon Sep 17 00:00:00 2001 From: Dom Del Nano Date: Thu, 29 Oct 2026 22:57:51 -0700 Subject: [PATCH] [carnot] Add ClickHouse source and export-sink planner IR Signed-off-by: Dom Del Nano --- .../planner/compiler_state/compiler_state.h | 8 +- src/carnot/planner/ir/BUILD.bazel | 17 + src/carnot/planner/ir/all_ir_nodes.h | 2 + .../planner/ir/clickhouse_export_sink_ir.cc | 135 +++++++ .../planner/ir/clickhouse_export_sink_ir.h | 74 ++++ .../ir/clickhouse_export_sink_ir_test.cc | 145 ++++++++ src/carnot/planner/ir/clickhouse_source_ir.cc | 334 ++++++++++++++++++ src/carnot/planner/ir/clickhouse_source_ir.h | 144 ++++++++ .../planner/ir/clickhouse_source_ir_test.cc | 132 +++++++ src/carnot/planner/ir/operators.inl | 2 + src/carnot/planner/ir/pattern_match.h | 6 +- src/carnot/planner/logical_planner.cc | 14 +- 12 files changed, 1009 insertions(+), 4 deletions(-) create mode 100644 src/carnot/planner/ir/clickhouse_export_sink_ir.cc create mode 100644 src/carnot/planner/ir/clickhouse_export_sink_ir.h create mode 100644 src/carnot/planner/ir/clickhouse_export_sink_ir_test.cc create mode 100644 src/carnot/planner/ir/clickhouse_source_ir.cc create mode 100644 src/carnot/planner/ir/clickhouse_source_ir.h create mode 100644 src/carnot/planner/ir/clickhouse_source_ir_test.cc diff --git a/src/carnot/planner/compiler_state/compiler_state.h b/src/carnot/planner/compiler_state/compiler_state.h index cd2e7902f0c..c25a14fe64d 100644 --- a/src/carnot/planner/compiler_state/compiler_state.h +++ b/src/carnot/planner/compiler_state/compiler_state.h @@ -119,7 +119,8 @@ class CompilerState : public NotCopyable { int64_t max_output_rows_per_table, std::string_view result_address, std::string_view result_ssl_targetname, const RedactionOptions& redaction_options, std::unique_ptr endpoint_config, - std::unique_ptr plugin_config, DebugInfo debug_info) + std::unique_ptr plugin_config, DebugInfo debug_info, + std::unique_ptr clickhouse_config = nullptr) : relation_map_(std::move(relation_map)), table_names_to_sensitive_columns_(table_names_to_sensitive_columns), registry_info_(registry_info), @@ -130,7 +131,8 @@ class CompilerState : public NotCopyable { redaction_options_(redaction_options), endpoint_config_(std::move(endpoint_config)), plugin_config_(std::move(plugin_config)), - debug_info_(std::move(debug_info)) {} + debug_info_(std::move(debug_info)), + clickhouse_config_(std::move(clickhouse_config)) {} CompilerState() = delete; @@ -175,6 +177,7 @@ class CompilerState : public NotCopyable { planpb::OTelEndpointConfig* endpoint_config() { return endpoint_config_.get(); } PluginConfig* plugin_config() { return plugin_config_.get(); } const DebugInfo& debug_info() { return debug_info_; } + planpb::ClickHouseConfig* clickhouse_config() { return clickhouse_config_.get(); } private: std::unique_ptr relation_map_; @@ -191,6 +194,7 @@ class CompilerState : public NotCopyable { std::unique_ptr endpoint_config_ = nullptr; std::unique_ptr plugin_config_ = nullptr; DebugInfo debug_info_; + std::unique_ptr clickhouse_config_ = nullptr; }; } // namespace planner diff --git a/src/carnot/planner/ir/BUILD.bazel b/src/carnot/planner/ir/BUILD.bazel index 55b3ac401d4..5940a65197e 100644 --- a/src/carnot/planner/ir/BUILD.bazel +++ b/src/carnot/planner/ir/BUILD.bazel @@ -47,6 +47,7 @@ pl_cc_library( "//src/carnot/planpb:plan_pl_cc_proto", "//src/shared/metadata:cc_library", "//src/shared/metadatapb:metadata_pl_cc_proto", + "@com_github_clickhouse_clickhouse_cpp//:clickhouse_cpp", "@com_github_vinzenz_libpypa//:libpypa", ], ) @@ -67,6 +68,22 @@ pl_cc_test( ], ) +pl_cc_test( + name = "clickhouse_export_sink_ir_test", + srcs = ["clickhouse_export_sink_ir_test.cc"], + deps = [ + "//src/carnot/planner/compiler:test_utils", + ], +) + +pl_cc_test( + name = "clickhouse_source_ir_test", + srcs = ["clickhouse_source_ir_test.cc"], + deps = [ + "//src/carnot/planner/compiler:test_utils", + ], +) + pl_cc_test( name = "pattern_match_test", srcs = ["pattern_match_test.cc"], diff --git a/src/carnot/planner/ir/all_ir_nodes.h b/src/carnot/planner/ir/all_ir_nodes.h index 5c0b49744cd..b5689d1389f 100644 --- a/src/carnot/planner/ir/all_ir_nodes.h +++ b/src/carnot/planner/ir/all_ir_nodes.h @@ -20,6 +20,8 @@ #include "src/carnot/planner/ir/blocking_agg_ir.h" #include "src/carnot/planner/ir/bool_ir.h" +#include "src/carnot/planner/ir/clickhouse_source_ir.h" +#include "src/carnot/planner/ir/clickhouse_export_sink_ir.h" #include "src/carnot/planner/ir/column_ir.h" #include "src/carnot/planner/ir/data_ir.h" #include "src/carnot/planner/ir/drop_ir.h" diff --git a/src/carnot/planner/ir/clickhouse_export_sink_ir.cc b/src/carnot/planner/ir/clickhouse_export_sink_ir.cc new file mode 100644 index 00000000000..d1bdf52227d --- /dev/null +++ b/src/carnot/planner/ir/clickhouse_export_sink_ir.cc @@ -0,0 +1,135 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "src/carnot/planner/ir/clickhouse_export_sink_ir.h" +#include +#include "src/carnot/planner/ir/ir.h" +#include "src/carnot/planpb/plan.pb.h" + +namespace px { +namespace carnot { +namespace planner { + +StatusOr>> +ClickHouseExportSinkIR::RequiredInputColumns() const { + return std::vector>{required_column_names_}; +} + +Status ClickHouseExportSinkIR::Init(OperatorIR* parent, const std::string& table_name, + const std::string& clickhouse_dsn) { + table_name_ = table_name; + + // Parse the ClickHouse DSN and initialize the config + PX_ASSIGN_OR_RETURN(auto config, ParseClickHouseDSN(clickhouse_dsn)); + clickhouse_config_ = std::make_unique(config); + + return AddParent(parent); +} + +StatusOr ClickHouseExportSinkIR::ParseClickHouseDSN( + const std::string& dsn) { + // Expected format: [clickhouse://]username:password@host:port/database + // The clickhouse:// prefix is optional + std::regex dsn_regex(R"((?:clickhouse://)?([^:]+):([^@]+)@([^:]+):(\d+)/(.+))"); + std::smatch matches; + + if (!std::regex_match(dsn, matches, dsn_regex)) { + return error::InvalidArgument( + "Invalid ClickHouse DSN format. Expected: " + "[clickhouse://]username:password@host:port/database"); + } + + planpb::ClickHouseConfig config; + + // Extract the components + config.set_username(matches[1].str()); + config.set_password(matches[2].str()); + config.set_host(matches[3].str()); + config.set_port(std::stoi(matches[4].str())); + config.set_database(matches[5].str()); + + // hostname will be set by the runtime + config.set_hostname(""); + + return config; +} + +Status ClickHouseExportSinkIR::ToProto(planpb::Operator* op) const { + op->set_op_type(planpb::CLICKHOUSE_EXPORT_SINK_OPERATOR); + auto clickhouse_op = op->mutable_clickhouse_sink_op(); + + // ClickHouse config must be set before calling ToProto + if (clickhouse_config_ == nullptr) { + return error::InvalidArgument("ClickHouse config not set"); + } + + // Set the ClickHouse configuration + *clickhouse_op->mutable_clickhouse_config() = *clickhouse_config_; + clickhouse_op->set_table_name(table_name_); + + // Map all input columns to ClickHouse columns + DCHECK(is_type_resolved()); + int64_t idx = 0; + for (const auto& [col_name, col_type] : *resolved_table_type()) { + DCHECK(col_type->IsValueType()); + auto value_type = std::static_pointer_cast(col_type); + + auto column_mapping = clickhouse_op->add_column_mappings(); + column_mapping->set_input_column_index(idx); + column_mapping->set_clickhouse_column_name(col_name); + column_mapping->set_column_type(value_type->data_type()); + idx++; + } + + return Status::OK(); +} + +Status ClickHouseExportSinkIR::CopyFromNodeImpl(const IRNode* node, + absl::flat_hash_map*) { + const ClickHouseExportSinkIR* source = static_cast(node); + table_name_ = source->table_name_; + required_column_names_ = source->required_column_names_; + if (source->clickhouse_config_ != nullptr) { + clickhouse_config_ = std::make_unique(*source->clickhouse_config_); + } + return Status::OK(); +} + +Status ClickHouseExportSinkIR::ResolveType(CompilerState* compiler_state) { + DCHECK_EQ(1U, parent_types().size()); + + auto parent_table_type = std::static_pointer_cast(parent_types()[0]); + + // Store ClickHouse config from compiler state only if not already set by Init() + if (clickhouse_config_ == nullptr && compiler_state->clickhouse_config() != nullptr) { + clickhouse_config_ = + std::make_unique(*compiler_state->clickhouse_config()); + } + + // Populate required column names + for (const auto& col_name : parent_table_type->ColumnNames()) { + required_column_names_.insert(col_name); + } + + // Export sink passes through the input schema + return SetResolvedType(parent_table_type); +} + +} // namespace planner +} // namespace carnot +} // namespace px diff --git a/src/carnot/planner/ir/clickhouse_export_sink_ir.h b/src/carnot/planner/ir/clickhouse_export_sink_ir.h new file mode 100644 index 00000000000..f4bc98246d6 --- /dev/null +++ b/src/carnot/planner/ir/clickhouse_export_sink_ir.h @@ -0,0 +1,74 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include "src/carnot/planner/compiler_state/compiler_state.h" +#include "src/carnot/planner/ir/column_ir.h" +#include "src/carnot/planner/ir/operator_ir.h" +#include "src/carnot/planpb/plan.pb.h" +#include "src/common/base/base.h" + +namespace px { +namespace carnot { +namespace planner { + +/** + * @brief The IR representation for the ClickHouseExportSink operator. + * Represents a configuration to export a DataFrame to a ClickHouse database. + */ +class ClickHouseExportSinkIR : public OperatorIR { + public: + explicit ClickHouseExportSinkIR(int64_t id) : OperatorIR(id, IRNodeType::kClickHouseExportSink) {} + + Status Init(OperatorIR* parent, const std::string& table_name, const std::string& clickhouse_dsn); + + StatusOr ParseClickHouseDSN(const std::string& dsn); + + Status ToProto(planpb::Operator* op) const override; + + Status CopyFromNodeImpl(const IRNode* node, + absl::flat_hash_map*) override; + + Status ResolveType(CompilerState* compiler_state); + inline bool IsBlocking() const override { return true; } + + StatusOr>> RequiredInputColumns() const override; + + const std::string& table_name() const { return table_name_; } + + protected: + StatusOr> PruneOutputColumnsToImpl( + const absl::flat_hash_set& /*kept_columns*/) override { + return error::Unimplemented("Unexpected call to ClickHouseExportSinkIR::PruneOutputColumnsTo."); + } + + private: + std::string table_name_; + absl::flat_hash_set required_column_names_; + std::unique_ptr clickhouse_config_; +}; + +} // namespace planner +} // namespace carnot +} // namespace px diff --git a/src/carnot/planner/ir/clickhouse_export_sink_ir_test.cc b/src/carnot/planner/ir/clickhouse_export_sink_ir_test.cc new file mode 100644 index 00000000000..df801178a6d --- /dev/null +++ b/src/carnot/planner/ir/clickhouse_export_sink_ir_test.cc @@ -0,0 +1,145 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include + +#include +#include +#include + +#include "src/carnot/planner/compiler/test_utils.h" +#include "src/carnot/planner/ir/clickhouse_export_sink_ir.h" +#include "src/carnot/planner/ir/memory_source_ir.h" +#include "src/carnot/planpb/plan.pb.h" +#include "src/common/testing/protobuf.h" +#include "src/table_store/table_store.h" + +namespace px { +namespace carnot { +namespace planner { + +using ClickHouseExportSinkTest = ASTVisitorTest; + +TEST_F(ClickHouseExportSinkTest, basic_export) { + // Create a simple relation with some columns + table_store::schema::Relation relation{ + {types::TIME64NS, types::STRING, types::INT64, types::FLOAT64}, + {"time_", "hostname", "count", "latency"}, + {types::ST_NONE, types::ST_NONE, types::ST_NONE, types::ST_DURATION_NS}}; + + (*compiler_state_->relation_map())["table"] = relation; + + auto src = MakeMemSource("table"); + EXPECT_OK(src->ResolveType(compiler_state_.get())); + + std::string clickhouse_dsn = "default:test_password@localhost:9000/default"; + ASSERT_OK_AND_ASSIGN(auto clickhouse_sink, graph->CreateNode( + src->ast(), src, "http_events", clickhouse_dsn)); + + clickhouse_sink->PullParentTypes(); + EXPECT_OK(clickhouse_sink->UpdateOpAfterParentTypesResolved()); + + // ResolveType will try to get config from compiler state, but we'll set it directly + // by creating a new CompilerState with ClickHouse config + auto new_relation_map = std::make_unique(); + (*new_relation_map)["table"] = relation; + + auto clickhouse_config = std::make_unique(); + clickhouse_config->set_host("localhost"); + clickhouse_config->set_port(9000); + clickhouse_config->set_username("default"); + clickhouse_config->set_password("test_password"); + clickhouse_config->set_database("default"); + + auto new_compiler_state = + std::make_unique(std::move(new_relation_map), SensitiveColumnMap{}, + compiler_state_->registry_info(), compiler_state_->time_now(), + 0, // max_output_rows_per_table + "", // result_address + "", // result_ssl_targetname + RedactionOptions{}, + nullptr, // endpoint_config + nullptr, // plugin_config + DebugInfo{}, std::move(clickhouse_config)); + + // ResolveType will copy the config from compiler state + EXPECT_OK(clickhouse_sink->ResolveType(new_compiler_state.get())); + + planpb::Operator pb; + EXPECT_OK(clickhouse_sink->ToProto(&pb)); + + EXPECT_EQ(pb.op_type(), planpb::CLICKHOUSE_EXPORT_SINK_OPERATOR); + EXPECT_EQ(pb.clickhouse_sink_op().table_name(), "http_events"); + EXPECT_EQ(pb.clickhouse_sink_op().column_mappings_size(), 4); + + // Verify column mappings + EXPECT_EQ(pb.clickhouse_sink_op().column_mappings(0).input_column_index(), 0); + EXPECT_EQ(pb.clickhouse_sink_op().column_mappings(0).clickhouse_column_name(), "time_"); + EXPECT_EQ(pb.clickhouse_sink_op().column_mappings(0).column_type(), types::TIME64NS); + + EXPECT_EQ(pb.clickhouse_sink_op().column_mappings(1).input_column_index(), 1); + EXPECT_EQ(pb.clickhouse_sink_op().column_mappings(1).clickhouse_column_name(), "hostname"); + EXPECT_EQ(pb.clickhouse_sink_op().column_mappings(1).column_type(), types::STRING); +} + +TEST_F(ClickHouseExportSinkTest, required_input_columns) { + table_store::schema::Relation relation{{types::TIME64NS, types::STRING, types::INT64}, + {"time_", "hostname", "count"}, + {types::ST_NONE, types::ST_NONE, types::ST_NONE}}; + + (*compiler_state_->relation_map())["table"] = relation; + + auto src = MakeMemSource("table"); + EXPECT_OK(src->ResolveType(compiler_state_.get())); + + std::string clickhouse_dsn = "default:test_password@localhost:9000/default"; + ASSERT_OK_AND_ASSIGN(auto clickhouse_sink, graph->CreateNode( + src->ast(), src, "http_events", clickhouse_dsn)); + + clickhouse_sink->PullParentTypes(); + EXPECT_OK(clickhouse_sink->UpdateOpAfterParentTypesResolved()); + + // Need to call ResolveType to populate required_column_names_ + auto clickhouse_config = std::make_unique(); + clickhouse_config->set_host("localhost"); + clickhouse_config->set_port(9000); + clickhouse_config->set_username("default"); + clickhouse_config->set_password("test_password"); + clickhouse_config->set_database("default"); + + auto new_relation_map = std::make_unique(); + (*new_relation_map)["table"] = relation; + + auto new_compiler_state = std::make_unique( + std::move(new_relation_map), SensitiveColumnMap{}, compiler_state_->registry_info(), + compiler_state_->time_now(), 0, "", "", RedactionOptions{}, nullptr, nullptr, DebugInfo{}, + std::move(clickhouse_config)); + + EXPECT_OK(clickhouse_sink->ResolveType(new_compiler_state.get())); + + ASSERT_OK_AND_ASSIGN(auto required_input_columns, clickhouse_sink->RequiredInputColumns()); + ASSERT_EQ(required_input_columns.size(), 1); + EXPECT_THAT(required_input_columns[0], + ::testing::UnorderedElementsAre("time_", "hostname", "count")); +} + +} // namespace planner +} // namespace carnot +} // namespace px diff --git a/src/carnot/planner/ir/clickhouse_source_ir.cc b/src/carnot/planner/ir/clickhouse_source_ir.cc new file mode 100644 index 00000000000..0d6035ae744 --- /dev/null +++ b/src/carnot/planner/ir/clickhouse_source_ir.cc @@ -0,0 +1,334 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "src/carnot/planner/ir/clickhouse_source_ir.h" + +#include + +#include "src/carnot/planner/ir/ir.h" + +namespace px { +namespace carnot { +namespace planner { + +std::string ClickHouseSourceIR::DebugString() const { + return absl::Substitute("$0(id=$1, table=$2)", type_string(), id(), table_name_); +} + +Status ClickHouseSourceIR::ToProto(planpb::Operator* op) const { + auto pb = op->mutable_clickhouse_source_op(); + op->set_op_type(planpb::CLICKHOUSE_SOURCE_OPERATOR); + + // Set ClickHouse connection parameters from stored values + pb->set_host(host_); + pb->set_port(port_); + pb->set_username(username_); + pb->set_password(password_); + pb->set_database(database_); + + if (!column_index_map_set()) { + return error::InvalidArgument("ClickHouseSource columns are not set."); + } + + DCHECK(is_type_resolved()); + DCHECK_EQ(column_index_map_.size(), resolved_table_type()->ColumnNames().size()); + + // Build the query with explicit column list to match output_descriptor_ order + std::vector column_list; + for (const auto& [idx, col_name] : Enumerate(resolved_table_type()->ColumnNames())) { + column_list.push_back(col_name); + pb->add_column_names(col_name); + auto val_type = std::static_pointer_cast( + resolved_table_type()->GetColumnType(col_name).ConsumeValueOrDie()); + pb->add_column_types(val_type->data_type()); + } + + // Generate SELECT with explicit columns instead of SELECT * to ensure correct column ordering + pb->set_query( + absl::Substitute("SELECT $0 FROM $1", absl::StrJoin(column_list, ", "), table_name_)); + + if (IsTimeStartSet()) { + pb->set_start_time(time_start_ns()); + } + + if (IsTimeStopSet()) { + pb->set_end_time(time_stop_ns()); + } + + // Set batch size + pb->set_batch_size(1024); + + // Set timestamp and partition columns from stored values + pb->set_timestamp_column(timestamp_column_); + pb->set_partition_column("hostname"); + + return Status::OK(); +} + +Status ClickHouseSourceIR::Init(const std::string& table_name, + const std::vector& select_columns, + const std::string& host, int port, const std::string& username, + const std::string& password, const std::string& database, + const std::string& timestamp_column) { + table_name_ = table_name; + column_names_ = select_columns; + host_ = host; + port_ = port; + username_ = username; + password_ = password; + database_ = database; + timestamp_column_ = timestamp_column; + return Status::OK(); +} + +StatusOr> ClickHouseSourceIR::PruneOutputColumnsToImpl( + const absl::flat_hash_set& output_colnames) { + DCHECK(column_index_map_set()); + DCHECK(is_type_resolved()); + std::vector new_col_names; + std::vector new_col_index_map; + + auto col_names = resolved_table_type()->ColumnNames(); + for (const auto& [idx, name] : Enumerate(col_names)) { + if (output_colnames.contains(name)) { + new_col_names.push_back(name); + new_col_index_map.push_back(column_index_map_[idx]); + } + } + if (new_col_names != resolved_table_type()->ColumnNames()) { + column_names_ = new_col_names; + } + column_index_map_ = new_col_index_map; + return output_colnames; +} + +Status ClickHouseSourceIR::CopyFromNodeImpl(const IRNode* node, + absl::flat_hash_map*) { + const ClickHouseSourceIR* source_ir = static_cast(node); + + table_name_ = source_ir->table_name_; + time_start_ns_ = source_ir->time_start_ns_; + time_stop_ns_ = source_ir->time_stop_ns_; + column_names_ = source_ir->column_names_; + column_index_map_set_ = source_ir->column_index_map_set_; + column_index_map_ = source_ir->column_index_map_; + + username_ = source_ir->username_; + password_ = source_ir->password_; + database_ = source_ir->database_; + port_ = source_ir->port_; + host_ = source_ir->host_; + + return Status::OK(); +} + +StatusOr ClickHouseSourceIR::ClickHouseTypeToPixieType( + const std::string& ch_type_name) { + // Integer types - Pixie only supports INT64 + if (ch_type_name == "UInt8" || ch_type_name == "UInt16" || ch_type_name == "UInt32" || + ch_type_name == "UInt64" || ch_type_name == "Int8" || ch_type_name == "Int16" || + ch_type_name == "Int32" || ch_type_name == "Int64") { + return types::DataType::INT64; + } + // UInt128 + if (ch_type_name == "UInt128") { + return types::DataType::UINT128; + } + // Floating point types - Pixie only supports FLOAT64 + if (ch_type_name == "Float32" || ch_type_name == "Float64") { + return types::DataType::FLOAT64; + } + // String types + if (ch_type_name == "String" || ch_type_name == "FixedString" || + absl::StartsWith(ch_type_name, "FixedString(")) { + return types::DataType::STRING; + } + // Date/time types + if (ch_type_name == "DateTime" || absl::StartsWith(ch_type_name, "DateTime64")) { + return types::DataType::TIME64NS; + } + // Boolean type (stored as UInt8 in ClickHouse) + if (ch_type_name == "Bool") { + return types::DataType::BOOLEAN; + } + return types::DataType::STRING; +} + +StatusOr ClickHouseSourceIR::InferRelationFromClickHouse( + CompilerState* compiler_state, const std::string& table_name) { + // Check if ClickHouse config is available + // TODO(ddelnano): Add this check in when the configuration plumbing is done. + auto* ch_config = compiler_state->clickhouse_config(); + PX_UNUSED(ch_config); + + // Use stored connection parameters from Init() + + clickhouse::ClientOptions options; + options.SetHost(host_); + options.SetPort(port_); + options.SetUser(username_); + options.SetPassword(password_); + options.SetDefaultDatabase(database_); + + // Create ClickHouse client + std::unique_ptr client; + try { + client = std::make_unique(options); + } catch (const std::exception& e) { + return error::Internal("Failed to connect to ClickHouse at $0:$1 - $2", host_, port_, e.what()); + } + + // Query ClickHouse for table schema using DESCRIBE TABLE + std::string describe_query = absl::Substitute("DESCRIBE TABLE $0", table_name); + + table_store::schema::Relation relation; + bool query_executed = false; + + try { + client->Select(describe_query, [&](const clickhouse::Block& block) { + query_executed = true; + // DESCRIBE TABLE returns columns: name, type, default_type, default_expression, comment, + // codec_expression, ttl_expression + size_t num_rows = block.GetRowCount(); + + if (num_rows == 0) { + return; + } + + // Get the column name and type columns + auto name_column = block[0]->As(); + auto type_column = block[1]->As(); + + for (size_t i = 0; i < num_rows; ++i) { + std::string col_name = std::string(name_column->At(i)); + std::string col_type = std::string(type_column->At(i)); + + // Convert ClickHouse type to Pixie type + auto pixie_type_or = ClickHouseTypeToPixieType(col_type); + if (!pixie_type_or.ok()) { + LOG(WARNING) << "Failed to convert ClickHouse type '" << col_type << "' for column '" + << col_name << "'. Using STRING as fallback."; + relation.AddColumn(types::DataType::STRING, col_name, types::SemanticType::ST_NONE); + } else { + types::DataType pixie_type = pixie_type_or.ConsumeValueOrDie(); + // Determine semantic type based on column name or type + types::SemanticType semantic_type = types::SemanticType::ST_NONE; + if (pixie_type == types::DataType::TIME64NS) { + semantic_type = types::SemanticType::ST_TIME_NS; + } + relation.AddColumn(pixie_type, col_name, semantic_type); + } + } + }); + } catch (const std::exception& e) { + return error::Internal("Failed to query ClickHouse table schema for '$0': $1", table_name, + e.what()); + } + + if (!query_executed || relation.NumColumns() == 0) { + return error::Internal("Table '$0' not found in ClickHouse or has no columns.", table_name); + } + + return relation; +} + +Status ClickHouseSourceIR::ResolveType(CompilerState* compiler_state) { + table_store::schema::Relation table_relation; + + auto existing_relation = false; + auto relation_it = compiler_state->relation_map()->find(table_name()); + if (relation_it == compiler_state->relation_map()->end()) { + // Table not found in relation_map, try to infer from ClickHouse + VLOG(1) << absl::Substitute( + "Table '$0' not found in relation_map. Attempting to infer schema from ClickHouse...", + table_name()); + + auto relation_or = InferRelationFromClickHouse(compiler_state, table_name()); + if (!relation_or.ok()) { + return CreateIRNodeError( + "Table '$0' not found in relation_map and failed to infer from ClickHouse: $1", + table_name_, relation_or.status().msg()); + } + + table_relation = relation_or.ConsumeValueOrDie(); + } else { + table_relation = relation_it->second; + existing_relation = true; + } + auto full_table_type = TableType::Create(table_relation); + if (select_all()) { + // For select_all, add all table columns plus ClickHouse-added columns (hostname, event_time) + std::vector column_indices; + int64_t table_column_count = static_cast(table_relation.NumColumns()); + + // Add all table columns + for (int64_t i = 0; i < table_column_count; ++i) { + column_indices.push_back(i); + } + + // Add ClickHouse-added columns + if (existing_relation) { + full_table_type->AddColumn( + "hostname", ValueType::Create(types::DataType::STRING, types::SemanticType::ST_NONE)); + column_indices.push_back(table_column_count); // hostname is after all table columns + + full_table_type->AddColumn("event_time", ValueType::Create(types::DataType::TIME64NS, + types::SemanticType::ST_TIME_NS)); + column_indices.push_back(table_column_count + 1); // event_time is after hostname + } + + SetColumnIndexMap(column_indices); + return SetResolvedType(full_table_type); + } + + std::vector column_indices; + auto new_table = TableType::Create(); + + // Calculate the index offset for ClickHouse-added columns (after all table columns) + int64_t table_column_count = static_cast(table_relation.NumColumns()); + auto next_count = 0; + + for (const auto& col_name : column_names_) { + // Handle special ClickHouse-added columns that don't exist in the source table + if (col_name == "hostname") { + new_table->AddColumn( + col_name, ValueType::Create(types::DataType::STRING, types::SemanticType::ST_NONE)); + // hostname is added by ClickHouse after all table columns + column_indices.push_back(table_column_count + (next_count++)); + continue; + } + if (col_name == "event_time") { + new_table->AddColumn( + col_name, ValueType::Create(types::DataType::TIME64NS, types::SemanticType::ST_TIME_NS)); + // event_time is added by ClickHouse after hostname + column_indices.push_back(table_column_count + (next_count++)); + continue; + } + + PX_ASSIGN_OR_RETURN(auto col_type, full_table_type->GetColumnType(col_name)); + new_table->AddColumn(col_name, col_type); + column_indices.push_back(table_relation.GetColumnIndex(col_name)); + } + + SetColumnIndexMap(column_indices); + return SetResolvedType(new_table); +} + +} // namespace planner +} // namespace carnot +} // namespace px diff --git a/src/carnot/planner/ir/clickhouse_source_ir.h b/src/carnot/planner/ir/clickhouse_source_ir.h new file mode 100644 index 00000000000..1f578e7bcef --- /dev/null +++ b/src/carnot/planner/ir/clickhouse_source_ir.h @@ -0,0 +1,144 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include + +#include "src/carnot/planner/compiler_state/compiler_state.h" +#include "src/carnot/planner/ir/expression_ir.h" +#include "src/carnot/planner/ir/operator_ir.h" +#include "src/carnot/planner/types/types.h" +#include "src/common/base/base.h" +#include "src/shared/types/types.h" +#include "src/table_store/schema/relation.h" + +namespace px { +namespace carnot { +namespace planner { + +/** + * @brief The ClickHouseSourceIR represents a source that reads data from a ClickHouse database. + */ +class ClickHouseSourceIR : public OperatorIR { + public: + ClickHouseSourceIR() = delete; + explicit ClickHouseSourceIR(int64_t id) : OperatorIR(id, IRNodeType::kClickHouseSource) {} + + /** + * @brief Initialize the ClickHouse source. + * + * @param table_name the table to load. + * @param select_columns the columns to select. If vector is empty, then select all columns. + * @param host the ClickHouse server host. + * @param port the ClickHouse server port. + * @param username the ClickHouse username. + * @param password the ClickHouse password. + * @param database the ClickHouse database. + * @return Status + */ + Status Init(const std::string& table_name, const std::vector& select_columns, + const std::string& host = "localhost", int port = 9000, + const std::string& username = "default", const std::string& password = "", + const std::string& database = "default", + const std::string& timestamp_column = "event_time"); + + std::string table_name() const { return table_name_; } + std::string host() const { return host_; } + int port() const { return port_; } + std::string username() const { return username_; } + std::string password() const { return password_; } + std::string database() const { return database_; } + std::string timestamp_column() const { return timestamp_column_; } + + void SetTimeStartNS(int64_t time_start_ns) { time_start_ns_ = time_start_ns; } + void SetTimeStopNS(int64_t time_stop_ns) { time_stop_ns_ = time_stop_ns; } + bool IsTimeStartSet() const { return time_start_ns_.has_value(); } + bool IsTimeStopSet() const { return time_stop_ns_.has_value(); } + + std::string DebugString() const override; + + int64_t time_start_ns() const { return time_start_ns_.value(); } + int64_t time_stop_ns() const { return time_stop_ns_.value(); } + + const std::vector& column_index_map() const { return column_index_map_; } + bool column_index_map_set() const { return column_index_map_set_; } + void SetColumnIndexMap(const std::vector& column_index_map) { + column_index_map_set_ = true; + column_index_map_ = column_index_map; + } + + Status ToProto(planpb::Operator*) const override; + + bool select_all() const { return column_names_.size() == 0; } + + Status CopyFromNodeImpl(const IRNode* node, + absl::flat_hash_map* copied_nodes_map) override; + const std::vector& column_names() const { return column_names_; } + + StatusOr>> RequiredInputColumns() const override { + return std::vector>{}; + } + + void SetColumnNames(const std::vector& col_names) { column_names_ = col_names; } + + bool IsSource() const override { return true; } + + Status ResolveType(CompilerState* compiler_state); + + protected: + // Helper method to query ClickHouse for table schema and create a Relation + StatusOr InferRelationFromClickHouse( + CompilerState* compiler_state, const std::string& table_name); + + // Helper method to convert ClickHouse type string to Pixie DataType + static StatusOr ClickHouseTypeToPixieType(const std::string& ch_type_name); + + StatusOr> PruneOutputColumnsToImpl( + const absl::flat_hash_set& output_colnames) override; + + private: + std::string table_name_; + + // ClickHouse connection parameters + std::string host_ = "localhost"; + int port_ = 9000; + std::string username_ = "default"; + std::string password_ = ""; + std::string database_ = "default"; + + // ClickHouse column configuration + std::string timestamp_column_ = "event_time"; + + std::optional time_start_ns_; + std::optional time_stop_ns_; + + // Hold of columns in the order that they are selected. + std::vector column_names_; + + // The mapping of the source's column indices to the current columns, as given by column_names_. + std::vector column_index_map_; + bool column_index_map_set_ = false; +}; + +} // namespace planner +} // namespace carnot +} // namespace px diff --git a/src/carnot/planner/ir/clickhouse_source_ir_test.cc b/src/carnot/planner/ir/clickhouse_source_ir_test.cc new file mode 100644 index 00000000000..e0664c92c04 --- /dev/null +++ b/src/carnot/planner/ir/clickhouse_source_ir_test.cc @@ -0,0 +1,132 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include + +#include + +#include "src/carnot/planner/compiler/test_utils.h" +#include "src/carnot/planner/ir/clickhouse_source_ir.h" +#include "src/carnot/planpb/plan.pb.h" +#include "src/table_store/table_store.h" + +namespace px { +namespace carnot { +namespace planner { + +using ClickHouseSourceTest = ASTVisitorTest; + +namespace { +// Seeds the compiler state's relation map with a small http_events-like table so +// ResolveType resolves against it instead of connecting to a live ClickHouse +// server (schema inference via DESCRIBE TABLE is exercised end-to-end by the +// clickhouse_source_node exec test). +table_store::schema::Relation HttpEventsRelation() { + return table_store::schema::Relation{ + {types::STRING, types::INT64}, {"req_path", "resp_status"}, {types::ST_NONE, types::ST_NONE}}; +} +} // namespace + +// Explicitly-selected columns, including the ClickHouse-provided hostname and +// event_time columns, serialize to a ClickHouseSourceOperator with matching +// column names/types and an explicit SELECT query. +TEST_F(ClickHouseSourceTest, explicit_columns_to_proto) { + (*compiler_state_->relation_map())["http_events"] = HttpEventsRelation(); + + std::vector columns = {"req_path", "hostname", "event_time"}; + ASSERT_OK_AND_ASSIGN(auto* source, graph->CreateNode( + ast, "http_events", columns, "localhost", 9000, "default", + "test_password", "default", "event_time")); + + EXPECT_OK(source->ResolveType(compiler_state_.get())); + + planpb::Operator pb; + ASSERT_OK(source->ToProto(&pb)); + + EXPECT_EQ(pb.op_type(), planpb::CLICKHOUSE_SOURCE_OPERATOR); + const auto& ch = pb.clickhouse_source_op(); + EXPECT_EQ(ch.host(), "localhost"); + EXPECT_EQ(ch.port(), 9000); + EXPECT_EQ(ch.username(), "default"); + EXPECT_EQ(ch.password(), "test_password"); + EXPECT_EQ(ch.database(), "default"); + EXPECT_EQ(ch.timestamp_column(), "event_time"); + EXPECT_EQ(ch.partition_column(), "hostname"); + + ASSERT_EQ(ch.column_names_size(), 3); + EXPECT_EQ(ch.column_names(0), "req_path"); + EXPECT_EQ(ch.column_names(1), "hostname"); + EXPECT_EQ(ch.column_names(2), "event_time"); + + ASSERT_EQ(ch.column_types_size(), 3); + EXPECT_EQ(ch.column_types(0), types::STRING); + EXPECT_EQ(ch.column_types(1), types::STRING); + EXPECT_EQ(ch.column_types(2), types::TIME64NS); + + EXPECT_EQ(ch.query(), "SELECT req_path, hostname, event_time FROM http_events"); +} + +// An empty column selection resolves to all table columns plus the +// ClickHouse-provided hostname and event_time columns. +TEST_F(ClickHouseSourceTest, select_all_to_proto) { + (*compiler_state_->relation_map())["http_events"] = HttpEventsRelation(); + + std::vector columns = {}; + ASSERT_OK_AND_ASSIGN(auto* source, graph->CreateNode( + ast, "http_events", columns, "localhost", 9000, "default", + "test_password", "default", "event_time")); + + EXPECT_OK(source->ResolveType(compiler_state_.get())); + + planpb::Operator pb; + ASSERT_OK(source->ToProto(&pb)); + + const auto& ch = pb.clickhouse_source_op(); + ASSERT_EQ(ch.column_names_size(), 4); + EXPECT_EQ(ch.column_names(0), "req_path"); + EXPECT_EQ(ch.column_names(1), "resp_status"); + EXPECT_EQ(ch.column_names(2), "hostname"); + EXPECT_EQ(ch.column_names(3), "event_time"); + EXPECT_EQ(ch.query(), "SELECT req_path, resp_status, hostname, event_time FROM http_events"); +} + +// Start/stop times set on the IR are propagated into the operator proto. +TEST_F(ClickHouseSourceTest, time_range_to_proto) { + (*compiler_state_->relation_map())["http_events"] = HttpEventsRelation(); + + std::vector columns = {"req_path", "hostname", "event_time"}; + ASSERT_OK_AND_ASSIGN(auto* source, graph->CreateNode( + ast, "http_events", columns, "localhost", 9000, "default", + "test_password", "default", "event_time")); + source->SetTimeStartNS(1000); + source->SetTimeStopNS(2000); + + EXPECT_OK(source->ResolveType(compiler_state_.get())); + + planpb::Operator pb; + ASSERT_OK(source->ToProto(&pb)); + + EXPECT_EQ(pb.clickhouse_source_op().start_time(), 1000); + EXPECT_EQ(pb.clickhouse_source_op().end_time(), 2000); +} + +} // namespace planner +} // namespace carnot +} // namespace px diff --git a/src/carnot/planner/ir/operators.inl b/src/carnot/planner/ir/operators.inl index 817295e3a6e..bb712c71c11 100644 --- a/src/carnot/planner/ir/operators.inl +++ b/src/carnot/planner/ir/operators.inl @@ -37,5 +37,7 @@ PX_CARNOT_IR_NODE(Rolling) PX_CARNOT_IR_NODE(Stream) PX_CARNOT_IR_NODE(EmptySource) PX_CARNOT_IR_NODE(OTelExportSink) +PX_CARNOT_IR_NODE(ClickHouseSource) +PX_CARNOT_IR_NODE(ClickHouseExportSink) #endif diff --git a/src/carnot/planner/ir/pattern_match.h b/src/carnot/planner/ir/pattern_match.h index f8c484f47b9..0eb386ddbc5 100644 --- a/src/carnot/planner/ir/pattern_match.h +++ b/src/carnot/planner/ir/pattern_match.h @@ -160,6 +160,10 @@ inline ClassMatch OTelExportSink() { return ClassMatch(); } +inline ClassMatch ClickHouseExportSink() { + return ClassMatch(); +} + inline ClassMatch EmptySource() { return ClassMatch(); } @@ -266,7 +270,7 @@ struct ResultSink : public ParentMatch { bool Match(const IRNode* node) const override { return ExternalGRPCSink().Match(node) || MemorySink().Match(node) || - OTelExportSink().Match(node); + OTelExportSink().Match(node) || ClickHouseExportSink().Match(node); } }; diff --git a/src/carnot/planner/logical_planner.cc b/src/carnot/planner/logical_planner.cc index 19ed07104cf..df840098a26 100644 --- a/src/carnot/planner/logical_planner.cc +++ b/src/carnot/planner/logical_planner.cc @@ -97,6 +97,18 @@ StatusOr> CreateCompilerState( for (const auto& debug_info_pb : logical_state.debug_info().otel_debug_attributes()) { debug_info.otel_debug_attrs.push_back({debug_info_pb.name(), debug_info_pb.value()}); } + + std::unique_ptr clickhouse_config = nullptr; + if (logical_state.has_clickhouse_config()) { + clickhouse_config = std::make_unique(); + clickhouse_config->set_hostname(logical_state.clickhouse_config().hostname()); + clickhouse_config->set_host(logical_state.clickhouse_config().host()); + clickhouse_config->set_port(logical_state.clickhouse_config().port()); + clickhouse_config->set_username(logical_state.clickhouse_config().username()); + clickhouse_config->set_password(logical_state.clickhouse_config().password()); + clickhouse_config->set_database(logical_state.clickhouse_config().database()); + } + // Create a CompilerState obj using the relation map and grabbing the current time. return std::make_unique( std::move(rel_map), sensitive_columns, registry_info, px::CurrentTimeNS(), @@ -105,7 +117,7 @@ StatusOr> CreateCompilerState( // TODO(philkuz) add an endpoint config to logical_state and pass that in here. RedactionOptionsFromPb(logical_state.redaction_options()), std::move(otel_endpoint_config), // TODO(philkuz) propagate the otel debug attributes here. - std::move(plugin_config), debug_info); + std::move(plugin_config), debug_info, std::move(clickhouse_config)); } StatusOr> LogicalPlanner::Create(const udfspb::UDFInfo& udf_info) {