Collect Iceberg data file statistics during insertion

Author: PingLiuPing

velox/dwio/common/CMakeLists.txt

@@ -42,6 +42,7 @@ velox_add_library(
   OnDemandUnitLoader.cpp
   InputStream.cpp
   IntDecoder.cpp
+  IcebergStatistics.cpp
   MetadataFilter.cpp
   Options.cpp
   OutputStream.cpp

velox/dwio/common/IcebergStatistics.cpp

@@ -0,0 +1,43 @@
+/*
+ * 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/dwio/common/IcebergStatistics.h"
+
+namespace facebook::velox::dwio::common {
+
+folly::dynamic IcebergDataFileStatistics::toJson() const {
+  folly::dynamic json = folly::dynamic::object;
+  json["recordCount"] = numRecords;
+
+  auto mapToJson = [](const auto& map) {
+    folly::dynamic result = folly::dynamic::object;
+    for (const auto& pair : map) {
+      result[folly::to<std::string>(pair.first)] = pair.second;
+    }
+    return result;
+  };
+
+  json["columnSizes"] = mapToJson(columnsSizes);
+  json["valueCounts"] = mapToJson(valueCounts);
+  json["nullValueCounts"] = mapToJson(nullValueCounts);
+  json["nanValueCounts"] = mapToJson(nanValueCounts);
+  json["lowerBounds"] = mapToJson(lowerBounds);
+  json["upperBounds"] = mapToJson(upperBounds);
+
+  return json;
+}
+
+} // namespace facebook::velox::dwio::common

velox/dwio/common/IcebergStatistics.h

@@ -0,0 +1,38 @@
+/*
+ * 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 <folly/json/dynamic.h>
+
+namespace facebook::velox::dwio::common {
+
+// Iceberg data_file struct fields.
+struct IcebergDataFileStatistics {
+  int64_t numRecords;
+  std::unordered_map<int32_t, int64_t> columnsSizes;
+  std::unordered_map<int32_t, int64_t> valueCounts;
+  std::unordered_map<int32_t, int64_t> nullValueCounts;
+  std::unordered_map<int32_t, int64_t> nanValueCounts;
+  std::unordered_map<int32_t, std::string> lowerBounds;
+  std::unordered_map<int32_t, std::string> upperBounds;
+
+  IcebergDataFileStatistics() : numRecords(0) {}
+
+  folly::dynamic toJson() const;
+};
+
+} // namespace facebook::velox::dwio::common

velox/dwio/common/Options.h

@@ -706,6 +706,12 @@ struct WriterOptions {
   std::string sessionTimezoneName;
   bool adjustTimestampToTimezone{false};
 
+  // This option controls whether collect iceberg data file statistics
+  // during write. It will be set to true only by IcebergDataSink.
+  bool collectIcebergStats{false};
+  std::shared_ptr<std::vector<std::pair<int32_t, bool>>> sourceColumnIndices{
+      nullptr};
+
   // WriterOption implementations can implement this function to specify how to
   // process format-specific session and connector configs.
   virtual void processConfigs(

velox/dwio/common/Statistics.h

@@ -18,6 +18,7 @@
 
 #include <folly/Hash.h>
 #include <folly/container/F14Map.h>
+#include <folly/json/dynamic.h>
 
 #include "velox/common/base/Exceptions.h"
 #include "velox/common/base/RuntimeMetrics.h"

velox/dwio/common/Writer.h

@@ -21,6 +21,8 @@
 #include <optional>
 #include <string>
 
+#include "velox/dwio/common/IcebergStatistics.h"
+#include "velox/dwio/common/Statistics.h"
 #include "velox/vector/ComplexVector.h"
 
 namespace facebook::velox::dwio::common {
@@ -79,6 +81,11 @@ class Writer {
   /// Data can no longer be written.
   virtual void abort() = 0;
 
+  /// Return statistics based on each Iceberg data file
+  std::shared_ptr<IcebergDataFileStatistics> dataFileStats() const {
+    return dataFileStats_;
+  };
+
  protected:
   bool isRunning() const;
   bool isFinishing() const;
@@ -92,6 +99,10 @@ class Writer {
   static void checkStateTransition(State oldState, State newState);
 
   State state_{State::kInit};
+  bool collectIcebergStats_{false};
+  std::shared_ptr<IcebergDataFileStatistics> dataFileStats_{nullptr};
+  std::shared_ptr<std::vector<std::pair<int32_t, bool>>> sourceColumnIndices_{
+      nullptr};
 };
 
 FOLLY_ALWAYS_INLINE std::ostream& operator<<(

velox/dwio/parquet/common/CMakeLists.txt

@@ -17,7 +17,8 @@ velox_add_library(
   BloomFilter.cpp
   XxHasher.cpp
   LevelComparison.cpp
-  LevelConversion.cpp)
+  LevelConversion.cpp
+  UnicodeUtil.cpp)
 
 velox_link_libraries(
   velox_dwio_parquet_common

velox/dwio/parquet/common/UnicodeUtil.cpp

@@ -0,0 +1,82 @@
+/*
+ * 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/dwio/parquet/common/UnicodeUtil.h"
+
+#include "velox/external/utf8proc/utf8proc.h"
+#include "velox/functions/lib/string/StringImpl.h"
+
+namespace facebook::velox::parquet {
+
+std::string UnicodeUtil::truncateStringMin(
+    const char* input,
+    int32_t inputLength,
+    int32_t numCodePoints) {
+  auto length = functions::stringImpl::cappedByteLength<false>(
+      StringView(input, inputLength), numCodePoints);
+  return std::string(input, length);
+}
+
+std::string UnicodeUtil::truncateStringMax(
+    const char* input,
+    int32_t inputLength,
+    int32_t numCodePoints) {
+  auto truncatedLength = functions::stringImpl::cappedByteLength<false>(
+      StringView(input, inputLength), numCodePoints);
+
+  if (truncatedLength == inputLength) {
+    return std::string(input, inputLength);
+  }
+
+  // Try to increment the last code point.
+  for (auto i = numCodePoints - 1; i >= 0; --i) {
+    const char* current = input;
+    int32_t currentCodePoint = 0;
+
+    // Find the i-th code point position.
+    while (current < input + truncatedLength && currentCodePoint < i) {
+      int32_t charLength;
+      utf8proc_codepoint(current, input + truncatedLength, charLength);
+      current += charLength;
+      currentCodePoint++;
+    }
+
+    if (current >= input + truncatedLength)
+      continue;
+
+    int32_t charLength;
+    auto codePoint =
+        utf8proc_codepoint(current, input + truncatedLength, charLength);
+    auto nextCodePoint = codePoint + 1;
+
+    // Check if the incremented code point is valid.
+    if (nextCodePoint != 0 && utf8proc_codepoint_valid(nextCodePoint)) {
+      std::string result;
+      result.reserve(truncatedLength + 4);
+      result.assign(input, current - input);
+      char buffer[4];
+      auto bytesWritten = utf8proc_encode_char(
+          nextCodePoint, reinterpret_cast<utf8proc_uint8_t*>(buffer));
+      result.append(buffer, bytesWritten);
+      return result;
+    }
+  }
+
+  // Return truncated string without intermediate string creation
+  return std::string(input, truncatedLength);
+}
+
+} // namespace facebook::velox::parquet

velox/dwio/parquet/common/UnicodeUtil.h

@@ -0,0 +1,45 @@
+/*
+ * 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 <cstdint>
+#include <optional>
+#include <string>
+
+namespace facebook::velox::parquet {
+
+class UnicodeUtil {
+ public:
+  static bool isCharHighSurrogate(char16_t ch) {
+    return (ch & 0xFC00) == 0xD800;
+  }
+
+  static std::string truncateStringMin(
+      const char* input,
+      int32_t inputLength,
+      int32_t numCodePoints);
+
+  static std::string truncateStringMax(
+      const char* input,
+      int32_t inputLength,
+      int32_t numCodePoints);
+
+ private:
+  UnicodeUtil() = delete;
+};
+
+} // namespace facebook::velox::parquet

velox/dwio/parquet/tests/common/CMakeLists.txt

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-add_executable(velox_dwio_parquet_common_test LevelConversionTest.cpp)
+add_executable(velox_dwio_parquet_common_test LevelConversionTest.cpp
+                                              UnicodeUtilTest.cpp)
 
 add_test(velox_dwio_parquet_common_test velox_dwio_parquet_common_test)
 target_link_libraries(