[carnot] Add ClickHouse source and export-sink plan operators

Author: ddelnano

src/carnot/plan/operators.cc

@@ -83,6 +83,10 @@ std::unique_ptr<Operator> Operator::FromProto(const planpb::Operator& pb, int64_
       return CreateOperator<UDTFSourceOperator>(id, pb.udtf_source_op());
     case planpb::EMPTY_SOURCE_OPERATOR:
       return CreateOperator<EmptySourceOperator>(id, pb.empty_source_op());
+    case planpb::CLICKHOUSE_SOURCE_OPERATOR:
+      return CreateOperator<ClickHouseSourceOperator>(id, pb.clickhouse_source_op());
+    case planpb::CLICKHOUSE_EXPORT_SINK_OPERATOR:
+      return CreateOperator<ClickHouseExportSinkOperator>(id, pb.clickhouse_sink_op());
     case planpb::OTEL_EXPORT_SINK_OPERATOR:
       return CreateOperator<OTelExportSinkOperator>(id, pb.otel_sink_op());
     default:
@@ -709,6 +713,67 @@ StatusOr<table_store::schema::Relation> EmptySourceOperator::OutputRelation(
   return r;
 }
 
+/**
+ * ClickHouseSourceOperator implementation.
+ */
+
+std::string ClickHouseSourceOperator::DebugString() const {
+  return absl::Substitute(R"(Op:ClickHouseSource(
+  host=$0
+  port=$1
+  username=$2
+  batch_size=$3
+  start_time=$4
+  end_time=$5
+  timestamp_column=$6
+  partition_column=$7
+)",
+                          pb_.host(), pb_.port(), pb_.username(), pb_.batch_size(),
+                          pb_.start_time(), pb_.end_time(), pb_.timestamp_column(),
+                          pb_.partition_column());
+}
+
+Status ClickHouseSourceOperator::Init(const planpb::ClickHouseSourceOperator& pb) {
+  pb_ = pb;
+  is_initialized_ = true;
+  return Status::OK();
+}
+
+StatusOr<table_store::schema::Relation> ClickHouseSourceOperator::OutputRelation(
+    const table_store::schema::Schema&, const PlanState&,
+    const std::vector<int64_t>& input_ids) const {
+  DCHECK(is_initialized_) << "Not initialized";
+  if (!input_ids.empty()) {
+    return error::InvalidArgument("Source operator cannot have any inputs");
+  }
+  table_store::schema::Relation r;
+  for (int i = 0; i < pb_.column_types_size(); ++i) {
+    r.AddColumn(static_cast<types::DataType>(pb_.column_types(i)), pb_.column_names(i));
+  }
+  return r;
+}
+
+/**
+ * ClickHouse Export Sink Operator Implementation.
+ */
+
+Status ClickHouseExportSinkOperator::Init(const planpb::ClickHouseExportSinkOperator& pb) {
+  pb_ = pb;
+  is_initialized_ = true;
+  return Status::OK();
+}
+
+StatusOr<table_store::schema::Relation> ClickHouseExportSinkOperator::OutputRelation(
+    const table_store::schema::Schema&, const PlanState&, const std::vector<int64_t>&) const {
+  DCHECK(is_initialized_) << "Not initialized";
+  // There are no outputs.
+  return table_store::schema::Relation();
+}
+
+std::string ClickHouseExportSinkOperator::DebugString() const {
+  return absl::Substitute("Op:ClickHouseExportSink(table=$0)", pb_.table_name());
+}
+
 /**
  * OTel Export Sink Operator Implementation.
  */

src/carnot/plan/operators.h

@@ -359,6 +359,69 @@ class EmptySourceOperator : public Operator {
   std::vector<int64_t> column_idxs_;
 };
 
+class ClickHouseSourceOperator : public Operator {
+ public:
+  explicit ClickHouseSourceOperator(int64_t id)
+      : Operator(id, planpb::CLICKHOUSE_SOURCE_OPERATOR) {}
+  ~ClickHouseSourceOperator() override = default;
+
+  StatusOr<table_store::schema::Relation> OutputRelation(
+      const table_store::schema::Schema& schema, const PlanState& state,
+      const std::vector<int64_t>& input_ids) const override;
+  Status Init(const planpb::ClickHouseSourceOperator& pb);
+  std::string DebugString() const override;
+
+  std::string host() const { return pb_.host(); }
+  int32_t port() const { return pb_.port(); }
+  std::string username() const { return pb_.username(); }
+  std::string password() const { return pb_.password(); }
+  std::string database() const { return pb_.database(); }
+  std::string query() const { return pb_.query(); }
+  int32_t batch_size() const { return pb_.batch_size(); }
+  bool streaming() const { return pb_.streaming(); }
+  std::vector<std::string> column_names() const {
+    return std::vector<std::string>(pb_.column_names().begin(), pb_.column_names().end());
+  }
+  std::vector<types::DataType> column_types() const {
+    std::vector<types::DataType> types;
+    types.reserve(pb_.column_types_size());
+    for (const auto& type : pb_.column_types()) {
+      types.push_back(static_cast<types::DataType>(type));
+    }
+    return types;
+  }
+  std::string timestamp_column() const { return pb_.timestamp_column(); }
+  std::string partition_column() const { return pb_.partition_column(); }
+  int64_t start_time() const { return pb_.start_time(); }
+  int64_t end_time() const { return pb_.end_time(); }
+
+ private:
+  planpb::ClickHouseSourceOperator pb_;
+};
+
+class ClickHouseExportSinkOperator : public Operator {
+ public:
+  explicit ClickHouseExportSinkOperator(int64_t id)
+      : Operator(id, planpb::CLICKHOUSE_EXPORT_SINK_OPERATOR) {}
+  ~ClickHouseExportSinkOperator() override = default;
+
+  StatusOr<table_store::schema::Relation> OutputRelation(
+      const table_store::schema::Schema& schema, const PlanState& state,
+      const std::vector<int64_t>& input_ids) const override;
+  Status Init(const planpb::ClickHouseExportSinkOperator& pb);
+  std::string DebugString() const override;
+
+  const planpb::ClickHouseConfig& clickhouse_config() const { return pb_.clickhouse_config(); }
+  const std::string& table_name() const { return pb_.table_name(); }
+  const ::google::protobuf::RepeatedPtrField<planpb::ClickHouseExportSinkOperator::ColumnMapping>&
+  column_mappings() const {
+    return pb_.column_mappings();
+  }
+
+ private:
+  planpb::ClickHouseExportSinkOperator pb_;
+};
+
 class OTelExportSinkOperator : public Operator {
  public:
   explicit OTelExportSinkOperator(int64_t id) : Operator(id, planpb::OTEL_EXPORT_SINK_OPERATOR) {}

src/carnot/plan/plan_fragment.cc

@@ -98,6 +98,13 @@ Status PlanFragmentWalker::CallWalkFn(const Operator& op) {
     case planpb::OperatorType::OTEL_EXPORT_SINK_OPERATOR:
       PX_RETURN_IF_ERROR(CallAs<OTelExportSinkOperator>(on_otel_sink_walk_fn_, op));
       break;
+    case planpb::OperatorType::CLICKHOUSE_SOURCE_OPERATOR:
+      PX_RETURN_IF_ERROR(CallAs<ClickHouseSourceOperator>(on_clickhouse_source_walk_fn_, op));
+      break;
+    case planpb::OperatorType::CLICKHOUSE_EXPORT_SINK_OPERATOR:
+      PX_RETURN_IF_ERROR(
+          CallAs<ClickHouseExportSinkOperator>(on_clickhouse_export_sink_walk_fn_, op));
+      break;
     default:
       LOG(FATAL) << absl::Substitute("Operator does not exist: $0", magic_enum::enum_name(op_type));
       return error::InvalidArgument("Operator does not exist: $0", magic_enum::enum_name(op_type));

src/carnot/plan/plan_fragment.h

@@ -76,6 +76,8 @@ class PlanFragmentWalker {
   using UDTFSourceWalkFn = std::function<Status(const UDTFSourceOperator&)>;
   using EmptySourceWalkFn = std::function<Status(const EmptySourceOperator&)>;
   using OTelSinkWalkFn = std::function<Status(const OTelExportSinkOperator&)>;
+  using ClickHouseSourceWalkFn = std::function<Status(const ClickHouseSourceOperator&)>;
+  using ClickHouseExportSinkWalkFn = std::function<Status(const ClickHouseExportSinkOperator&)>;
 
   /**
    * Register callback for when a memory source operator is encountered.
@@ -181,6 +183,17 @@ class PlanFragmentWalker {
     on_otel_sink_walk_fn_ = fn;
     return *this;
   }
+
+  PlanFragmentWalker& OnClickHouseSource(const ClickHouseSourceWalkFn& fn) {
+    on_clickhouse_source_walk_fn_ = fn;
+    return *this;
+  }
+
+  PlanFragmentWalker& OnClickHouseExportSink(const ClickHouseExportSinkWalkFn& fn) {
+    on_clickhouse_export_sink_walk_fn_ = fn;
+    return *this;
+  }
+
   /**
    * Perform a walk of the plan fragment operators in a topologically-sorted order.
    * @param plan_fragment The plan fragment to walk.
@@ -206,6 +219,8 @@ class PlanFragmentWalker {
   UDTFSourceWalkFn on_udtf_source_walk_fn_;
   EmptySourceWalkFn on_empty_source_walk_fn_;
   OTelSinkWalkFn on_otel_sink_walk_fn_;
+  ClickHouseSourceWalkFn on_clickhouse_source_walk_fn_;
+  ClickHouseExportSinkWalkFn on_clickhouse_export_sink_walk_fn_;
 };
 
 }  // namespace plan