fix(parquet): Fix inserting decimal type values precision and scale

Author: jkhaliqi

velox/connectors/Connector.h

@@ -160,6 +160,8 @@ class ConnectorInsertTableHandle : public ISerializable {
 
   virtual std::string toString() const = 0;
 
+  virtual const std::vector<TypePtr> getColumnHandleDataTypes() const = 0;
+
   folly::dynamic serialize() const override {
     VELOX_NYI();
   }

velox/connectors/hive/HiveDataSink.h

@@ -337,6 +337,14 @@ class HiveInsertTableHandle : public ConnectorInsertTableHandle {
 
   std::string toString() const override;
 
+  const std::vector<TypePtr> getColumnHandleDataTypes() const override {
+    std::vector<TypePtr> tableTypes;
+    for (const auto& i : inputColumns_) {
+      tableTypes.emplace_back(i->dataType());
+    }
+    return tableTypes;
+  }
+
  protected:
   const std::vector<std::shared_ptr<const HiveColumnHandle>> inputColumns_;
   const std::shared_ptr<const LocationHandle> locationHandle_;

velox/connectors/hive/tests/HiveDataSinkTest.cpp

@@ -1204,6 +1204,37 @@ TEST_F(HiveDataSinkTest, insertTableHandleToString) {
       "HiveInsertTableHandle [dwrf zstd], [inputColumns: [ HiveColumnHandle [name: c0, columnType: Regular, dataType: BIGINT, requiredSubfields: [ ]] HiveColumnHandle [name: c1, columnType: Regular, dataType: INTEGER, requiredSubfields: [ ]] HiveColumnHandle [name: c2, columnType: Regular, dataType: SMALLINT, requiredSubfields: [ ]] HiveColumnHandle [name: c3, columnType: Regular, dataType: REAL, requiredSubfields: [ ]] HiveColumnHandle [name: c4, columnType: Regular, dataType: DOUBLE, requiredSubfields: [ ]] HiveColumnHandle [name: c5, columnType: PartitionKey, dataType: VARCHAR, requiredSubfields: [ ]] HiveColumnHandle [name: c6, columnType: PartitionKey, dataType: BOOLEAN, requiredSubfields: [ ]] ], locationHandle: LocationHandle [targetPath: /path/to/test, writePath: /path/to/test, tableType: kNew, tableFileName: ], bucketProperty: \nHiveBucketProperty[<HIVE_COMPATIBLE 4>\n\tBucket Columns:\n\t\tc5\n\tBucket Types:\n\t\tVARCHAR\n\tSortedBy Columns:\n\t\t[COLUMN[c5] ORDER[DESC NULLS LAST]]\n]\n, fileNameGenerator: HiveInsertFileNameGenerator]");
 }
 
+TEST_F(HiveDataSinkTest, insertTableHandleGetColumnHandleDataTypes) {
+  const int32_t kNumBuckets = 4;
+  auto bucketProperty = std::make_shared<HiveBucketProperty>(
+      HiveBucketProperty::Kind::kHiveCompatible,
+      kNumBuckets,
+      std::vector<std::string>{"c5"},
+      std::vector<TypePtr>{VARCHAR()},
+      std::vector<std::shared_ptr<const HiveSortingColumn>>{
+          std::make_shared<HiveSortingColumn>(
+              "c5", core::SortOrder{false, false})});
+  auto rowType =
+      ROW({"c0", "c5", "c6"}, {DECIMAL(10, 2), VARCHAR(), BOOLEAN()});
+  auto insertTableHandle = createHiveInsertTableHandle(
+      rowType,
+      TempDirectoryPath::create()->getPath(),
+      dwio::common::FileFormat::PARQUET,
+      {"c5", "c6"},
+      bucketProperty);
+  auto [precision, scale] = getDecimalPrecisionScale(
+      *insertTableHandle->getColumnHandleDataTypes()[0]);
+  ASSERT_EQ(precision, 10);
+  ASSERT_EQ(scale, 2);
+  ASSERT_EQ(
+      insertTableHandle->getColumnHandleDataTypes()[0]->toString(),
+      "DECIMAL(10, 2)");
+  ASSERT_EQ(
+      insertTableHandle->getColumnHandleDataTypes()[1]->toString(), "VARCHAR");
+  ASSERT_EQ(
+      insertTableHandle->getColumnHandleDataTypes()[2]->toString(), "BOOLEAN");
+}
+
 #ifdef VELOX_ENABLE_PARQUET
 TEST_F(HiveDataSinkTest, flushPolicyWithParquet) {
   const auto outputDirectory = TempDirectoryPath::create();

velox/dwio/parquet/writer/Writer.cpp

@@ -229,7 +229,18 @@ std::shared_ptr<::arrow::Field> updateFieldNameRecursive(
         arrowMapType->item_field(), *mapType.valueType());
     return newField->WithType(
         ::arrow::map(newKeyField->type(), newValueField->type()));
-  } else if (name != "") {
+  } else if (type.isDecimal()) {
+    // Parquet type is set from the column type rather than inferred from the
+    // field data.
+    auto precisionScale = getDecimalPrecisionScale(type);
+    if (!name.empty()) {
+      auto newField = field->WithName(name);
+      return newField->WithType(
+          ::arrow::decimal(precisionScale.first, precisionScale.second));
+    }
+    return field->WithType(
+        ::arrow::decimal(precisionScale.first, precisionScale.second));
+  } else if (!name.empty()) {
     return field->WithName(name);
   } else {
     return field;
@@ -318,6 +329,72 @@ std::optional<std::string> getParquetCreatedBy(
   return std::nullopt;
 }
 
+bool equivalentWithSimilarDecimal(
+    const Type& dataType,
+    const Type& schemaType) {
+  if (!(typeid(dataType) == typeid(schemaType))) {
+    return false;
+  }
+  if (schemaType.isRow()) {
+    const auto& schemaTyped = schemaType.asRow();
+    if (schemaTyped.size() != dataType.size()) {
+      return false;
+    }
+    for (size_t i = 0; i < dataType.size(); ++i) {
+      if (schemaTyped.childAt(i)->isDecimal()) {
+        auto presicisonScaleSchema =
+            getDecimalPrecisionScale(*schemaTyped.childAt(i));
+        auto precisionScaleData =
+            getDecimalPrecisionScale(*dataType.childAt(i));
+        // Table precision can be larger than or equal to the inserted data, but
+        // the scale must be the same.
+        if (presicisonScaleSchema.first < precisionScaleData.first ||
+            presicisonScaleSchema.second != precisionScaleData.second) {
+          return false;
+        }
+      } else if (!equivalentWithSimilarDecimal(
+                     *dataType.childAt(i), *schemaTyped.childAt(i))) {
+        return false;
+      }
+    }
+  } else if (schemaType.isArray()) {
+    const auto& schemaTyped = schemaType.asArray();
+    const auto& dataTyped = dataType.asArray();
+    if (schemaTyped.childAt(0)->isDecimal()) {
+      auto presicisonScaleSchema =
+          getDecimalPrecisionScale(*schemaTyped.childAt(0));
+      auto precisionScaleData = getDecimalPrecisionScale(*dataType.childAt(0));
+      if (presicisonScaleSchema.first < precisionScaleData.first ||
+          presicisonScaleSchema.second != precisionScaleData.second) {
+        return false;
+      }
+    }
+    return equivalentWithSimilarDecimal(
+        *dataTyped.elementType(), *schemaTyped.elementType());
+  } else if (schemaType.isMap()) {
+    auto& schemaTyped = schemaType.asMap();
+    auto& dataTyped = dataType.asMap();
+    return equivalentWithSimilarDecimal(
+               *dataTyped.keyType(), *schemaTyped.keyType()) &&
+        equivalentWithSimilarDecimal(
+               *dataTyped.valueType(), *schemaTyped.valueType());
+  } else if (schemaType.isFunction()) {
+    auto& schemaTyped = schemaType.asFunction();
+    auto& dataTyped = dataType.asFunction();
+    if (dataTyped.size() != schemaTyped.size()) {
+      return false;
+    }
+    for (auto i = 0; i < schemaTyped.size(); ++i) {
+      if (!equivalentWithSimilarDecimal(
+              *dataTyped.childAt(i), *schemaTyped.childAt(i))) {
+        return false;
+      }
+    }
+    return true;
+  }
+  return true;
+}
+
 } // namespace
 
 Writer::Writer(
@@ -425,7 +502,7 @@ dwio::common::StripeProgress getStripeProgress(
  */
 void Writer::write(const VectorPtr& data) {
   VELOX_USER_CHECK(
-      data->type()->equivalent(*schema_),
+      equivalentWithSimilarDecimal(*data->type(), *schema_),
       "The file schema type should be equal with the input rowvector type.");
 
   VectorPtr exportData = data;

velox/exec/TableWriter.cpp

@@ -72,15 +72,17 @@ TableWriter::TableWriter(
 void TableWriter::setTypeMappings(
     const core::TableWriteNodePtr& tableWriteNode) {
   auto outputNames = tableWriteNode->columnNames();
-  auto outputTypes = tableWriteNode->columns()->children();
+  auto tableTypes = tableWriteNode->insertTableHandle()
+                        ->connectorInsertTableHandle()
+                        ->getColumnHandleDataTypes();
 
   const auto& inputType = tableWriteNode->sources()[0]->outputType();
 
   // Ids that map input to output columns.
-  inputMapping_.reserve(outputTypes.size());
+  inputMapping_.reserve(tableTypes.size());
   std::vector<TypePtr> inputTypes;
 
-  // Generate mappings between input and output types. Note that column names
+  // Generate mappings between input and table types. Note that column names
   // must match, but in some case the types won't, for example, when writing a
   // struct (ROW) as a flat map (MAP).
   for (const auto& name : tableWriteNode->columns()->names()) {
@@ -89,7 +91,7 @@ void TableWriter::setTypeMappings(
     inputTypes.emplace_back(inputType->childAt(idx));
   }
 
-  mappedOutputType_ = ROW(folly::copy(outputNames), std::move(outputTypes));
+  mappedOutputType_ = ROW(folly::copy(outputNames), std::move(tableTypes));
   mappedInputType_ = ROW(std::move(outputNames), std::move(inputTypes));
 }
 

velox/experimental/cudf/connectors/hive/CudfHiveDataSink.h

@@ -253,6 +253,14 @@ class CudfHiveInsertTableHandle : public ConnectorInsertTableHandle {
 
   static void registerSerDe();
 
+  const std::vector<TypePtr> getColumnHandleDataTypes() const override {
+    std::vector<TypePtr> tableTypes;
+    for (const auto& i : inputColumns_) {
+      tableTypes.emplace_back(i->type());
+    }
+    return tableTypes;
+  }
+
   std::string toString() const override;
 
  private: