fix: Align TimestampUnit and TimestampPrecision in unit tests

facebookincubator · Nov 30, 2024 · e1e50f9 · e1e50f9
1 parent ac5c15e
commit e1e50f9
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/velox/dwio/parquet/tests/reader/ParquetTableScanTest.cpp b/velox/dwio/parquet/tests/reader/ParquetTableScanTest.cpp
@@ -99,7 +99,13 @@ class ParquetTableScanTest : public HiveConnectorTestBase {
                 rowType, subfieldFilters, remainingFilter, nullptr, assignments)
             .planNode();
 
-    assertQuery(plan, splits_, sql);
+    AssertQueryBuilder(plan, duckDbQueryRunner_)
+        .connectorSessionProperty(
+            kHiveConnectorId,
+            HiveConfig::kReadTimestampUnitSession,
+            std::to_string(static_cast<int>(timestampPrecision_)))
+        .splits(splits_)
+        .assertResults(sql);
   }
 
   void assertSelectWithAgg(
@@ -211,6 +217,11 @@ class ParquetTableScanTest : public HiveConnectorTestBase {
         rootPool_->addAggregateChild("ParquetTableScanTest.Writer");
     options.memoryPool = childPool.get();
 
+    if (options.parquetWriteTimestampUnit) {
+      timestampPrecision_ =
+          static_cast<TimestampPrecision>(*options.parquetWriteTimestampUnit);
+    }
+
     auto writer = std::make_unique<Writer>(
         std::move(sink), options, asRowType(data[0]->type()));
 
@@ -298,6 +309,7 @@ class ParquetTableScanTest : public HiveConnectorTestBase {
 
   RowTypePtr rowType_;
   std::vector<std::shared_ptr<connector::ConnectorSplit>> splits_;
+  TimestampPrecision timestampPrecision_ = TimestampPrecision::kMilliseconds;
 };
 
 TEST_F(ParquetTableScanTest, basic) {
@@ -820,13 +832,15 @@ TEST_F(ParquetTableScanTest, timestampInt96Dictionary) {
   WriterOptions options;
   options.writeInt96AsTimestamp = true;
   options.enableDictionary = true;
+  options.parquetWriteTimestampUnit = TimestampUnit::kMilli;
   testInt96TimestampRead(options);
 }
 
 TEST_F(ParquetTableScanTest, timestampInt96Plain) {
   WriterOptions options;
   options.writeInt96AsTimestamp = true;
   options.enableDictionary = false;
+  options.parquetWriteTimestampUnit = TimestampUnit::kMilli;
   testInt96TimestampRead(options);
 }