Support timestamp reader for Parquet file format (4680)

rui-mo · rui-mo · commit 8c9af765670e · 2023-09-11T11:03:21.000+08:00
diff --git a/velox/dwio/common/SelectiveColumnReader.cpp b/velox/dwio/common/SelectiveColumnReader.cpp
@@ -214,6 +214,9 @@ void SelectiveColumnReader::getIntValues(
           VELOX_FAIL("Unsupported value size: {}", valueSize_);
       }
       break;
+    case TypeKind::TIMESTAMP:
+      getFlatValues<Timestamp, Timestamp>(rows, result, requestedType);
+      break;
     default:
       VELOX_FAIL(
           "Not a valid type for integer reader: {}", requestedType->toString());
diff --git a/velox/dwio/parquet/reader/PageReader.cpp b/velox/dwio/parquet/reader/PageReader.cpp
@@ -524,6 +524,51 @@ void PageReader::prepareDictionary(const PageHeader& pageHeader) {
       }
       break;
     }
+    case thrift::Type::INT96: {
+      auto numVeloxBytes = dictionary_.numValues * sizeof(Timestamp);
+      dictionary_.values = AlignedBuffer::allocate<char>(numVeloxBytes, &pool_);
+      auto numBytes = dictionary_.numValues * sizeof(Int96Timestamp);
+      if (pageData_) {
+        memcpy(dictionary_.values->asMutable<char>(), pageData_, numBytes);
+      } else {
+        dwio::common::readBytes(
+            numBytes,
+            inputStream_.get(),
+            dictionary_.values->asMutable<char>(),
+            bufferStart_,
+            bufferEnd_);
+      }
+      // Expand the Parquet type length values to Velox type length.
+      // We start from the end to allow in-place expansion.
+      auto values = dictionary_.values->asMutable<Timestamp>();
+      auto parquetValues = dictionary_.values->asMutable<char>();
+      static constexpr int64_t kJulianToUnixEpochDays = 2440588LL;
+      static constexpr int64_t kSecondsPerDay = 86400LL;
+      static constexpr int64_t kNanosPerSecond =
+          Timestamp::kNanosecondsInMillisecond *
+          Timestamp::kMillisecondsInSecond;
+      for (auto i = dictionary_.numValues - 1; i >= 0; --i) {
+        // Convert the timestamp into seconds and nanos since the Unix epoch,
+        // 00:00:00.000000 on 1 January 1970.
+        uint64_t nanos;
+        memcpy(
+            &nanos,
+            parquetValues + i * sizeof(Int96Timestamp),
+            sizeof(uint64_t));
+        int32_t days;
+        memcpy(
+            &days,
+            parquetValues + i * sizeof(Int96Timestamp) + sizeof(uint64_t),
+            sizeof(int32_t));
+        int64_t seconds = (days - kJulianToUnixEpochDays) * kSecondsPerDay;
+        if (nanos > Timestamp::kMaxNanos) {
+          seconds += nanos / kNanosPerSecond;
+          nanos -= (nanos / kNanosPerSecond) * kNanosPerSecond;
+        }
+        values[i] = Timestamp(seconds, nanos);
+      }
+      break;
+    }
     case thrift::Type::BYTE_ARRAY: {
       dictionary_.values =
           AlignedBuffer::allocate<StringView>(dictionary_.numValues, &pool_);
@@ -614,7 +659,6 @@ void PageReader::prepareDictionary(const PageHeader& pageHeader) {
       VELOX_UNSUPPORTED(
           "Parquet type {} not supported for dictionary", parquetType);
     }
-    case thrift::Type::INT96:
     default:
       VELOX_UNSUPPORTED(
           "Parquet type {} not supported for dictionary", parquetType);
@@ -641,6 +685,8 @@ int32_t parquetTypeBytes(thrift::Type::type type) {
     case thrift::Type::INT64:
     case thrift::Type::DOUBLE:
       return 8;
+    case thrift::Type::INT96:
+      return 12;
     default:
       VELOX_FAIL("Type does not have a byte width {}", type);
   }
diff --git a/velox/dwio/parquet/reader/ParquetColumnReader.cpp b/velox/dwio/parquet/reader/ParquetColumnReader.cpp
@@ -28,6 +28,7 @@
 #include "velox/dwio/parquet/reader/StructColumnReader.h"
 
 #include "velox/dwio/parquet/reader/Statistics.h"
+#include "velox/dwio/parquet/reader/TimestampColumnReader.h"
 #include "velox/dwio/parquet/thrift/ParquetThriftTypes.h"
 
 namespace facebook::velox::parquet {
@@ -90,6 +91,10 @@ std::unique_ptr<dwio::common::SelectiveColumnReader> ParquetColumnReader::build(
     case TypeKind::BOOLEAN:
       return std::make_unique<BooleanColumnReader>(dataType, params, scanSpec);
 
+    case TypeKind::TIMESTAMP:
+      return std::make_unique<TimestampColumnReader>(
+          dataType, params, scanSpec);
+
     default:
       VELOX_FAIL(
           "buildReader unhandled type: " +
diff --git a/velox/dwio/parquet/reader/ParquetReader.cpp b/velox/dwio/parquet/reader/ParquetReader.cpp
@@ -534,7 +534,7 @@ TypePtr ReaderBase::convertType(
       case thrift::Type::type::INT64:
         return BIGINT();
       case thrift::Type::type::INT96:
-        return DOUBLE(); // TODO: Lose precision
+        return TIMESTAMP();
       case thrift::Type::type::FLOAT:
         return REAL();
       case thrift::Type::type::DOUBLE:
diff --git a/velox/dwio/parquet/reader/TimestampColumnReader.h b/velox/dwio/parquet/reader/TimestampColumnReader.h
@@ -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.
+ */
+
+#pragma once
+
+#include "velox/dwio/parquet/reader/IntegerColumnReader.h"
+#include "velox/dwio/parquet/reader/ParquetColumnReader.h"
+
+namespace facebook::velox::parquet {
+
+class TimestampColumnReader : public IntegerColumnReader {
+ public:
+  TimestampColumnReader(
+      const std::shared_ptr<const dwio::common::TypeWithId>& nodeType,
+      ParquetParams& params,
+      common::ScanSpec& scanSpec)
+      : IntegerColumnReader(nodeType, nodeType, params, scanSpec) {}
+
+  void read(
+      vector_size_t offset,
+      RowSet rows,
+      const uint64_t* /*incomingNulls*/) override {
+    auto& data = formatData_->as<ParquetData>();
+    // Use int128_t as a workaroud. Timestamp in Velox is of 16-byte length.
+    prepareRead<int128_t>(offset, rows, nullptr);
+    readCommon<IntegerColumnReader>(rows);
+  }
+};
+
+} // namespace facebook::velox::parquet
diff --git a/velox/dwio/parquet/tests/examples/timestamp_int96.parquet b/velox/dwio/parquet/tests/examples/timestamp_int96.parquet
diff --git a/velox/dwio/parquet/tests/reader/ParquetTableScanTest.cpp b/velox/dwio/parquet/tests/reader/ParquetTableScanTest.cpp
@@ -72,6 +72,34 @@ class ParquetTableScanTest : public HiveConnectorTestBase {
     assertQuery(plan, splits_, sql);
   }
 
+  void assertSelectWithFilter(
+      std::vector<std::string>&& outputColumnNames,
+      const std::vector<std::string>& subfieldFilters,
+      const std::string& remainingFilter,
+      const std::string& sql,
+      bool isFilterPushdownEnabled) {
+    auto rowType = getRowType(std::move(outputColumnNames));
+    parse::ParseOptions options;
+    options.parseDecimalAsDouble = false;
+
+    auto plan = PlanBuilder(pool_.get())
+                    .setParseOptions(options)
+                    // Function extractFiltersFromRemainingFilter will extract
+                    // filters to subfield filters, but for some types, filter
+                    // pushdown is not supported.
+                    .tableScan(
+                        "hive_table",
+                        rowType,
+                        {},
+                        subfieldFilters,
+                        remainingFilter,
+                        nullptr,
+                        isFilterPushdownEnabled)
+                    .planNode();
+
+    assertQuery(plan, splits_, sql);
+  }
+
   void assertSelectWithAgg(
       std::vector<std::string>&& outputColumnNames,
       const std::vector<std::string>& aggregates,
diff --git a/velox/exec/tests/utils/PlanBuilder.cpp b/velox/exec/tests/utils/PlanBuilder.cpp
@@ -83,7 +83,8 @@ PlanBuilder& PlanBuilder::tableScan(
     const std::unordered_map<std::string, std::string>& columnAliases,
     const std::vector<std::string>& subfieldFilters,
     const std::string& remainingFilter,
-    const RowTypePtr& dataColumns) {
+    const RowTypePtr& dataColumns,
+    bool isFilterPushdownEnabled) {
   std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
       assignments;
   std::unordered_map<std::string, core::TypedExprPtr> typedMapping;
@@ -143,7 +144,7 @@ PlanBuilder& PlanBuilder::tableScan(
   auto tableHandle = std::make_shared<HiveTableHandle>(
       kHiveConnectorId,
       tableName,
-      true,
+      isFilterPushdownEnabled,
       std::move(filters),
       remainingFilterExpr,
       dataColumns);
diff --git a/velox/exec/tests/utils/PlanBuilder.h b/velox/exec/tests/utils/PlanBuilder.h
@@ -144,7 +144,8 @@ class PlanBuilder {
       const std::unordered_map<std::string, std::string>& columnAliases = {},
       const std::vector<std::string>& subfieldFilters = {},
       const std::string& remainingFilter = "",
-      const RowTypePtr& dataColumns = nullptr);
+      const RowTypePtr& dataColumns = nullptr,
+      bool isFilterPushdownEnabled = true);
 
   /// Add a TableScanNode using a connector-specific table handle and
   /// assignments. Supports any connector, not just Hive connector.
diff --git a/velox/type/Type.h b/velox/type/Type.h
@@ -47,6 +47,11 @@ namespace facebook::velox {
 using int128_t = __int128_t;
 using int256_t = boost::multiprecision::int256_t;
 
+struct __attribute__((__packed__)) Int96Timestamp {
+  int32_t days;
+  uint64_t nanos;
+};
+
 /// Velox type system supports a small set of SQL-compatible composeable types:
 /// BOOLEAN, TINYINT, SMALLINT, INTEGER, BIGINT, HUGEINT, REAL, DOUBLE, VARCHAR,
 /// VARBINARY, TIMESTAMP, ARRAY, MAP, ROW
diff --git a/velox/vector/arrow/Bridge.cpp b/velox/vector/arrow/Bridge.cpp
@@ -247,7 +247,7 @@ const char* exportArrowFormatStr(
     case TypeKind::TIMESTAMP:
       // TODO: need to figure out how we'll map this since in Velox we currently
       // store timestamps as two int64s (epoch in sec and nanos).
-      return "ttn"; // time64 [nanoseconds]
+      return "tsu:"; // timestamp [microseconds]
     // Complex/nested types.
     case TypeKind::ARRAY:
       static_assert(sizeof(vector_size_t) == 4);
@@ -366,6 +366,7 @@ void exportValues(
   out.n_buffers = 2;
   // Short decimals need to be converted to 128 bit values as they are mapped
   // to Arrow Decimal128.
+  // Timestamps need to be converted to micros.
   if (!rows.changed() && !vec.type()->isShortDecimal()) {
     holder.setBuffer(1, vec.values());
     return;
@@ -380,6 +381,33 @@ void exportValues(
   holder.setBuffer(1, values);
 }
 
+void exportTimestamps(
+    const BaseVector& vec,
+    const Selection& rows,
+    ArrowArray& out,
+    memory::MemoryPool* pool,
+    VeloxToArrowBridgeHolder& holder) {
+  out.n_buffers = 2;
+  auto size = vec.type()->cppSizeInBytes();
+  auto values = AlignedBuffer::allocate<uint8_t>(
+      checkedMultiply<size_t>(out.length, size), pool);
+  const Buffer& buf = *vec.values();
+  const auto& tsSrc = buf.as<Timestamp>();
+  Buffer& outBuffer = *values;
+  auto dst = outBuffer.asMutable<uint8_t>();
+  vector_size_t j = 0; // index into dst
+  rows.apply([&](vector_size_t i) {
+    int64_t value = 0;
+    if (!vec.mayHaveNulls() || !vec.isNullAt(i)) {
+      // Prevent null here because the use of toMicros on null causes integer
+      // overflow.
+      value = tsSrc[i].toMicros();
+    }
+    memcpy(dst + (j++) * sizeof(int64_t), &value, sizeof(int64_t));
+  });
+  holder.setBuffer(1, values);
+}
+
 void exportStrings(
     const FlatVector<StringView>& vec,
     const Selection& rows,
@@ -432,6 +460,9 @@ void exportFlat(
     case TypeKind::DOUBLE:
       exportValues(vec, rows, out, pool, holder);
       break;
+    case TypeKind::TIMESTAMP:
+      exportTimestamps(vec, rows, out, pool, holder);
+      break;
     case TypeKind::VARCHAR:
     case TypeKind::VARBINARY:
       exportStrings(
@@ -729,6 +760,11 @@ void exportToArrow(const VectorPtr& vec, ArrowSchema& arrowSchema) {
 TypePtr importFromArrow(const ArrowSchema& arrowSchema) {
   const char* format = arrowSchema.format;
   VELOX_CHECK_NOT_NULL(format);
+  const std::string formatStr(format);
+  // TODO: Timezone and unit are not handled.
+  if (formatStr.rfind("ts", 0) == 0) {
+    return TIMESTAMP();
+  }
 
   switch (format[0]) {
     case 'b':
@@ -946,6 +982,34 @@ VectorPtr createStringFlatVector(
       optionalNullCount(nullCount));
 }
 
+VectorPtr createTimestampVector(
+    memory::MemoryPool* pool,
+    const TypePtr& type,
+    BufferPtr nulls,
+    const ArrowSchema& arrowSchema,
+    const ArrowArray& arrowArray,
+    WrapInBufferViewFunc wrapInBufferView) {
+  auto valueBuf = wrapInBufferView(
+      arrowArray.buffers[1], arrowArray.length * sizeof(Timestamp));
+
+  auto src = valueBuf->as<uint8_t>();
+
+  VectorPtr base = BaseVector::create(type, arrowArray.length, pool);
+  base->setNulls(nulls);
+
+  auto flatVector = std::dynamic_pointer_cast<FlatVector<Timestamp>>(base);
+
+  for (int i = 0; i < arrowArray.length; i++) {
+    if (!base->isNullAt(i)) {
+      int64_t result;
+      memcpy(&result, src + i * sizeof(int64_t), sizeof(int64_t));
+      flatVector->set(i, Timestamp::fromMicros(result));
+    }
+  }
+
+  return flatVector;
+}
+
 VectorPtr importFromArrowImpl(
     ArrowSchema& arrowSchema,
     ArrowArray& arrowArray,
@@ -1152,6 +1216,12 @@ VectorPtr importFromArrowImpl(
     return createMapVector(
         pool, type, nulls, arrowSchema, arrowArray, isViewer, wrapInBufferView);
   }
+
+  if (type->kind() == TypeKind::TIMESTAMP) {
+    return createTimestampVector(
+        pool, type, nulls, arrowSchema, arrowArray, wrapInBufferView);
+  }
+
   // Other primitive types.
   VELOX_CHECK(
       type->isPrimitiveType(),
diff --git a/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp b/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp
@@ -661,7 +661,7 @@ TEST_F(ArrowBridgeArrayExportTest, unsupported) {
 
   // Timestamps.
   vector = vectorMaker_.flatVectorNullable<Timestamp>({});
-  EXPECT_THROW(exportToArrow(vector, arrowArray, pool_.get()), VeloxException);
+  exportToArrow(vector, arrowArray, pool_.get());
 
   // Constant encoding.
   vector = BaseVector::createConstant(INTEGER(), variant(10), 10, pool_.get());
diff --git a/velox/vector/arrow/tests/ArrowBridgeSchemaTest.cpp b/velox/vector/arrow/tests/ArrowBridgeSchemaTest.cpp
@@ -120,7 +120,7 @@ TEST_F(ArrowBridgeSchemaExportTest, scalar) {
   testScalarType(VARCHAR(), "u");
   testScalarType(VARBINARY(), "z");
 
-  testScalarType(TIMESTAMP(), "ttn");
+  testScalarType(TIMESTAMP(), "tsu:");
   testScalarType(DATE(), "tdD");
 
   testScalarType(DECIMAL(10, 4), "d:10,4");
