[carnot] Add ClickHouse source and export-sink execution nodes

src/carnot/carnot.cc

@@ -181,6 +181,8 @@ Status CarnotImpl::RegisterUDFsInPlanFragment(exec::ExecState* exec_state, plan:
       .OnUDTFSource(no_op)
       .OnEmptySource(no_op)
       .OnOTelSink(no_op)
+      .OnClickHouseSource(no_op)
+      .OnClickHouseExportSink(no_op)
       .Walk(pf);
 }
 

src/carnot/exec/BUILD.bazel

@@ -33,6 +33,7 @@ pl_cc_library(
         ],
     ),
     hdrs = [
+        "clickhouse_source_node.h",
         "exec_node.h",
         "exec_state.h",
     ],
@@ -42,10 +43,12 @@ pl_cc_library(
         "//src/carnot/plan:cc_library",
         "//src/carnot/planpb:plan_pl_cc_proto",
         "//src/carnot/udf:cc_library",
+        "//src/common/ch:cc_library",
         "//src/common/uuid:cc_library",
         "//src/shared/types:cc_library",
         "//src/table_store/table:cc_library",
         "@com_github_apache_arrow//:arrow",
+        "@com_github_clickhouse_clickhouse_cpp//:clickhouse_cpp",
         "@com_github_grpc_grpc//:grpc++",
         "@com_github_opentelemetry_proto//:logs_service_grpc_cc",
         "@com_github_opentelemetry_proto//:metrics_service_grpc_cc",
@@ -300,3 +303,50 @@ pl_cc_test(
         "@com_github_grpc_grpc//:grpc++_test",
     ],
 )
+
+pl_cc_test(
+    name = "clickhouse_source_node_test",
+    timeout = "long",
+    srcs = ["clickhouse_source_node_test.cc"],
+    data = [
+        "//src/stirling/source_connectors/socket_tracer/testing/container_images:clickhouse.tar",
+    ],
+    flaky = True,
+    tags = [
+        "exclusive",
+        "requires_root",
+    ],
+    deps = [
+        ":cc_library",
+        ":exec_node_test_helpers",
+        ":test_utils",
+        "//src/carnot/planpb:plan_testutils",
+        "//src/common/ch:cc_library",
+        "//src/common/testing/test_utils:cc_library",
+        "@com_github_clickhouse_clickhouse_cpp//:clickhouse_cpp",
+    ],
+)
+
+pl_cc_test(
+    name = "clickhouse_export_sink_node_test",
+    timeout = "long",
+    srcs = ["clickhouse_export_sink_node_test.cc"],
+    data = [
+        "//src/stirling/source_connectors/socket_tracer/testing/container_images:clickhouse.tar",
+    ],
+    flaky = True,
+    tags = [
+        "exclusive",
+        "requires_root",
+    ],
+    deps = [
+        ":cc_library",
+        ":exec_node_test_helpers",
+        ":test_utils",
+        "//src/carnot/plan:cc_library",
+        "//src/carnot/planpb:plan_pl_cc_proto",
+        "//src/common/ch:cc_library",
+        "//src/common/testing/test_utils:cc_library",
+        "@com_github_clickhouse_clickhouse_cpp//:clickhouse_cpp",
+    ],
+)

src/carnot/exec/clickhouse_export_sink_node.cc

@@ -0,0 +1,181 @@
+/*
+ * 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/exec/clickhouse_export_sink_node.h"
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <absl/numeric/int128.h>
+#include <absl/strings/substitute.h>
+#include <sole.hpp>
+#include "glog/logging.h"
+#include "src/carnot/planpb/plan.pb.h"
+#include "src/common/base/macros.h"
+#include "src/common/ch/utils.h"
+#include "src/shared/types/typespb/types.pb.h"
+#include "src/table_store/table_store.h"
+
+namespace px {
+namespace carnot {
+namespace exec {
+
+using table_store::schema::RowBatch;
+using table_store::schema::RowDescriptor;
+
+std::string ClickHouseExportSinkNode::DebugStringImpl() {
+  return absl::Substitute("Exec::ClickHouseExportSinkNode: $0", plan_node_->DebugString());
+}
+
+Status ClickHouseExportSinkNode::InitImpl(const plan::Operator& plan_node) {
+  CHECK(plan_node.op_type() == planpb::OperatorType::CLICKHOUSE_EXPORT_SINK_OPERATOR);
+  if (input_descriptors_.size() != 1) {
+    return error::InvalidArgument(
+        "ClickHouse Export operator expects a single input relation, got $0",
+        input_descriptors_.size());
+  }
+
+  input_descriptor_ = std::make_unique<RowDescriptor>(input_descriptors_[0]);
+  const auto* sink_plan_node = static_cast<const plan::ClickHouseExportSinkOperator*>(&plan_node);
+  plan_node_ = std::make_unique<plan::ClickHouseExportSinkOperator>(*sink_plan_node);
+  return Status::OK();
+}
+
+Status ClickHouseExportSinkNode::PrepareImpl(ExecState*) { return Status::OK(); }
+
+Status ClickHouseExportSinkNode::OpenImpl(ExecState* /*exec_state*/) {
+  // Connect to ClickHouse using config from plan node
+  const auto& config = plan_node_->clickhouse_config();
+
+  clickhouse::ClientOptions options;
+  options.SetHost(config.host());
+  options.SetPort(config.port());
+  options.SetUser(config.username());
+  options.SetPassword(config.password());
+  options.SetDefaultDatabase(config.database());
+
+  PX_ASSIGN_OR_RETURN(clickhouse_client_, ch::CreateClient(options));
+
+  return Status::OK();
+}
+
+Status ClickHouseExportSinkNode::CloseImpl(ExecState* exec_state) {
+  if (sent_eos_) {
+    return Status::OK();
+  }
+
+  LOG(INFO) << absl::Substitute(
+      "Closing ClickHouseExportSinkNode $0 in query $1 before receiving EOS", plan_node_->id(),
+      exec_state->query_id().str());
+
+  return Status::OK();
+}
+
+Status ClickHouseExportSinkNode::ConsumeNextImpl(ExecState* /*exec_state*/, const RowBatch& rb,
+                                                 size_t /*parent_index*/) {
+  // Skip insertion if the batch is empty
+  if (rb.num_rows() == 0) {
+    if (rb.eos()) {
+      sent_eos_ = true;
+    }
+    return Status::OK();
+  }
+
+  // Build an INSERT query with the data from the row batch
+  clickhouse::Block block;
+
+  // Create columns based on the column mappings
+  for (const auto& mapping : plan_node_->column_mappings()) {
+    auto arrow_col = rb.ColumnAt(mapping.input_column_index());
+    int64_t num_rows = arrow_col->length();
+
+    // Create ClickHouse column based on data type
+    switch (mapping.column_type()) {
+      case types::INT64: {
+        auto col = std::make_shared<clickhouse::ColumnInt64>();
+        for (int64_t i = 0; i < num_rows; ++i) {
+          col->Append(types::GetValueFromArrowArray<types::INT64>(arrow_col.get(), i));
+        }
+        block.AppendColumn(mapping.clickhouse_column_name(), col);
+        break;
+      }
+      case types::FLOAT64: {
+        auto col = std::make_shared<clickhouse::ColumnFloat64>();
+        for (int64_t i = 0; i < num_rows; ++i) {
+          col->Append(types::GetValueFromArrowArray<types::FLOAT64>(arrow_col.get(), i));
+        }
+        block.AppendColumn(mapping.clickhouse_column_name(), col);
+        break;
+      }
+      case types::STRING: {
+        auto col = std::make_shared<clickhouse::ColumnString>();
+        for (int64_t i = 0; i < num_rows; ++i) {
+          col->Append(types::GetValueFromArrowArray<types::STRING>(arrow_col.get(), i));
+        }
+        block.AppendColumn(mapping.clickhouse_column_name(), col);
+        break;
+      }
+      case types::TIME64NS: {
+        auto col = std::make_shared<clickhouse::ColumnDateTime64>(9);
+        for (int64_t i = 0; i < num_rows; ++i) {
+          int64_t ns_val = types::GetValueFromArrowArray<types::TIME64NS>(arrow_col.get(), i);
+          col->Append(ns_val);
+        }
+        block.AppendColumn(mapping.clickhouse_column_name(), col);
+        break;
+      }
+      case types::BOOLEAN: {
+        auto col = std::make_shared<clickhouse::ColumnUInt8>();
+        for (int64_t i = 0; i < num_rows; ++i) {
+          col->Append(types::GetValueFromArrowArray<types::BOOLEAN>(arrow_col.get(), i) ? 1 : 0);
+        }
+        block.AppendColumn(mapping.clickhouse_column_name(), col);
+        break;
+      }
+      case types::UINT128: {
+        // UINT128 is exported as STRING (UUID format)
+        auto col = std::make_shared<clickhouse::ColumnString>();
+        for (int64_t i = 0; i < num_rows; ++i) {
+          auto val = types::GetValueFromArrowArray<types::UINT128>(arrow_col.get(), i);
+          std::string uuid_str =
+              sole::rebuild(absl::Uint128High64(val), absl::Uint128Low64(val)).str();
+          col->Append(uuid_str);
+        }
+        block.AppendColumn(mapping.clickhouse_column_name(), col);
+        break;
+      }
+      default:
+        return error::InvalidArgument("Unsupported data type for ClickHouse export: $0",
+                                      types::ToString(mapping.column_type()));
+    }
+  }
+
+  // Insert the block into ClickHouse
+  PX_RETURN_IF_ERROR(ch::Insert(clickhouse_client_.get(), plan_node_->table_name(), block));
+
+  if (rb.eos()) {
+    sent_eos_ = true;
+  }
+
+  return Status::OK();
+}
+
+}  // namespace exec
+}  // namespace carnot
+}  // namespace px

src/carnot/exec/clickhouse_export_sink_node.h

@@ -0,0 +1,55 @@
+/*
+ * 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 <clickhouse/client.h>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "src/carnot/exec/exec_node.h"
+#include "src/carnot/planpb/plan.pb.h"
+#include "src/common/base/base.h"
+#include "src/shared/types/types.h"
+
+namespace px {
+namespace carnot {
+namespace exec {
+
+class ClickHouseExportSinkNode : public SinkNode {
+ public:
+  virtual ~ClickHouseExportSinkNode() = default;
+
+ protected:
+  std::string DebugStringImpl() override;
+  Status InitImpl(const plan::Operator& plan_node) override;
+  Status PrepareImpl(ExecState* exec_state) override;
+  Status OpenImpl(ExecState* exec_state) override;
+  Status CloseImpl(ExecState* exec_state) override;
+  Status ConsumeNextImpl(ExecState* exec_state, const table_store::schema::RowBatch& rb,
+                         size_t parent_index) override;
+
+ private:
+  std::unique_ptr<table_store::schema::RowDescriptor> input_descriptor_;
+  std::unique_ptr<clickhouse::Client> clickhouse_client_;
+  std::unique_ptr<plan::ClickHouseExportSinkOperator> plan_node_;
+};
+
+}  // namespace exec
+}  // namespace carnot
+}  // namespace px