Use IcebergColumnHandle

Author: PingLiuPing

presto-native-execution/presto_cpp/main/connectors/IcebergPrestoToVeloxConnector.cpp

@@ -150,6 +150,18 @@ toVeloxIcebergPartitionSpec(
       spec.specId, fields);
 }
 
+velox::parquet::ParquetFieldId toParquetField(
+    const protocol::iceberg::ColumnIdentity& column) {
+  std::vector<velox::parquet::ParquetFieldId> children;
+  if (!column.children.empty()) {
+    children.reserve(column.children.size());
+    for (const auto& child : column.children) {
+      children.push_back(toParquetField(child));
+    }
+  }
+  return velox::parquet::ParquetFieldId(column.id, children);
+}
+
 } // namespace
 
 std::unique_ptr<velox::connector::ConnectorSplit>
@@ -234,13 +246,13 @@ IcebergPrestoToVeloxConnector::toVeloxColumnHandle(
     columnParseParameters.partitionDateValueFormat = velox::connector::hive::
         HiveColumnHandle::ColumnParseParameters::kDaysSinceEpoch;
   }
-  return std::make_unique<velox::connector::hive::HiveColumnHandle>(
+
+  return std::make_unique<velox::connector::hive::iceberg::IcebergColumnHandle>(
       icebergColumn->columnIdentity.name,
       toHiveColumnType(icebergColumn->columnType),
       type,
-      type,
-      toRequiredSubfields(icebergColumn->requiredSubfields),
-      columnParseParameters);
+      toParquetField(icebergColumn->columnIdentity),
+      toRequiredSubfields(icebergColumn->requiredSubfields));
 }
 
 std::unique_ptr<velox::connector::ConnectorTableHandle>
@@ -324,7 +336,7 @@ IcebergPrestoToVeloxConnector::toVeloxInsertTableHandle(
       createHandle->handle.connectorHandle->_type);
 
   const auto inputColumns =
-      toHiveColumns(icebergOutputTableHandle->inputColumns, typeParser);
+      toIcebergColumns(icebergOutputTableHandle->inputColumns, typeParser);
 
   return std::make_unique<
       velox::connector::hive::iceberg::IcebergInsertTableHandle>(
@@ -354,7 +366,7 @@ IcebergPrestoToVeloxConnector::toVeloxInsertTableHandle(
       insertHandle->handle.connectorHandle->_type);
 
   const auto inputColumns =
-      toHiveColumns(icebergInsertTableHandle->inputColumns, typeParser);
+      toIcebergColumns(icebergInsertTableHandle->inputColumns, typeParser);
 
   return std::make_unique<
       velox::connector::hive::iceberg::IcebergInsertTableHandle>(
@@ -370,18 +382,20 @@ IcebergPrestoToVeloxConnector::toVeloxInsertTableHandle(
           toFileCompressionKind(icebergInsertTableHandle->compressionCodec)));
 }
 
-std::vector<velox::connector::hive::HiveColumnHandlePtr>
-IcebergPrestoToVeloxConnector::toHiveColumns(
+std::vector<velox::connector::hive::iceberg::IcebergColumnHandlePtr>
+IcebergPrestoToVeloxConnector::toIcebergColumns(
     const protocol::List<protocol::iceberg::IcebergColumnHandle>& inputColumns,
     const TypeParser& typeParser) const {
-  std::vector<velox::connector::hive::HiveColumnHandlePtr> hiveColumns;
-  hiveColumns.reserve(inputColumns.size());
+  std::vector<velox::connector::hive::iceberg::IcebergColumnHandlePtr>
+      icebergColumns;
+  icebergColumns.reserve(inputColumns.size());
   for (const auto& columnHandle : inputColumns) {
-    hiveColumns.emplace_back(
-        std::dynamic_pointer_cast<velox::connector::hive::HiveColumnHandle>(
+    icebergColumns.emplace_back(
+        std::dynamic_pointer_cast<
+            velox::connector::hive::iceberg::IcebergColumnHandle>(
             std::shared_ptr(toVeloxColumnHandle(&columnHandle, typeParser))));
   }
-  return hiveColumns;
+  return icebergColumns;
 }
 
 } // namespace facebook::presto

presto-native-execution/presto_cpp/main/connectors/IcebergPrestoToVeloxConnector.h

@@ -17,6 +17,8 @@
 #include "presto_cpp/main/connectors/PrestoToVeloxConnector.h"
 #include "presto_cpp/presto_protocol/connector/iceberg/presto_protocol_iceberg.h"
 
+#include "velox/connectors/hive/iceberg/IcebergColumnHandle.h"
+
 namespace facebook::presto {
 
 class IcebergPrestoToVeloxConnector final : public PrestoToVeloxConnector {
@@ -52,7 +54,8 @@ class IcebergPrestoToVeloxConnector final : public PrestoToVeloxConnector {
       const TypeParser& typeParser) const final;
 
  private:
-  std::vector<velox::connector::hive::HiveColumnHandlePtr> toHiveColumns(
+  std::vector<velox::connector::hive::iceberg::IcebergColumnHandlePtr>
+  toIcebergColumns(
       const protocol::List<protocol::iceberg::IcebergColumnHandle>&
           inputColumns,
       const TypeParser& typeParser) const;

presto-native-execution/presto_cpp/main/types/tests/PrestoToVeloxConnectorTest.cpp

@@ -21,6 +21,7 @@
 #include "velox/common/base/tests/GTestUtils.h"
 #include "velox/connectors/hive/HiveConnector.h"
 #include "velox/connectors/hive/TableHandle.h"
+#include "velox/connectors/hive/iceberg/IcebergColumnHandle.h"
 
 using namespace facebook::presto;
 using namespace facebook::velox;
@@ -185,3 +186,110 @@ TEST_F(PrestoToVeloxConnectorTest, hiveLowercasesColumnNames) {
   EXPECT_EQ(dataColumnsType->nameOf(0), "mixedcasecol1");
   EXPECT_EQ(dataColumnsType->nameOf(1), "uppercasecol2");
 }
+
+namespace {
+
+protocol::iceberg::IcebergColumnHandle createIcebergColumnHandle(
+    const std::string& name,
+    int32_t fieldId,
+    const std::string& type,
+    protocol::iceberg::TypeCategory typeCategory =
+        protocol::iceberg::TypeCategory::PRIMITIVE,
+    const std::vector<protocol::iceberg::ColumnIdentity>& children = {}) {
+  protocol::iceberg::IcebergColumnHandle column;
+  column.columnIdentity.name = name;
+  column.columnIdentity.id = fieldId;
+  column.columnIdentity.typeCategory = typeCategory;
+  column.columnIdentity.children = children;
+  column.type = type;
+  column.columnType = protocol::hive::ColumnType::REGULAR;
+  return column;
+}
+
+} // namespace
+
+TEST_F(PrestoToVeloxConnectorTest, icebergColumnHandleSimple) {
+  auto icebergColumn = createIcebergColumnHandle("col1", 1, "integer");
+
+  IcebergPrestoToVeloxConnector icebergConnector("iceberg");
+  auto handle =
+      icebergConnector.toVeloxColumnHandle(&icebergColumn, *typeParser_);
+  auto* icebergHandle =
+      dynamic_cast<connector::hive::iceberg::IcebergColumnHandle*>(
+          handle.get());
+  ASSERT_NE(icebergHandle, nullptr);
+
+  EXPECT_EQ(icebergHandle->name(), "col1");
+  EXPECT_EQ(icebergHandle->dataType()->kind(), TypeKind::INTEGER);
+  EXPECT_EQ(icebergHandle->field().fieldId, 1);
+  EXPECT_TRUE(icebergHandle->field().children.empty());
+}
+
+TEST_F(PrestoToVeloxConnectorTest, icebergColumnHandleNested) {
+  protocol::iceberg::ColumnIdentity child1;
+  child1.name = "child1";
+  child1.id = 2;
+  child1.typeCategory = protocol::iceberg::TypeCategory::PRIMITIVE;
+
+  protocol::iceberg::ColumnIdentity child2;
+  child2.name = "child2";
+  child2.id = 3;
+  child2.typeCategory = protocol::iceberg::TypeCategory::PRIMITIVE;
+
+  auto icebergColumn = createIcebergColumnHandle(
+      "struct_col",
+      1,
+      "row(child1 integer, child2 varchar)",
+      protocol::iceberg::TypeCategory::STRUCT,
+      {child1, child2});
+
+  IcebergPrestoToVeloxConnector icebergConnector("iceberg");
+  auto handle =
+      icebergConnector.toVeloxColumnHandle(&icebergColumn, *typeParser_);
+  auto* icebergHandle =
+      dynamic_cast<connector::hive::iceberg::IcebergColumnHandle*>(
+          handle.get());
+  ASSERT_NE(icebergHandle, nullptr);
+
+  EXPECT_EQ(icebergHandle->name(), "struct_col");
+  EXPECT_EQ(icebergHandle->dataType()->kind(), TypeKind::ROW);
+  EXPECT_EQ(icebergHandle->field().fieldId, 1);
+  ASSERT_EQ(icebergHandle->field().children.size(), 2);
+  EXPECT_EQ(icebergHandle->field().children[0].fieldId, 2);
+  EXPECT_EQ(icebergHandle->field().children[1].fieldId, 3);
+}
+
+TEST_F(PrestoToVeloxConnectorTest, icebergColumnHandleDeeplyNested) {
+  protocol::iceberg::ColumnIdentity inner;
+  inner.name = "inner";
+  inner.id = 3;
+  inner.typeCategory = protocol::iceberg::TypeCategory::PRIMITIVE;
+
+  protocol::iceberg::ColumnIdentity middle;
+  middle.name = "middle";
+  middle.id = 2;
+  middle.typeCategory = protocol::iceberg::TypeCategory::STRUCT;
+  middle.children = {inner};
+
+  auto icebergColumn = createIcebergColumnHandle(
+      "outer",
+      1,
+      "row(middle row(inner bigint))",
+      protocol::iceberg::TypeCategory::STRUCT,
+      {middle});
+
+  IcebergPrestoToVeloxConnector icebergConnector("iceberg");
+  auto handle =
+      icebergConnector.toVeloxColumnHandle(&icebergColumn, *typeParser_);
+  auto* icebergHandle =
+      dynamic_cast<connector::hive::iceberg::IcebergColumnHandle*>(
+          handle.get());
+  ASSERT_NE(icebergHandle, nullptr);
+
+  EXPECT_EQ(icebergHandle->name(), "outer");
+  EXPECT_EQ(icebergHandle->field().fieldId, 1);
+  ASSERT_EQ(icebergHandle->field().children.size(), 1);
+  EXPECT_EQ(icebergHandle->field().children[0].fieldId, 2);
+  ASSERT_EQ(icebergHandle->field().children[0].children.size(), 1);
+  EXPECT_EQ(icebergHandle->field().children[0].children[0].fieldId, 3);
+}