Fix

facebookincubator · Nov 20, 2024 · 717711b · 717711b
1 parent 1ac4c1c
commit 717711b
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 12 deletions.
diff --git a/velox/dwio/dwrf/test/ColumnWriterTest.cpp b/velox/dwio/dwrf/test/ColumnWriterTest.cpp
@@ -327,7 +327,8 @@ template <typename T>
 void testDataTypeWriter(
     const TypePtr& type,
     std::vector<std::optional<T>>& data,
-    const uint32_t sequence = 0) {
+    const uint32_t sequence = 0,
+    DwrfFormat format = DwrfFormat::kDwrf) {
   // Generate a seed and randomly shuffle the data
   uint32_t seed = Random::rand32();
   std::shuffle(data.begin(), data.end(), std::default_random_engine(seed));
@@ -340,7 +341,8 @@ void testDataTypeWriter(
   auto dataTypeWithId = TypeWithId::create(type, 1);
 
   // write
-  auto writer = BaseColumnWriter::create(context, *dataTypeWithId, sequence);
+  auto writer = BaseColumnWriter::create(
+      context, *dataTypeWithId, sequence, nullptr, format);
   auto size = data.size();
   auto batch = populateBatch(data, pool.get(), type);
   const size_t stripeCount = 2;
@@ -459,6 +461,7 @@ TEST_F(ColumnWriterTest, TestNullBooleanWriter) {
 }
 
 TEST_F(ColumnWriterTest, testDecimalWriter) {
+  const auto format = DwrfFormat::kOrc;
   auto genShortDecimals = [&](bool hasNull) {
     std::vector<std::optional<int64_t>> shortDecimals;
     for (auto i = 0; i < ITERATIONS; ++i) {
@@ -472,9 +475,9 @@ TEST_F(ColumnWriterTest, testDecimalWriter) {
   };
 
   auto shortValues = genShortDecimals(false);
-  testDataTypeWriter(DECIMAL(10, 2), shortValues);
+  testDataTypeWriter(DECIMAL(10, 2), shortValues, 0 /*sequence*/, format);
   shortValues = genShortDecimals(true);
-  testDataTypeWriter(DECIMAL(10, 2), shortValues);
+  testDataTypeWriter(DECIMAL(10, 2), shortValues, 0 /*sequence*/, format);
 
   auto genLongDecimals = [&](bool hasNull) {
     std::vector<std::optional<int128_t>> longDecimals;
@@ -489,9 +492,9 @@ TEST_F(ColumnWriterTest, testDecimalWriter) {
   };
 
   auto longValues = genLongDecimals(false);
-  testDataTypeWriter(DECIMAL(38, 4), longValues);
+  testDataTypeWriter(DECIMAL(38, 4), longValues, 0 /*sequence*/, format);
   longValues = genLongDecimals(true);
-  testDataTypeWriter(DECIMAL(38, 4), longValues);
+  testDataTypeWriter(DECIMAL(38, 4), longValues, 0 /*sequence*/, format);
 }
 
 TEST_F(ColumnWriterTest, TestTimestampEpochWriter) {

diff --git a/velox/dwio/dwrf/writer/ColumnWriter.cpp b/velox/dwio/dwrf/writer/ColumnWriter.cpp
@@ -2098,7 +2098,8 @@ std::unique_ptr<BaseColumnWriter> BaseColumnWriter::create(
     WriterContext& context,
     const TypeWithId& type,
     uint32_t sequence,
-    std::function<void(IndexBuilder&)> onRecordPosition) {
+    std::function<void(IndexBuilder&)> onRecordPosition,
+    DwrfFormat format) {
   const auto flatMapEnabled = context.getConfig(Config::FLATTEN_MAP) &&
       type.parent() != nullptr && (type.parent()->id() == 0);
   bool isFlatMapColumn{false};
@@ -2149,16 +2150,19 @@ std::unique_ptr<BaseColumnWriter> BaseColumnWriter::create(
       return std::make_unique<IntegerColumnWriter<int32_t>>(
           context, type, sequence, onRecordPosition);
     case TypeKind::BIGINT: {
-      if (type.type()->isDecimal()) {
+      if (format == DwrfFormat::kOrc && type.type()->isDecimal()) {
         return std::make_unique<DecimalColumnWriter>(
             context, type, sequence, onRecordPosition);
       }
       return std::make_unique<IntegerColumnWriter<int64_t>>(
           context, type, sequence, onRecordPosition);
     }
-    case TypeKind::HUGEINT:
-      return std::make_unique<DecimalColumnWriter>(
-          context, type, sequence, onRecordPosition);
+    case TypeKind::HUGEINT: {
+      if (format == DwrfFormat::kOrc) {
+        return std::make_unique<DecimalColumnWriter>(
+            context, type, sequence, onRecordPosition);
+      }
+    }
     case TypeKind::REAL:
       return std::make_unique<FloatColumnWriter<float>>(
           context, type, sequence, onRecordPosition);

diff --git a/velox/dwio/dwrf/writer/ColumnWriter.h b/velox/dwio/dwrf/writer/ColumnWriter.h
@@ -151,7 +151,8 @@ class BaseColumnWriter : public ColumnWriter {
       WriterContext& context,
       const dwio::common::TypeWithId& type,
       const uint32_t sequence = 0,
-      std::function<void(IndexBuilder&)> onRecordPosition = nullptr);
+      std::function<void(IndexBuilder&)> onRecordPosition = nullptr,
+      DwrfFormat format = DwrfFormat::kDwrf);
 
  protected:
   BaseColumnWriter(

diff --git a/velox/dwio/dwrf/writer/Writer.h b/velox/dwio/dwrf/writer/Writer.h
@@ -44,6 +44,7 @@ struct WriterOptions : public dwio::common::WriterOptions {
       columnWriterFactory;
   const tz::TimeZone* sessionTimezone{nullptr};
   bool adjustTimestampToTimezone{false};
+  DwrfFormat format{DwrfFormat::kDwrf};
 };
 
 class Writer : public dwio::common::Writer {