feat: Add support for evaluating Iceberg partition transforms (facebookincubator#15477)

Summary:
Introduce infrastructure for evaluating Iceberg partition transforms.
- TransformExprBuilder converts Iceberg partition specifications into Velox expressions
- TransformEvaluator evaluates multiple transform expressions in a single pass using compiled ExprSet.

Pull Request resolved: facebookincubator#15477

Reviewed By: mbasmanova

Differential Revision: D86893852

Pulled By: Yuhta

fbshipit-source-id: 3f839a3a8ea6eeb7e152fbd960e3023d208e122d

Author: PingLiuPing

velox/connectors/hive/iceberg/CMakeLists.txt

@@ -20,6 +20,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/IcebergDataSink.cpp

@@ -19,6 +19,14 @@
 
 namespace facebook::velox::connector::hive::iceberg {
 
+void registerIcebergInternalFunctions(const std::string_view& prefix) {
+  static std::once_flag registerFlag;
+
+  std::call_once(registerFlag, [prefix]() {
+    functions::iceberg::registerFunctions(std::string(prefix));
+  });
+}
+
 IcebergInsertTableHandle::IcebergInsertTableHandle(
     std::vector<HiveColumnHandlePtr> inputColumns,
     LocationHandlePtr locationHandle,
@@ -57,7 +65,11 @@ IcebergDataSink::IcebergDataSink(
           commitStrategy,
           hiveConfig,
           0,
-          nullptr) {}
+          nullptr) {
+  static constexpr std::string_view kDefaultIcebergFunctionPrefix{
+      "$internal$.iceberg."};
+  registerIcebergInternalFunctions(kDefaultIcebergFunctionPrefix.data());
+}
 
 std::vector<std::string> IcebergDataSink::commitMessage() const {
   std::vector<std::string> commitTasks;

velox/connectors/hive/iceberg/IcebergDataSink.h

@@ -18,9 +18,16 @@
 
 #include "velox/connectors/hive/HiveDataSink.h"
 #include "velox/connectors/hive/iceberg/PartitionSpec.h"
+#include "velox/functions/iceberg/Register.h"
 
 namespace facebook::velox::connector::hive::iceberg {
 
+/// Registers Iceberg partition transform functions with prefix.
+/// NOTE: These functions are registered for internal transform usage only.
+/// Upstream engines such as Prestissimo and Gluten should register the same
+/// functions with different prefixes to avoid conflicts.
+void registerIcebergInternalFunctions(const std::string_view& prefix);
+
 /// Represents a request for Iceberg write.
 class IcebergInsertTableHandle final : public HiveInsertTableHandle {
  public:

velox/connectors/hive/iceberg/PartitionSpec.h

@@ -113,6 +113,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,111 @@
+/*
+ * 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.
+/// @param icebergFuncPrefix Prefix of iceberg transform function names.
+/// @return Typed expression representing the transform.
+core::TypedExprPtr toExpression(
+    const IcebergPartitionSpec::Field& field,
+    const std::string& inputFieldName,
+    const std::string& icebergFuncPrefix) {
+  // 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 = icebergFuncPrefix + "bucket";
+      break;
+    case TransformType::kTruncate:
+      functionName = icebergFuncPrefix + "truncate";
+      break;
+    case TransformType::kYear:
+      functionName = icebergFuncPrefix + "years";
+      break;
+    case TransformType::kMonth:
+      functionName = icebergFuncPrefix + "months";
+      break;
+    case TransformType::kDay:
+      functionName = icebergFuncPrefix + "days";
+      break;
+    case TransformType::kHour:
+      functionName = icebergFuncPrefix + "hours";
+      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,
+    const std::string& icebergFuncPrefix) {
+  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),
+        icebergFuncPrefix));
+  }
+
+  return transformExprs;
+}
+
+} // namespace facebook::velox::connector::hive::iceberg

velox/connectors/hive/iceberg/TransformExprBuilder.h

@@ -0,0 +1,51 @@
+/*
+ * 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.
+  /// @param icebergFuncPrefix Prefix for Iceberg transform function names.
+  /// @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,
+      const std::string& icebergFuncPrefix);
+};
+
+} // namespace facebook::velox::connector::hive::iceberg

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

@@ -64,6 +64,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)