Add iceberg partition transform evaluator and utility

Author: PingLiuPing

velox/connectors/hive/iceberg/CMakeLists.txt

@@ -15,10 +15,20 @@
 velox_add_library(
   velox_hive_iceberg_splitreader
   IcebergDataSink.cpp
+  IcebergPartitionPath.cpp
   IcebergSplit.cpp
   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/IcebergPartitionPath.cpp

@@ -0,0 +1,85 @@
+/*
+ * 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/IcebergPartitionPath.h"
+#include "velox/common/encode/Base64.h"
+
+namespace facebook::velox::connector::hive::iceberg {
+
+constexpr int32_t kEpochYear = 1970;
+
+std::string IcebergPartitionPath::toPartitionString(
+    int32_t value,
+    const TypePtr& type) const {
+  switch (transformType_) {
+    case TransformType::kIdentity: {
+      if (type->isDate()) {
+        return DATE()->toString(value);
+      }
+      return folly::to<std::string>(value);
+    }
+    case TransformType::kDay:
+      return DATE()->toString(value);
+    case TransformType::kYear:
+      return fmt::format("{:04d}", kEpochYear + value);
+    case TransformType::kMonth: {
+      int32_t year = kEpochYear + value / 12;
+      int32_t month = 1 + value % 12;
+      if (month <= 0) {
+        month += 12;
+        year -= 1;
+      }
+      return fmt::format("{:04d}-{:02d}", year, month);
+    }
+    case TransformType::kHour: {
+      int64_t seconds = static_cast<int64_t>(value) * 3600;
+      std::tm tmValue;
+      VELOX_USER_CHECK(
+          Timestamp::epochToCalendarUtc(seconds, tmValue),
+          "Failed to convert seconds to time: {}",
+          seconds);
+      return fmt::format(
+          "{:04d}-{:02d}-{:02d}-{:02d}",
+          tmValue.tm_year + 1900,
+          tmValue.tm_mon + 1,
+          tmValue.tm_mday,
+          tmValue.tm_hour);
+    }
+    default:
+      return folly::to<std::string>(value);
+  }
+}
+
+std::string IcebergPartitionPath::toPartitionString(
+    Timestamp value,
+    const TypePtr& type) const {
+  TimestampToStringOptions options;
+  options.precision = TimestampPrecision::kMilliseconds;
+  options.zeroPaddingYear = true;
+  options.skipTrailingZeros = true;
+  options.leadingPositiveSign = true;
+  return value.toString(options);
+}
+
+std::string IcebergPartitionPath::toPartitionString(
+    StringView value,
+    const TypePtr& type) const {
+  if (type->isVarbinary()) {
+    return encoding::Base64::encode(value.data(), value.size());
+  }
+  return std::string(value);
+}
+
+} // namespace facebook::velox::connector::hive::iceberg

velox/connectors/hive/iceberg/IcebergPartitionPath.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/HivePartitionUtil.h"
+#include "velox/connectors/hive/iceberg/PartitionSpec.h"
+
+namespace facebook::velox::connector::hive::iceberg {
+
+/// Converts a partition value to its string representation for use in
+/// partition directory path. The format follows the Iceberg specification
+/// for partition path encoding.
+class IcebergPartitionPath : public HivePartitionUtil {
+ public:
+  explicit IcebergPartitionPath(TransformType transformType)
+      : transformType_(transformType) {}
+
+  ~IcebergPartitionPath() override = default;
+
+  using HivePartitionUtil::toPartitionString;
+
+  std::string toPartitionString(int32_t value, const TypePtr& type)
+      const override;
+
+  std::string toPartitionString(Timestamp value, const TypePtr& type)
+      const override;
+
+  std::string toPartitionString(StringView value, const TypePtr& type)
+      const override;
+
+ private:
+  const TransformType transformType_;
+};
+
+using IcebergPartitionPathPtr = std::shared_ptr<const IcebergPartitionPath>;
+
+} // namespace facebook::velox::connector::hive::iceberg

velox/connectors/hive/iceberg/TransformEvaluator.cpp

@@ -0,0 +1,62 @@
+/*
+ * 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 <mutex>
+
+#include "velox/expression/Expr.h"
+#include "velox/functions/iceberg/Register.h"
+
+namespace facebook::velox::connector::hive::iceberg {
+
+namespace {
+
+constexpr char const* kIcebergFunctionPrefix{"iceberg_"};
+
+}
+
+TransformEvaluator::TransformEvaluator(
+    const std::vector<core::TypedExprPtr>& expressions,
+    const ConnectorQueryCtx* connectorQueryCtx)
+    : connectorQueryCtx_(connectorQueryCtx) {
+  VELOX_CHECK_NOT_NULL(connectorQueryCtx_);
+  // Register Iceberg functions once for expression evaluation.
+  static std::once_flag registerFlag;
+  std::call_once(registerFlag, []() {
+    functions::iceberg::registerFunctions(kIcebergFunctionPrefix);
+  });
+
+  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,68 @@
+/*
+ * 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. Iceberg partition transform functions with 'iceberg_'
+/// prefix are automatically registered during create first TransformEvaluator
+/// instance via std::call_once.
+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,102 @@
+/*
+ * 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 {
+
+constexpr char const* kBucketFunction{"iceberg_bucket"};
+constexpr char const* kTruncateFunction{"iceberg_truncate"};
+constexpr char const* kYearFunction{"iceberg_years"};
+constexpr char const* kMonthFunction{"iceberg_months"};
+constexpr char const* kDayFunction{"iceberg_days"};
+constexpr char const* kHourFunction{"iceberg_hours"};
+
+} // namespace
+
+core::TypedExprPtr TransformExprBuilder::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);
+}
+
+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