feat: Add support for evaluating Iceberg partition transforms

velox/connectors/hive/HiveConnector.cpp

@@ -21,6 +21,7 @@
 #include "velox/connectors/hive/HiveDataSource.h"
 #include "velox/connectors/hive/HivePartitionFunction.h"
 #include "velox/connectors/hive/iceberg/IcebergDataSink.h"
+#include "velox/functions/iceberg/Register.h"
 
 #include <boost/lexical_cast.hpp>
 #include <memory>
@@ -42,6 +43,7 @@ HiveConnector::HiveConnector(
               : nullptr,
           std::make_unique<FileHandleGenerator>(hiveConfig_->config())),
       ioExecutor_(ioExecutor) {
+  iceberg::registerIcebergFunctions();
   if (hiveConfig_->isFileHandleCacheEnabled()) {
     LOG(INFO) << "Hive connector " << connectorId()
               << " created with maximum of "

velox/connectors/hive/iceberg/CMakeLists.txt

@@ -19,6 +19,15 @@ velox_add_library(
   IcebergSplitReader.cpp
   PartitionSpec.cpp
   PositionalDeleteFileReader.cpp
+  TransformEvaluator.cpp
+  TransformExprBuilder.cpp
+)
+
+velox_link_libraries(
+  velox_hive_iceberg_splitreader
+  velox_connector
+  velox_functions_iceberg
+  Folly::folly
 )
 
 velox_link_libraries(velox_hive_iceberg_splitreader velox_connector Folly::folly)

velox/connectors/hive/iceberg/PartitionSpec.cpp

@@ -16,6 +16,7 @@
 
 #include "velox/connectors/hive/iceberg/PartitionSpec.h"
 
+#include "velox/functions/iceberg/Register.h"
 #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h"
 
 namespace facebook::velox::connector::hive::iceberg {
@@ -102,6 +103,14 @@ VELOX_DEFINE_ENUM_NAME(TransformType, transformTypeNames);
 
 VELOX_DEFINE_ENUM_NAME(TransformCategory, transformCategoryNames);
 
+void registerIcebergFunctions() {
+  static std::once_flag registerFlag;
+
+  std::call_once(registerFlag, []() {
+    functions::iceberg::registerFunctions(kIcebergFunctionPrefix);
+  });
+}
+
 void IcebergPartitionSpec::checkCompatibility() const {
   folly::F14FastMap<std::string_view, std::vector<TransformType>>
       columnTransforms;

velox/connectors/hive/iceberg/PartitionSpec.h

@@ -20,6 +20,25 @@
 
 namespace facebook::velox::connector::hive::iceberg {
 
+inline constexpr char const* kIcebergFunctionPrefix{"iceberg_"};
+
+inline const std::string kBucketFunction =
+    std::string(kIcebergFunctionPrefix) + "bucket";
+inline const std::string kTruncateFunction =
+    std::string(kIcebergFunctionPrefix) + "truncate";
+inline const std::string kYearFunction =
+    std::string(kIcebergFunctionPrefix) + "years";
+inline const std::string kMonthFunction =
+    std::string(kIcebergFunctionPrefix) + "months";
+inline const std::string kDayFunction =
+    std::string(kIcebergFunctionPrefix) + "days";
+inline const std::string kHourFunction =
+    std::string(kIcebergFunctionPrefix) + "hours";
+
+/// Registers Iceberg partition transform functions with prefix
+/// kIcebergFunctionPrefix.
+void registerIcebergFunctions();
+
 /// Partition transform types.
 /// Defines how source column values are converted into partition keys.
 /// See https://iceberg.apache.org/spec/#partition-transforms.
@@ -113,6 +132,7 @@ struct IcebergPartitionSpec {
         case TransformType::kTruncate:
           return type;
       }
+      VELOX_UNREACHABLE("Unknown transform type");
     }
   };
 

velox/connectors/hive/iceberg/TransformEvaluator.cpp

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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.
+ */
+
+#include "velox/connectors/hive/iceberg/TransformEvaluator.h"
+
+#include "velox/expression/Expr.h"
+
+namespace facebook::velox::connector::hive::iceberg {
+
+TransformEvaluator::TransformEvaluator(
+    const std::vector<core::TypedExprPtr>& expressions,
+    const ConnectorQueryCtx* connectorQueryCtx)
+    : connectorQueryCtx_(connectorQueryCtx) {
+  VELOX_CHECK_NOT_NULL(connectorQueryCtx_);
+  exprSet_ = connectorQueryCtx_->expressionEvaluator()->compile(expressions);
+  VELOX_CHECK_NOT_NULL(exprSet_);
+}
+
+std::vector<VectorPtr> TransformEvaluator::evaluate(
+    const RowVectorPtr& input) const {
+  const auto numRows = input->size();
+  const auto numExpressions = exprSet_->exprs().size();
+
+  std::vector<VectorPtr> results(numExpressions);
+  SelectivityVector rows(numRows);
+
+  // Evaluate all expressions in one pass.
+  connectorQueryCtx_->expressionEvaluator()->evaluate(
+      exprSet_.get(), rows, *input, results);
+
+  return results;
+}
+
+} // namespace facebook::velox::connector::hive::iceberg

velox/connectors/hive/iceberg/TransformEvaluator.h

@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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.
+ */
+#pragma once
+
+#include "velox/connectors/Connector.h"
+#include "velox/core/QueryCtx.h"
+#include "velox/expression/Expr.h"
+
+namespace facebook::velox::connector::hive::iceberg {
+
+/// Evaluates multiple expressions efficiently using batch evaluation.
+/// Expressions are compiled once in the constructor and reused across multiple
+/// input batches.
+class TransformEvaluator {
+ public:
+  /// Creates an evaluator with the given expressions and connector query
+  /// context. Compiles the expressions once for reuse across multiple
+  /// evaluations.
+  ///
+  /// @param expressions Vector of typed expressions to evaluate. These are
+  /// typically built using TransformExprBuilder::toExpressions() for Iceberg
+  /// partition transforms, but can be any valid Velox expressions. The
+  /// expressions are compiled once during construction.
+  /// @param connectorQueryCtx Connector query context providing access to the
+  /// expression evaluator (for compilation and evaluation) and memory pool.
+  /// Must remain valid for the lifetime of this TransformEvaluator.
+  TransformEvaluator(
+      const std::vector<core::TypedExprPtr>& expressions,
+      const ConnectorQueryCtx* connectorQueryCtx);
+
+  /// Evaluates all expressions on the input data in a single pass.
+  /// Uses the pre-compiled ExprSet from the constructor for efficiency.
+  ///
+  /// The input RowType must match the RowType used when building the
+  /// expressions (passed to TransformExprBuilder::toExpressions). The column
+  /// positions, names and types must align. Create new TransformEvaluator for
+  /// input that has different RowType with the one when building the
+  /// expressions.
+  ///
+  /// @param input Input row vector containing the source data. Must have the
+  /// same RowType (column positions, names and types) as used when building the
+  /// expressions in the constructor.
+  /// @return Vector of result columns, one for each expression, in the same
+  /// order as the expressions provided to the constructor.
+  std::vector<VectorPtr> evaluate(const RowVectorPtr& input) const;
+
+ private:
+  const ConnectorQueryCtx* connectorQueryCtx_;
+  std::unique_ptr<exec::ExprSet> exprSet_;
+};
+
+} // namespace facebook::velox::connector::hive::iceberg

velox/connectors/hive/iceberg/TransformExprBuilder.cpp

@@ -0,0 +1,106 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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.
+ */
+#include "velox/connectors/hive/iceberg/TransformExprBuilder.h"
+#include "velox/core/Expressions.h"
+
+namespace facebook::velox::connector::hive::iceberg {
+
+namespace {
+
+/// Converts a single partition field to a typed expression.
+///
+/// Builds an expression tree for one partition transform. Identity transforms
+/// become FieldAccessTypedExpr, while other transforms (bucket, truncate,
+/// year, month, day, hour) become CallTypedExpr with appropriate function
+/// names and parameters.
+///
+/// @param field Partition field containing transform type, source column
+/// type, and optional parameter (e.g., bucket count, truncate width).
+/// @param inputFieldName Name of the source column in the input RowVector.
+/// @return Typed expression representing the transform.
+core::TypedExprPtr toExpression(
+    const IcebergPartitionSpec::Field& field,
+    const std::string& inputFieldName) {
+  // For identity transform, just return a field access expression.
+  if (field.transformType == TransformType::kIdentity) {
+    return std::make_shared<core::FieldAccessTypedExpr>(
+        field.type, inputFieldName);
+  }
+
+  // For other transforms, build a CallTypedExpr with the appropriate function.
+  std::string functionName;
+  switch (field.transformType) {
+    case TransformType::kBucket:
+      functionName = kBucketFunction;
+      break;
+    case TransformType::kTruncate:
+      functionName = kTruncateFunction;
+      break;
+    case TransformType::kYear:
+      functionName = kYearFunction;
+      break;
+    case TransformType::kMonth:
+      functionName = kMonthFunction;
+      break;
+    case TransformType::kDay:
+      functionName = kDayFunction;
+      break;
+    case TransformType::kHour:
+      functionName = kHourFunction;
+      break;
+    case TransformType::kIdentity:
+      break;
+  }
+
+  // Build the expression arguments.
+  std::vector<core::TypedExprPtr> exprArgs;
+  if (field.parameter.has_value()) {
+    exprArgs.emplace_back(
+        std::make_shared<core::ConstantTypedExpr>(
+            INTEGER(), Variant(field.parameter.value())));
+  }
+  exprArgs.emplace_back(
+      std::make_shared<core::FieldAccessTypedExpr>(field.type, inputFieldName));
+
+  return std::make_shared<core::CallTypedExpr>(
+      field.resultType(), std::move(exprArgs), functionName);
+}
+
+} // namespace
+
+std::vector<core::TypedExprPtr> TransformExprBuilder::toExpressions(
+    const IcebergPartitionSpecPtr& partitionSpec,
+    const std::vector<column_index_t>& partitionChannels,
+    const RowTypePtr& inputType) {
+  VELOX_CHECK_EQ(
+      partitionSpec->fields.size(),
+      partitionChannels.size(),
+      "Number of partition fields must match number of partition channels");
+
+  const auto numTransforms = partitionChannels.size();
+  std::vector<core::TypedExprPtr> transformExprs;
+  transformExprs.reserve(numTransforms);
+
+  for (auto i = 0; i < numTransforms; i++) {
+    const auto channel = partitionChannels[i];
+    transformExprs.emplace_back(
+        toExpression(partitionSpec->fields.at(i), inputType->nameOf(channel)));
+  }
+
+  return transformExprs;
+}
+
+} // namespace facebook::velox::connector::hive::iceberg

velox/connectors/hive/iceberg/TransformExprBuilder.h

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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.
+ */
+
+#pragma once
+
+#include "velox/connectors/hive/iceberg/PartitionSpec.h"
+#include "velox/expression/Expr.h"
+
+namespace facebook::velox::connector::hive::iceberg {
+
+/// Converts Iceberg partition specification to Velox expressions.
+class TransformExprBuilder {
+ public:
+  /// Converts partition specification to a list of typed expressions.
+  ///
+  /// @param partitionSpec Iceberg partition specification containing transform
+  /// definitions for each partition field.
+  /// @param partitionChannels Column indices (0-based) in the input RowVector
+  /// that correspond to each partition field. Must have the same size as
+  /// partitionSpec->fields. Provides the positional mapping from partition spec
+  /// fields to input RowVector columns.
+  /// @param inputType The row type of the input data. This is necessary for
+  /// building expressions because the column names in partitionSpec reference
+  /// table schema names, which might not match the column names in inputType
+  /// (e.g., inputType may use generated names like c0, c1, c2). The
+  /// FieldAccessTypedExpr must be built using the actual column names from
+  /// inputType that will be present at runtime. The partitionChannels provide
+  /// the positional mapping to locate the correct columns.
+  /// @return Vector of typed expressions, one for each partition field.
+  static std::vector<core::TypedExprPtr> toExpressions(
+      const IcebergPartitionSpecPtr& partitionSpec,
+      const std::vector<column_index_t>& partitionChannels,
+      const RowTypePtr& inputType);
+};
+
+} // namespace facebook::velox::connector::hive::iceberg

velox/connectors/hive/iceberg/tests/CMakeLists.txt

@@ -63,6 +63,7 @@ if(NOT VELOX_DISABLE_GOOGLETEST)
     IcebergTestBase.cpp
     Main.cpp
     PartitionSpecTest.cpp
+    TransformTest.cpp
   )
 
   add_test(velox_hive_iceberg_insert_test velox_hive_iceberg_insert_test)