Cleanup ArrowFederationColumnHandle protocol

Author: pdabre12

presto-native-execution/presto_cpp/main/connectors/arrow_flight/arrow_federation/ArrowFederationConnector.cpp

@@ -40,19 +40,20 @@ velox::connector::ColumnHandleMap toArrowFlightColumnHandleMap(
   return arrowFlightColumnHandles;
 }
 
-// Json conversion helpers for ColumnHandleMap
-void to_json(nlohmann::json& j, const ColumnHandleMap& map) {
-  j = nlohmann::json::array();
-  for (const auto& [name, handle] : map) {
-    if (!handle) {
-      continue; // skip null handles
-    }
-    auto handleStr = folly::toJson(handle->serialize());
-
-    // Base64 encode
-    auto encoded = folly::base64Encode(handleStr);
-    j.push_back(encoded);
+const std::vector<std::string> toArrowFederationColumnHandleList(
+    const velox::connector::ColumnHandleMap& columnHandles) {
+  std::vector<std::string> arrowFederationColumnHandles;
+
+  for (const auto& [name, handle] : columnHandles) {
+    auto federationColumnHandle =
+        std::dynamic_pointer_cast<const ArrowFederationColumnHandle>(handle);
+    VELOX_CHECK(
+        handle,
+        "ArrowFederationDataSource received wrong type of column handle");
+    arrowFederationColumnHandles.push_back(
+        federationColumnHandle->columnHandleBytes());
   }
+  return arrowFederationColumnHandles;
 }
 } // namespace
 
@@ -82,7 +83,8 @@ void ArrowFederationDataSource::addSplit(
   nlohmann::json request;
   request["connectorId"] = federationSplit->connectorId;
   request["splitBytes"] = federationSplit->splitBytes_;
-  to_json(request["columnHandlesBytes"], columnHandles_);
+  request["columnHandlesBytes"] =
+      toArrowFederationColumnHandleList(columnHandles_);
 
   arrow::flight::FlightEndpoint flightEndpoint{request.dump()};
 

presto-native-execution/presto_cpp/main/connectors/arrow_flight/arrow_federation/ArrowFederationConnector.h

@@ -26,13 +26,6 @@ class Location;
 
 namespace facebook::presto {
 
-struct ArrowFederationTypeHandle {
-  int jdbcType;
-  std::string jdbcTypeName;
-  int columnSize;
-  int decimalDigits;
-};
-
 struct ArrowFederationSplit : public velox::connector::ConnectorSplit {
   /// @param connectorId
   /// @param splitBytes Base64 Serialized Split for Arrow Federation Flight
@@ -48,64 +41,20 @@ struct ArrowFederationSplit : public velox::connector::ConnectorSplit {
 class ArrowFederationColumnHandle : public velox::connector::ColumnHandle {
  public:
   explicit ArrowFederationColumnHandle(
-      const std::string& columnName,
-      const std::string& connectorId,
-      const std::string& columnType,
-      bool nullable,
-      const ArrowFederationTypeHandle& arrowFederationTypeHandle)
-      : columnName_(columnName),
-        connectorId_(connectorId),
-        columnType_(columnType),
-        nullable_(nullable),
-        arrowFederationTypeHandle_(arrowFederationTypeHandle) {}
+      const std::string& columnHandleBytes,
+      const std::string& columnName)
+      : columnName_(columnName), columnHandleBytes_(columnHandleBytes) {}
 
   const std::string& name() const {
     return columnName_;
   }
 
-  const std::string& connectorId() const {
-    return connectorId_;
-  }
-
-  const std::string& columnType() const {
-    return columnType_;
-  }
-
-  bool nullable() const {
-    return nullable_;
+  const std::string& columnHandleBytes() const {
+    return columnHandleBytes_;
   }
 
-  const ArrowFederationTypeHandle& arrowFederationTypeHandle() const {
-    return arrowFederationTypeHandle_;
-  }
-
-  folly::dynamic serializeArrowFederationTypeHandle(
-      const ArrowFederationTypeHandle& handle) const {
-    folly::dynamic obj = folly::dynamic::object;
-    obj["jdbcType"] = handle.jdbcType;
-    obj["jdbcTypeName"] = handle.jdbcTypeName;
-    obj["columnSize"] = handle.columnSize;
-    obj["decimalDigits"] = handle.decimalDigits;
-    return obj;
-  }
-
-  folly::dynamic serialize() const {
-    folly::dynamic obj = folly::dynamic::object;
-    obj["connectorId"] = connectorId();
-    obj["columnName"] = name();
-    obj["columnType"] = columnType();
-    obj["nullable"] = nullable();
-    obj["jdbcTypeHandle"] =
-        serializeArrowFederationTypeHandle(arrowFederationTypeHandle());
-    return obj;
-  }
-
- private:
-  std::string columnName_;
-  std::string connectorId_;
-  std::string columnType_;
-  bool nullable_;
-  const ArrowFederationTypeHandle arrowFederationTypeHandle_;
+  const std::string columnName_;
+  const std::string columnHandleBytes_;
 };
 
 class ArrowFederationTableHandle

presto-native-execution/presto_cpp/main/connectors/arrow_flight/arrow_federation/ArrowFederationPrestoToVeloxConnector.cpp

@@ -40,17 +40,8 @@ ArrowFederationPrestoToVeloxConnector::toVeloxColumnHandle(
   VELOX_CHECK_NOT_NULL(
       arrowColumn, "Unexpected column handle type {}", column->_type);
 
-  auto arrowJdbcTypeHandle = presto::ArrowFederationTypeHandle{
-      arrowColumn->jdbcTypeHandle.jdbcType,
-      arrowColumn->jdbcTypeHandle.jdbcTypeName,
-      arrowColumn->jdbcTypeHandle.columnSize,
-      arrowColumn->jdbcTypeHandle.decimalDigits};
   return std::make_unique<presto::ArrowFederationColumnHandle>(
-      arrowColumn->columnName,
-      arrowColumn->connectorId,
-      arrowColumn->columnType,
-      arrowColumn->nullable,
-      arrowJdbcTypeHandle);
+      arrowColumn->columnHandleBytes, arrowColumn->columnName);
 }
 
 std::unique_ptr<velox::connector::ConnectorTableHandle>

presto-native-execution/presto_cpp/presto_protocol/connector/arrow_federation/presto_protocol-json-hpp.mustache

@@ -89,90 +89,33 @@ void from_json(const json& j, ArrowFederationTableHandle& p) {
 namespace facebook::presto::protocol::arrow_federation {
 
 struct ArrowFederationColumnHandle : public ColumnHandle {
-  String connectorId = {};
+  String columnHandleBytes = {};
   String columnName = {};
-  String columnType = {};
-  boolean nullable = {};
-  JdbcTypeHandle jdbcTypeHandle = {};
 };
 
 void to_json(json& j, const ArrowFederationColumnHandle& p) {
-  j = json::object();
   j["@type"] = "arrow-federation";
-  to_json_key(
-      j,
-      "connectorId",
-      p.connectorId,
-      "ArrowFederationColumnHandle",
-      "ConnectorId",
-      "connectorId");
   to_json_key(
       j,
       "columnName",
       p.columnName,
-      "ArrowFederationColumnHandle",
+      "ArrowFederationTableHandle",
       "ColumnName",
       "columnName");
-  to_json_key(
-      j,
-      "columnType",
-      p.columnType,
-      "ArrowFederationColumnHandle",
-      "ColumnType",
-      "columnType");
-    to_json_key(
-      j,
-      "nullable",
-      p.nullable,
-      "ArrowFederationColumnHandle",
-      "Nullable",
-      "nullable");
-    to_json_key(
-      j,
-      "jdbcTypeHandle",
-      p.jdbcTypeHandle,
-      "ArrowFederationColumnHandle",
-      "JdbcTypeHandle",
-      "jdbcTypeHandle");
+  auto columnHandle = folly::base64Decode(p.columnHandleBytes);
+  j = json::parse(columnHandle);
 }
 
 void from_json(const json& j, ArrowFederationColumnHandle& p) {
   p._type = j["@type"];
-  from_json_key(
-      j,
-      "connectorId",
-      p.connectorId,
-      "ArrowFederationColumnHandle",
-      "ConnectorId",
-      "connectorId");
   from_json_key(
       j,
       "columnName",
       p.columnName,
-      "ArrowFederationColumnHandle",
+      "ArrowFederationTableHandle",
       "ColumnName",
       "columnName");
-  from_json_key(
-      j,
-      "columnType",
-      p.columnType,
-      "ArrowFederationColumnHandle",
-      "ColumnType",
-      "columnType");
-    from_json_key(
-      j,
-      "nullable",
-      p.nullable,
-      "ArrowFederationColumnHandle",
-      "Nullable",
-      "nullable");
-    from_json_key(
-      j,
-      "jdbcTypeHandle",
-      p.jdbcTypeHandle,
-      "ArrowFederationColumnHandle",
-      "JdbcTypeHandle",
-      "jdbcTypeHandle");
+  p.columnHandleBytes = folly::base64Encode(j.dump());
 }
 } // namespace facebook::presto::protocol::arrow_federation