wgtmac commented on code in PR #15240:
URL: https://github.com/apache/arrow/pull/15240#discussion_r1065955959
##########
cpp/src/parquet/arrow/arrow_reader_writer_test.cc:
##########
@@ -4727,5 +4727,111 @@ std::vector<NestedFilterTestCase>
GenerateMapFilteredTestCases() {
INSTANTIATE_TEST_SUITE_P(MapFilteredReads, TestNestedSchemaFilteredReader,
::testing::ValuesIn(GenerateMapFilteredTestCases()));
+template <typename TestType>
+class TestBufferedParquetIO : public TestParquetIO<TestType> {
+ public:
+ void WriteBufferedFile(const std::shared_ptr<Array>& values, int64_t
batch_size) {
+ std::shared_ptr<GroupNode> schema =
+ MakeSimpleSchema(*values->type(), Repetition::OPTIONAL);
+ SchemaDescriptor descriptor;
+ ASSERT_NO_THROW(descriptor.Init(schema));
+ std::shared_ptr<::arrow::Schema> arrow_schema;
+ ArrowReaderProperties props;
+ ASSERT_OK_NO_THROW(FromParquetSchema(&descriptor, props, &arrow_schema));
+
+ std::unique_ptr<FileWriter> writer;
+ ASSERT_OK_NO_THROW(FileWriter::Make(::arrow::default_memory_pool(),
+ this->MakeWriter(schema), arrow_schema,
+ default_arrow_writer_properties(),
&writer));
+ for (int i = 0; i < 4; i++) {
+ if (i % 2 == 0) {
+ ASSERT_OK_NO_THROW(writer->NewBufferedRowGroup());
+ }
+ std::shared_ptr<Array> sliced_array = values->Slice(i * batch_size,
batch_size);
+ std::vector<std::shared_ptr<Array>> arrays = {sliced_array};
+ auto batch = ::arrow::RecordBatch::Make(arrow_schema, batch_size,
arrays);
+ ASSERT_OK_NO_THROW(writer->WriteRecordBatch(*batch));
+ }
+ ASSERT_OK_NO_THROW(writer->Close());
+ }
+};
+
+TYPED_TEST_SUITE(TestBufferedParquetIO, TestTypes);
+
+TYPED_TEST(TestBufferedParquetIO, SingleColumnOptionalBufferedWriteSmall) {
+ constexpr int64_t batch_size = SMALL_SIZE / 4;
+ std::shared_ptr<Array> values;
+ ASSERT_OK(NullableArray<TypeParam>(SMALL_SIZE, 10, kDefaultSeed, &values));
+ this->WriteBufferedFile(values, batch_size);
+ ASSERT_NO_FATAL_FAILURE(this->ReadAndCheckSingleColumnFile(*values));
Review Comment:
Changed the writer to respect `WriterProperties.max_row_group_length` in the
`WriterRecordBatch`. Also, I have changed the test case to check number of row
groups created.
--
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]
