[oap ] Add support to read plain encoded INT96 timestamp from Parquet…

… file
oap-project · Jul 25, 2024 · b138d15 · b138d15
1 parent 81b8578
commit b138d15
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 2 deletions.
diff --git a/velox/dwio/common/DirectDecoder.h b/velox/dwio/common/DirectDecoder.h
@@ -92,7 +92,17 @@ class DirectDecoder : public IntDecoder<isSigned> {
       } else if constexpr (std::is_same_v<
                                typename Visitor::DataType,
                                int128_t>) {
-        toSkip = visitor.process(super::template readInt<int128_t>(), atEnd);
+        if (super::numBytes == 12 /* INT96 */) {
+          int128_t encoded = super::template readInt<int128_t>();
+          int32_t days = encoded & ((1ULL << 32) - 1);
+          uint64_t nanos = static_cast<uint64_t>(encoded >> 32);
+
+          auto timestamp = Timestamp::fromDaysAndNanos(days, nanos);
+          toSkip =
+              visitor.process(*reinterpret_cast<int128_t*>(&timestamp), atEnd);
+        } else {
+          toSkip = visitor.process(super::template readInt<int128_t>(), atEnd);
+        }
       } else {
         toSkip = visitor.process(super::template readInt<int64_t>(), atEnd);
       }

diff --git a/velox/dwio/common/IntDecoder.h b/velox/dwio/common/IntDecoder.h
@@ -154,6 +154,9 @@ class IntDecoder {
   template <typename T>
   T readInt();
 
+  template <typename T>
+  T readInt96();
+
   template <typename T>
   T readVInt();
 
@@ -438,12 +441,44 @@ inline T IntDecoder<isSigned>::readInt() {
     return readLittleEndianFromBigEndian<T>();
   } else {
     if constexpr (std::is_same_v<T, int128_t>) {
-      VELOX_NYI();
+      if (numBytes == 12) {
+        // TODO:: Do we need to handle useVInts case?
+        return readInt96<T>();
+      } else {
+        VELOX_NYI();
+      }
     }
     return readLongLE();
   }
 }
 
+template <bool isSigned>
+template <typename T>
+inline T IntDecoder<isSigned>::readInt96() {
+  int64_t offset = 0;
+  unsigned char ch;
+
+  // read unsigned byte 64
+  uint64_t part1 = 0;
+  for (uint32_t i = 0; i < 8; ++i) {
+    ch = readByte();
+    part1 |= (ch & BASE_256_MASK) << offset;
+    offset += 8;
+  }
+
+  // read signed byte 32
+  int32_t part2 = 0;
+  offset = 0;
+  for (uint32_t i = 0; i < 4; ++i) {
+    ch = readByte();
+    part2 |= (ch & BASE_256_MASK) << offset;
+    offset += 8;
+  }
+
+  int128_t result = part1;
+  return (result << 32) | part2;
+}
+
 template <bool isSigned>
 template <typename T>
 inline T IntDecoder<isSigned>::readVInt() {

diff --git a/velox/dwio/parquet/tests/examples/timestamp_dict_int96.parquet b/velox/dwio/parquet/tests/examples/timestamp_dict_int96.parquet
diff --git a/velox/dwio/parquet/tests/examples/timestamp_plain_int96.parquet b/velox/dwio/parquet/tests/examples/timestamp_plain_int96.parquet
diff --git a/velox/dwio/parquet/tests/reader/ParquetTableScanTest.cpp b/velox/dwio/parquet/tests/reader/ParquetTableScanTest.cpp
@@ -898,6 +898,34 @@ TEST_F(ParquetTableScanTest, timestampPrecisionMicrosecond) {
   assertEqualResults({expected}, result.second);
 }
 
+
+TEST_F(ParquetTableScanTest, timestampINT96) {
+  auto a = makeFlatVector<Timestamp>({Timestamp(1, 0), Timestamp(2, 0)});
+  auto expected = makeRowVector({"time"}, {a});
+  createDuckDbTable("expected", {expected});
+
+  auto vector = makeArrayVector<Timestamp>({{}});
+  loadData(
+      getExampleFilePath("timestamp_dict_int96.parquet"),
+      ROW({"time"}, {TIMESTAMP()}),
+      makeRowVector(
+          {"time"},
+          {
+              vector,
+          }));
+  assertSelect({"time"}, "SELECT time from expected");
+
+  loadData(
+      getExampleFilePath("timestamp_plain_int96.parquet"),
+      ROW({"time"}, {TIMESTAMP()}),
+      makeRowVector(
+          {"time"},
+          {
+              vector,
+          }));
+  assertSelect({"time"}, "SELECT time from expected");
+}
+
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   folly::Init init{&argc, &argv, false};