wjones127 commented on code in PR #15244:
URL: https://github.com/apache/arrow/pull/15244#discussion_r1065082767
##########
cpp/src/parquet/arrow/arrow_reader_writer_test.cc:
##########
@@ -4727,5 +4726,74 @@ std::vector<NestedFilterTestCase>
GenerateMapFilteredTestCases() {
INSTANTIATE_TEST_SUITE_P(MapFilteredReads, TestNestedSchemaFilteredReader,
::testing::ValuesIn(GenerateMapFilteredTestCases()));
+template <typename TestType>
+class TestIntegerAnnotateDecimalTypeParquetIO : public TestParquetIO<TestType>
{
+ public:
+ template <typename ArrayType>
+ void WriteColumn(const std::shared_ptr<ArrayType>& values) {
+ auto arrow_schema = ::arrow::schema({::arrow::field("a", values->type())});
+
+ parquet::WriterProperties::Builder builder;
+ // Enforce integer type to annotate decimal type
+ auto writer_properties =
builder.enable_integer_annotate_decimal()->build();
+ std::shared_ptr<SchemaDescriptor> parquet_schema;
+ ASSERT_OK_NO_THROW(ToParquetSchema(arrow_schema.get(), *writer_properties,
+ *default_arrow_writer_properties(),
+ &parquet_schema));
+
+ this->sink_ = CreateOutputStream();
+ auto schema_node =
std::static_pointer_cast<GroupNode>(parquet_schema->schema_root());
+
+ std::unique_ptr<FileWriter> writer;
+ ASSERT_OK_NO_THROW(FileWriter::Make(
+ ::arrow::default_memory_pool(),
+ ParquetFileWriter::Open(this->sink_, schema_node, writer_properties),
+ arrow_schema, default_arrow_writer_properties(), &writer));
+ ASSERT_OK_NO_THROW(writer->NewRowGroup(values->length()));
+ ASSERT_OK_NO_THROW(writer->WriteColumnChunk(*values));
+ ASSERT_OK_NO_THROW(writer->Close());
+ }
+
+ void ReadAndCheckSingleDecimalColumnFile(const Array& values) {
+ std::shared_ptr<Array> out;
+ std::unique_ptr<FileReader> reader;
+ this->ReaderFromSink(&reader);
+ this->ReadSingleColumnFile(std::move(reader), &out);
+
+ // Reader always read values as DECIMAL128 type
+ ASSERT_EQ(out->type()->id(), ::arrow::Type::DECIMAL128);
+
+ if (values.type()->id() == ::arrow::Type::DECIMAL128) {
+ AssertArraysEqual(values, *out);
+ } else {
+ auto& expected_values = dynamic_cast<const
::arrow::Decimal256Array&>(values);
+ auto& read_values = dynamic_cast<const ::arrow::Decimal128Array&>(*out);
+ ASSERT_EQ(expected_values.null_count(), read_values.null_count());
+ ASSERT_EQ(0, read_values.null_count());
+ ASSERT_EQ(expected_values.length(), read_values.length());
+ for (int64_t i = 0; i < expected_values.length(); ++i) {
+ ASSERT_EQ(::arrow::Decimal256(expected_values.Value(i)).ToString(0),
+ ::arrow::Decimal128(read_values.Value(i)).ToString(0));
+ }
+ }
+ }
+};
+
+typedef ::testing::Types<
+ DecimalWithPrecisionAndScale<1>, DecimalWithPrecisionAndScale<5>,
+ DecimalWithPrecisionAndScale<10>, DecimalWithPrecisionAndScale<18>,
+ Decimal256WithPrecisionAndScale<1>, Decimal256WithPrecisionAndScale<5>,
+ Decimal256WithPrecisionAndScale<10>, Decimal256WithPrecisionAndScale<18>>
+ DecimalTestTypes;
+
+TYPED_TEST_SUITE(TestIntegerAnnotateDecimalTypeParquetIO, DecimalTestTypes);
+
+TYPED_TEST(TestIntegerAnnotateDecimalTypeParquetIO, SingleDecimalColumn) {
Review Comment:
Could we also have a test with null values?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
