Add serialization logic to ArrowFederationColumnHandle

Author: pdabre12

presto-flight-shim/etc/catalog/postgresql_test.properties

@@ -0,0 +1,4 @@
+connector.name=postgresql
+connection-url=jdbc:postgresql://localhost:5432/testdb
+connection-user=presto
+connection-password=presto

presto-flight-shim/etc/flightshim.properties

@@ -1,9 +1,9 @@
 flight-shim.server=localhost
 flight-shim.server.port=9999
-flight-shim.server-ssl-certificate-file=src/test/resources/server.crt
-flight-shim.server-ssl-key-file=src/test/resources/server.key
+#flight-shim.server-ssl-certificate-file=src/test/resources/server.crt
+flight-shim.server-ssl-enabled=false
 
 plugin.bundles=\
-  ../presto-oracle/pom, \
+  ../presto-oracle/pom.xml, \
   ../presto-postgresql/pom.xml, \
   ../presto-mysql/pom.xml

presto-main/etc/catalog/postgresql.properties

@@ -0,0 +1,4 @@
+connector.name=postgresql
+connection-url=jdbc:postgresql://localhost:5432/testdb
+connection-user=presto
+connection-password=presto

presto-main/etc/catalog/postgresql_test.properties

@@ -1,2 +1,3 @@
 connector.name=tpcds
 tpcds.splits-per-node=4
+tpcds.use-varchar-type=true

presto-main/etc/catalog/tpcds.properties

@@ -7,7 +7,7 @@
 
 # sample nodeId to provide consistency across test runs
 node.id=ffffffff-ffff-ffff-ffff-ffffffffffff
-node.environment=test
+node.environment=testing
 http-server.http.port=8080
 
 discovery-server.enabled=true
@@ -55,4 +55,5 @@ plugin.bundles=\
   ../presto-sql-helpers/presto-sql-invoked-functions-plugin/pom.xml
 
 presto.version=testversion
-node-scheduler.include-coordinator=true
+node-scheduler.include-coordinator=false
+native-execution-enabled=true

presto-main/etc/config.properties

@@ -35,10 +35,25 @@ velox::connector::ColumnHandleMap toArrowFlightColumnHandleMap(
   for (const auto& [name, handle] : columnHandles) {
     arrowFlightColumnHandles[name] =
         std::make_shared<facebook::presto::ArrowFlightColumnHandle>(
-            "arrow-flight");
+            handle->name());
   }
   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);
+  }
+}
 } // namespace
 
 ArrowFederationDataSource::ArrowFederationDataSource(
@@ -54,7 +69,8 @@ ArrowFederationDataSource::ArrowFederationDataSource(
           authenticator,
           connectorQueryCtx,
           flightConfig,
-          clientOpts) {}
+          clientOpts),
+      columnHandles_(columnHandles) {}
 
 void ArrowFederationDataSource::addSplit(
     std::shared_ptr<ConnectorSplit> split) {
@@ -64,16 +80,17 @@ void ArrowFederationDataSource::addSplit(
       "ArrowFederationDataSource received wrong type of split");
 
   nlohmann::json request;
-  request["split"] = federationSplit->splitBytes_;
-  request["columns"] = columnMapping_;
+  request["connectorId"] = federationSplit->connectorId;
+  request["splitBytes"] = federationSplit->splitBytes_;
+  to_json(request["columnHandlesBytes"], columnHandles_);
 
   arrow::flight::FlightEndpoint flightEndpoint{request.dump()};
 
   std::string flightEndpointBytes;
   AFC_ASSIGN_OR_RAISE(flightEndpointBytes, flightEndpoint.SerializeToString());
 
   auto flightSplit = std::make_shared<ArrowFlightSplit>(
-      federationSplit->connectorId, flightEndpointBytes);
+      federationSplit->connectorId, folly::base64Encode(flightEndpointBytes));
 
   ArrowFlightDataSource::addSplit(flightSplit);
 }

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

@@ -26,6 +26,13 @@ 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
@@ -40,15 +47,65 @@ struct ArrowFederationSplit : public velox::connector::ConnectorSplit {
 
 class ArrowFederationColumnHandle : public velox::connector::ColumnHandle {
  public:
-  explicit ArrowFederationColumnHandle(const std::string& columnName)
-      : columnName_(columnName) {}
+  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& name() const {
     return columnName_;
   }
 
+  const std::string& connectorId() const {
+    return connectorId_;
+  }
+
+  const std::string& columnType() const {
+    return columnType_;
+  }
+
+  bool nullable() const {
+    return nullable_;
+  }
+
+  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_;
 };
 
 class ArrowFederationTableHandle
@@ -81,6 +138,9 @@ class ArrowFederationDataSource : public ArrowFlightDataSource {
   std::optional<velox::RowVectorPtr> next(
       uint64_t size,
       velox::ContinueFuture& /* unused */) override;
+
+ private:
+  const velox::connector::ColumnHandleMap columnHandles_;
 };
 
 class ArrowFederationConnector : public ArrowFlightConnector {

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

@@ -39,8 +39,18 @@ ArrowFederationPrestoToVeloxConnector::toVeloxColumnHandle(
       const protocol::arrow_federation::ArrowFederationColumnHandle*>(column);
   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->columnName,
+      arrowColumn->connectorId,
+      arrowColumn->columnType,
+      arrowColumn->nullable,
+      arrowJdbcTypeHandle);
 }
 
 std::unique_ptr<velox::connector::ConnectorTableHandle>

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

@@ -89,21 +89,90 @@ void from_json(const json& j, ArrowFederationTableHandle& p) {
 namespace facebook::presto::protocol::arrow_federation {
 
 struct ArrowFederationColumnHandle : public ColumnHandle {
+  String connectorId = {};
   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, "columnName", p.columnName, "ArrowFederationColumnHandle", "ColumnName", "columnName");
-
+      j,
+      "connectorId",
+      p.connectorId,
+      "ArrowFederationColumnHandle",
+      "ConnectorId",
+      "connectorId");
+  to_json_key(
+      j,
+      "columnName",
+      p.columnName,
+      "ArrowFederationColumnHandle",
+      "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");
 }
 
 void from_json(const json& j, ArrowFederationColumnHandle& p) {
   p._type = j["@type"];
   from_json_key(
-      j, "columnName", p.columnName, "ArrowFederationColumnHandle", "ColumnName", "columnName");
+      j,
+      "connectorId",
+      p.connectorId,
+      "ArrowFederationColumnHandle",
+      "ConnectorId",
+      "connectorId");
+  from_json_key(
+      j,
+      "columnName",
+      p.columnName,
+      "ArrowFederationColumnHandle",
+      "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");
 }
 } // namespace facebook::presto::protocol::arrow_federation