fix(iceberg): Date partition value parse issue

Author: nmahadevuni

velox/connectors/hive/HiveConnectorUtil.cpp

@@ -634,12 +634,18 @@ namespace {
 bool applyPartitionFilter(
     const TypePtr& type,
     const std::string& partitionValue,
+    bool isPartitionValueDaysSinceEpoch,
     common::Filter* filter) {
   if (type->isDate()) {
-    const auto result = util::fromDateString(
-        StringView(partitionValue), util::ParseMode::kPrestoCast);
-    VELOX_CHECK(!result.hasError());
-    return applyFilter(*filter, result.value());
+    int32_t result = 0;
+    // Iceberg partition values are already in daysSinceEpoch, no need to
+    // convert.
+    if (isPartitionValueDaysSinceEpoch) {
+      result = folly::to<int32_t>(partitionValue);
+    } else {
+      result = DATE()->toDays(static_cast<folly::StringPiece>(partitionValue));
+    }
+    return applyFilter(*filter, result);
   }
 
   switch (type->kind()) {
@@ -681,7 +687,8 @@ bool testFilters(
     const std::unordered_map<std::string, std::optional<std::string>>&
         partitionKeys,
     const std::unordered_map<std::string, std::shared_ptr<HiveColumnHandle>>&
-        partitionKeysHandle) {
+        partitionKeysHandle,
+    bool isPartitionValueDaysSinceEpoch) {
   const auto totalRows = reader->numberOfRows();
   const auto& fileTypeWithId = reader->typeWithId();
   const auto& rowType = reader->rowType();
@@ -701,6 +708,7 @@ bool testFilters(
           return applyPartitionFilter(
               handlesIter->second->dataType(),
               iter->second.value(),
+              isPartitionValueDaysSinceEpoch,
               child->filter());
         }
         // Column is missing, most likely due to schema evolution. Or it's a

velox/connectors/hive/HiveConnectorUtil.h

@@ -20,6 +20,7 @@
 
 #include "velox/connectors/Connector.h"
 #include "velox/connectors/hive/FileHandle.h"
+#include "velox/connectors/hive/SplitReader.h"
 #include "velox/dwio/common/BufferedInput.h"
 #include "velox/dwio/common/Reader.h"
 
@@ -97,7 +98,8 @@ bool testFilters(
     const std::unordered_map<std::string, std::optional<std::string>>&
         partitionKey,
     const std::unordered_map<std::string, std::shared_ptr<HiveColumnHandle>>&
-        partitionKeysHandle);
+        partitionKeysHandle,
+    bool isPartitionValueDaysSinceEpoch = false);
 
 std::unique_ptr<dwio::common::BufferedInput> createBufferedInput(
     const FileHandle& fileHandle,

velox/connectors/hive/SplitReader.cpp

@@ -35,14 +35,22 @@ VectorPtr newConstantFromString(
     const std::optional<std::string>& value,
     vector_size_t size,
     velox::memory::MemoryPool* pool,
-    const std::string& sessionTimezone) {
+    const std::string& sessionTimezone,
+    bool isPartitionValueDaysSinceEpoch) {
   using T = typename TypeTraits<kind>::NativeType;
   if (!value.has_value()) {
     return std::make_shared<ConstantVector<T>>(pool, size, true, type, T());
   }
 
   if (type->isDate()) {
-    auto days = DATE()->toDays(static_cast<folly::StringPiece>(value.value()));
+    int32_t days = 0;
+    // For Iceberg, the date partition values are already in daysSinceEpoch
+    // form.
+    if (isPartitionValueDaysSinceEpoch) {
+      days = folly::to<int32_t>(value.value());
+    } else {
+      days = DATE()->toDays(static_cast<folly::StringPiece>(value.value()));
+    }
     return std::make_shared<ConstantVector<int32_t>>(
         pool, size, false, type, std::move(days));
   }
@@ -101,7 +109,8 @@ std::unique_ptr<SplitReader> SplitReader::create(
         ioStats,
         fileHandleFactory,
         executor,
-        scanSpec));
+        scanSpec,
+        false));
   }
 }
 
@@ -116,12 +125,14 @@ SplitReader::SplitReader(
     const std::shared_ptr<io::IoStatistics>& ioStats,
     FileHandleFactory* fileHandleFactory,
     folly::Executor* executor,
-    const std::shared_ptr<common::ScanSpec>& scanSpec)
+    const std::shared_ptr<common::ScanSpec>& scanSpec,
+    bool isPartitionValueDaysSinceEpoch)
     : hiveSplit_(hiveSplit),
       hiveTableHandle_(hiveTableHandle),
       partitionKeys_(partitionKeys),
       connectorQueryCtx_(connectorQueryCtx),
       hiveConfig_(hiveConfig),
+      isPartitionValueDaysSinceEpoch_(isPartitionValueDaysSinceEpoch),
       readerOutputType_(readerOutputType),
       ioStats_(ioStats),
       fileHandleFactory_(fileHandleFactory),
@@ -272,7 +283,8 @@ bool SplitReader::filterOnStats(
           baseReader_.get(),
           hiveSplit_->filePath,
           hiveSplit_->partitionKeys,
-          *partitionKeys_)) {
+          *partitionKeys_,
+          isPartitionValueDaysSinceEpoch_)) {
     return true;
   }
   ++runtimeStats.skippedSplits;
@@ -335,7 +347,8 @@ std::vector<TypePtr> SplitReader::adaptColumns(
           iter->second,
           1,
           connectorQueryCtx_->memoryPool(),
-          connectorQueryCtx_->sessionTimezone());
+          connectorQueryCtx_->sessionTimezone(),
+          {});
       childSpec->setConstantValue(constant);
     } else if (
         childSpec->columnType() == common::ScanSpec::ColumnType::kRegular) {
@@ -389,7 +402,8 @@ void SplitReader::setPartitionValue(
       value,
       1,
       connectorQueryCtx_->memoryPool(),
-      connectorQueryCtx_->sessionTimezone());
+      connectorQueryCtx_->sessionTimezone(),
+      isPartitionValueDaysSinceEpoch_);
   spec->setConstantValue(constant);
 }
 

velox/connectors/hive/SplitReader.h

@@ -109,7 +109,8 @@ class SplitReader {
       const std::shared_ptr<io::IoStatistics>& ioStats,
       FileHandleFactory* fileHandleFactory,
       folly::Executor* executor,
-      const std::shared_ptr<common::ScanSpec>& scanSpec);
+      const std::shared_ptr<common::ScanSpec>& scanSpec,
+      bool isPartitionValueDaysSinceEpoch);
 
   /// Create the dwio::common::Reader object baseReader_, which will be used to
   /// read the data file's metadata and schema
@@ -154,6 +155,7 @@ class SplitReader {
       std::shared_ptr<HiveColumnHandle>>* const partitionKeys_;
   const ConnectorQueryCtx* connectorQueryCtx_;
   const std::shared_ptr<const HiveConfig> hiveConfig_;
+  bool isPartitionValueDaysSinceEpoch_;
 
   const RowTypePtr readerOutputType_;
   const std::shared_ptr<io::IoStatistics> ioStats_;

velox/connectors/hive/iceberg/IcebergSplitReader.cpp

@@ -46,7 +46,8 @@ IcebergSplitReader::IcebergSplitReader(
           ioStats,
           fileHandleFactory,
           executor,
-          scanSpec),
+          scanSpec,
+          true),
       baseReadOffset_(0),
       splitOffset_(0),
       deleteBitmap_(nullptr),

velox/connectors/hive/iceberg/tests/IcebergReadTest.cpp

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include "velox/common/base/tests/GTestUtils.h"
 #include "velox/common/file/FileSystems.h"
 #include "velox/connectors/hive/HiveConnectorSplit.h"
 #include "velox/connectors/hive/iceberg/IcebergDeleteFile.h"
@@ -225,6 +226,35 @@ class HiveIcebergTest : public HiveConnectorTestBase {
     ASSERT_TRUE(it->second.peakMemoryBytes > 0);
   }
 
+  void assertQuery(
+      RowTypePtr rowType,
+      const std::vector<RowVectorPtr>& dataVectors,
+      const std::string duckDbSql,
+      const std::unordered_map<std::string, std::optional<std::string>>
+          partitionKeys = {},
+      const std::vector<std::string> filters = {}) {
+    VELOX_CHECK(!duckDbSql.empty(), "DuckDb sql is empty");
+    auto dataFilePath = TempFilePath::create();
+
+    writeToFile(
+        dataFilePath->getPath(), dataVectors, config_, flushPolicyFactory_);
+    std::vector<std::shared_ptr<ConnectorSplit>> splits;
+    splits.emplace_back(
+        makeIcebergSplit(dataFilePath->getPath(), {}, partitionKeys));
+
+    std::unordered_set<std::string> partitionColumns;
+
+    for (auto partitionKey : partitionKeys) {
+      partitionColumns.insert(partitionKey.first);
+    }
+
+    auto plan =
+        PlanBuilder(pool_.get())
+            .tableScan(rowType, filters, "", nullptr, {}, partitionColumns)
+            .planNode();
+    HiveConnectorTestBase::assertQuery(plan, splits, duckDbSql);
+  }
+
   const static int rowCount = 20000;
 
  private:
@@ -337,8 +367,9 @@ class HiveIcebergTest : public HiveConnectorTestBase {
 
   std::shared_ptr<ConnectorSplit> makeIcebergSplit(
       const std::string& dataFilePath,
-      const std::vector<IcebergDeleteFile>& deleteFiles = {}) {
-    std::unordered_map<std::string, std::optional<std::string>> partitionKeys;
+      const std::vector<IcebergDeleteFile>& deleteFiles = {},
+      const std::unordered_map<std::string, std::optional<std::string>>
+          partitionKeys = {}) {
     std::unordered_map<std::string, std::string> customSplitInfo;
     customSplitInfo["table_format"] = "hive-iceberg";
 
@@ -660,4 +691,25 @@ TEST_F(HiveIcebergTest, positionalDeletesMultipleSplits) {
   assertMultipleSplits({}, 10, 3);
 }
 
+TEST_F(HiveIcebergTest, testPartitionedRead) {
+  RowTypePtr rowType{ROW({"c0", "ds"}, {BIGINT(), DateType::get()})};
+  std::unordered_map<std::string, std::optional<std::string>> partitionKeys;
+  // Iceberg API sets partition values for dates to daysSinceEpoch, so
+  // in velox, we do not need to convert it to days.
+  // Date = 2018-04-06, daysSinceEpoch = 17627
+  partitionKeys["ds"] = "17627";
+
+  std::vector<RowVectorPtr> dataVectors;
+  VectorPtr c0 = makeFlatVector<int64_t>((std::vector<int64_t>){1});
+  VectorPtr ds = makeFlatVector<int32_t>((std::vector<int32_t>){17627});
+  dataVectors.push_back(makeRowVector({"c0", "ds"}, {c0, ds}));
+
+  assertQuery(
+      rowType, dataVectors, "SELECT 1, '2018-04-06'", partitionKeys, {});
+
+  std::vector<std::string> filters = {"ds = date'2018-04-06'"};
+  assertQuery(
+      rowType, dataVectors, "SELECT 1, '2018-04-06'", partitionKeys, filters);
+}
+
 } // namespace facebook::velox::connector::hive::iceberg

velox/exec/tests/utils/PlanBuilder.cpp

@@ -74,14 +74,16 @@ PlanBuilder& PlanBuilder::tableScan(
     const RowTypePtr& dataColumns,
     const std::unordered_map<
         std::string,
-        std::shared_ptr<connector::ColumnHandle>>& assignments) {
+        std::shared_ptr<connector::ColumnHandle>>& assignments,
+    const std::unordered_set<std::string>& partitionKeys) {
   return TableScanBuilder(*this)
       .outputType(outputType)
       .assignments(assignments)
       .subfieldFilters(subfieldFilters)
       .remainingFilter(remainingFilter)
       .dataColumns(dataColumns)
       .assignments(assignments)
+      .partitionKeys(partitionKeys)
       .endTableScan();
 }
 
@@ -170,13 +172,14 @@ core::PlanNodePtr PlanBuilder::TableScanBuilder::build(core::PlanNodeId id) {
     }
 
     if (!hasAssignments) {
+      auto columnType = HiveColumnHandle::ColumnType::kRegular;
+      if (partitionKeys_.count(name) > 0) {
+        columnType = HiveColumnHandle::ColumnType::kPartitionKey;
+      }
       assignments_.insert(
           {name,
            std::make_shared<HiveColumnHandle>(
-               hiveColumnName,
-               HiveColumnHandle::ColumnType::kRegular,
-               type,
-               type)});
+               hiveColumnName, columnType, type, type)});
     }
   }
 

velox/exec/tests/utils/PlanBuilder.h

@@ -131,14 +131,16 @@ class PlanBuilder {
   /// you define the output types only. See 'missingColumns' test in
   /// 'TableScanTest'.
   /// @param assignments Optional ColumnHandles.
+  /// @param partitionKeys Optional partition keys.
   PlanBuilder& tableScan(
       const RowTypePtr& outputType,
       const std::vector<std::string>& subfieldFilters = {},
       const std::string& remainingFilter = "",
       const RowTypePtr& dataColumns = nullptr,
       const std::unordered_map<
           std::string,
-          std::shared_ptr<connector::ColumnHandle>>& assignments = {});
+          std::shared_ptr<connector::ColumnHandle>>& assignments = {},
+      const std::unordered_set<std::string>& partitionKeys = {});
 
   /// Add a TableScanNode to scan a Hive table.
   ///
@@ -277,6 +279,12 @@ class PlanBuilder {
       return *this;
     }
 
+    TableScanBuilder& partitionKeys(
+        std::unordered_set<std::string> partitionKeys) {
+      partitionKeys_ = std::move(partitionKeys);
+      return *this;
+    }
+
     /// Stop the TableScanBuilder.
     PlanBuilder& endTableScan() {
       planBuilder_.planNode_ = build(planBuilder_.nextPlanNodeId());
@@ -298,6 +306,7 @@ class PlanBuilder {
     std::shared_ptr<connector::ConnectorTableHandle> tableHandle_;
     std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
         assignments_;
+    std::unordered_set<std::string> partitionKeys_;
   };
 
   /// Start a TableScanBuilder.