Integrate Iceberg data file statistics and adding unit test.

Author: PingLiuPing

velox/connectors/hive/HiveDataSink.h

@@ -679,7 +679,7 @@ class HiveDataSink : public DataSink {
   // Invoked to write 'input' to the specified file writer.
   void write(size_t index, RowVectorPtr input);
 
-  void closeInternal();
+  virtual void closeInternal();
 
   // IMPORTANT NOTE: these are passed to writers as raw pointers. HiveDataSink
   // owns the lifetime of these objects, and therefore must destroy them last.

velox/connectors/hive/iceberg/CMakeLists.txt

@@ -14,8 +14,10 @@
 
 velox_add_library(
   velox_hive_iceberg_splitreader
+  DataFileStatsCollector.cpp
   EqualityDeleteFileReader.cpp
   FilterUtil.cpp
+  IcebergColumnHandle.cpp
   IcebergDataSink.cpp
   IcebergDeleteFile.cpp
   IcebergPartitionIdGenerator.cpp

velox/connectors/hive/iceberg/DataFileStatsCollector.cpp

@@ -0,0 +1,121 @@
+/*
+ * 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/DataFileStatsCollector.h"
+#include "velox/common/base/Exceptions.h"
+#include "velox/common/encode/Base64.h"
+#include "velox/dwio/parquet/writer/arrow/Metadata.h"
+#include "velox/dwio/parquet/writer/arrow/Statistics.h"
+
+namespace facebook::velox::connector::hive::iceberg {
+
+using namespace facebook::velox::parquet;
+
+DataFileStatsCollector::DataFileStatsCollector(
+    std::shared_ptr<
+        std::vector<std::unique_ptr<dwio::common::DataFileStatsSettings>>>
+        settings)
+    : FileStatsCollector(std::move(settings)) {}
+
+void DataFileStatsCollector::collectStats(
+    const void* metadata,
+    const std::shared_ptr<dwio::common::DataFileStatistics>& dataFileStats) {
+  const auto& fileMetadata =
+      *static_cast<const std::shared_ptr<parquet::arrow::FileMetaData>*>(
+          metadata);
+  VELOX_CHECK_NOT_NULL(fileMetadata);
+
+  std::unordered_set<int32_t> skipBoundsFields;
+  std::function<int32_t(IcebergDataFileStatsSettings*)> processFields =
+      [&skipBoundsFields,
+       &processFields](IcebergDataFileStatsSettings* field) -> int32_t {
+    if (field->skipBounds) {
+      skipBoundsFields.insert(field->fieldId);
+    }
+    if (field->children.empty()) {
+      return 1;
+    }
+    int32_t count = 0;
+    for (const auto& child : field->children) {
+      count += processFields(child.get());
+    }
+    return count;
+  };
+
+  // numFields is not the number of columns in Iceberg table's schema,
+  // e.g., schema_->size(). It also contains the sub-fields when there are
+  // nested types in table's schema.
+  int32_t numFields = 0;
+  for (const auto& field : *statsSetting_) {
+    auto* icebergField =
+        static_cast<IcebergDataFileStatsSettings*>(field.get());
+    numFields += processFields(icebergField);
+  }
+
+  std::unordered_map<int32_t, std::shared_ptr<arrow::Statistics>>
+      globalMinStats;
+  std::unordered_map<int32_t, std::shared_ptr<arrow::Statistics>>
+      globalMaxStats;
+
+  dataFileStats->numRecords = fileMetadata->num_rows();
+  const auto numRowGroups = fileMetadata->num_row_groups();
+  for (auto i = 0; i < numRowGroups; ++i) {
+    const auto rgm = fileMetadata->RowGroup(i);
+    VELOX_CHECK_EQ(numFields, rgm->num_columns());
+    dataFileStats->splitOffsets.emplace_back(rgm->file_offset());
+
+    for (auto j = 0; j < numFields; ++j) {
+      const auto columnChunkMetadata = rgm->ColumnChunk(j);
+      const auto fieldId = columnChunkMetadata->field_id();
+      const auto numValues = columnChunkMetadata->num_values();
+
+      dataFileStats->valueCounts[fieldId] += numValues;
+      dataFileStats->columnsSizes[fieldId] +=
+          columnChunkMetadata->total_compressed_size();
+
+      const auto columnChunkStats = columnChunkMetadata->statistics();
+      if (columnChunkStats->nan_count() > 0) {
+        dataFileStats->nanValueCounts[fieldId] += columnChunkStats->nan_count();
+      }
+      dataFileStats->nullValueCounts[fieldId] += columnChunkStats->null_count();
+
+      if (columnChunkStats->HasMinMax() &&
+          !skipBoundsFields.contains(fieldId)) {
+        if (globalMaxStats.find(fieldId) == globalMaxStats.end()) {
+          globalMinStats[fieldId] = columnChunkStats;
+          globalMaxStats[fieldId] = columnChunkStats;
+        } else {
+          globalMaxStats[fieldId] = arrow::Statistics::CompareAndGetMax(
+              globalMaxStats[fieldId], columnChunkStats);
+          globalMinStats[fieldId] = arrow::Statistics::CompareAndGetMin(
+              globalMinStats[fieldId], columnChunkStats);
+        }
+      }
+    }
+  }
+
+  for (const auto& [fieldId, minStats] : globalMinStats) {
+    const auto lowerBound = minStats->MinValue();
+    dataFileStats->lowerBounds[fieldId] =
+        encoding::Base64::encode(lowerBound.data(), lowerBound.size());
+  }
+  for (const auto& [fieldId, maxStats] : globalMaxStats) {
+    const auto upperBound = maxStats->MaxValue();
+    dataFileStats->upperBounds[fieldId] =
+        encoding::Base64::encode(upperBound.data(), upperBound.size());
+  }
+}
+
+} // namespace facebook::velox::connector::hive::iceberg

velox/connectors/hive/iceberg/DataFileStatsCollector.h

@@ -0,0 +1,48 @@
+/*
+ * 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/dwio/common/DataFileStatsCollector.h"
+
+namespace facebook::velox::connector::hive::iceberg {
+
+/// Settings for collecting Iceberg parquet data file statistics.
+/// Holds the Iceberg source field id and whether to skip bounds
+/// collection for this field. For nested field, it contains child fields.
+struct IcebergDataFileStatsSettings
+    : public dwio::common::DataFileStatsSettings {
+  int32_t fieldId;
+  bool skipBounds;
+  std::vector<std::unique_ptr<IcebergDataFileStatsSettings>> children;
+
+  IcebergDataFileStatsSettings(int32_t id, bool skip)
+      : fieldId(id), skipBounds(skip), children() {}
+};
+
+class DataFileStatsCollector : public dwio::common::FileStatsCollector {
+ public:
+  explicit DataFileStatsCollector(
+      std::shared_ptr<
+          std::vector<std::unique_ptr<dwio::common::DataFileStatsSettings>>>
+          settings);
+
+  void collectStats(
+      const void* metadata,
+      const std::shared_ptr<dwio::common::DataFileStatistics>& fileStats)
+      override;
+};
+
+} // namespace facebook::velox::connector::hive::iceberg

velox/connectors/hive/iceberg/IcebergColumnHandle.cpp

@@ -0,0 +1,42 @@
+/*
+ * 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/IcebergColumnHandle.h"
+
+namespace facebook::velox::connector::hive::iceberg {
+
+IcebergColumnHandle::IcebergColumnHandle(
+    const std::string& name,
+    ColumnType columnType,
+    TypePtr dataType,
+    TypePtr hiveType,
+    const IcebergNestedField& nestedField,
+    std::vector<common::Subfield> requiredSubfields,
+    ColumnParseParameters columnParseParameters)
+    : HiveColumnHandle(
+          name,
+          columnType,
+          dataType,
+          hiveType,
+          std::move(requiredSubfields),
+          columnParseParameters),
+      nestedField_(nestedField) {}
+
+const IcebergNestedField& IcebergColumnHandle::nestedField() const {
+  return nestedField_;
+}
+
+} // namespace facebook::velox::connector::hive::iceberg

velox/connectors/hive/iceberg/IcebergColumnHandle.h

@@ -0,0 +1,44 @@
+/*
+ * 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/TableHandle.h"
+
+namespace facebook::velox::connector::hive::iceberg {
+
+struct IcebergNestedField {
+  int32_t id;
+  std::vector<IcebergNestedField> children;
+};
+
+class IcebergColumnHandle : public HiveColumnHandle {
+ public:
+  IcebergColumnHandle(
+      const std::string& name,
+      ColumnType columnType,
+      TypePtr dataType,
+      TypePtr hiveType,
+      const IcebergNestedField& nestedField,
+      std::vector<common::Subfield> requiredSubfields = {},
+      ColumnParseParameters columnParseParameters = {});
+
+  const IcebergNestedField& nestedField() const;
+
+ private:
+  const IcebergNestedField nestedField_;
+};
+
+} // namespace facebook::velox::connector::hive::iceberg