KarateSnowMachine commented on code in PR #36191:
URL: https://github.com/apache/arrow/pull/36191#discussion_r1236978025
##########
cpp/src/parquet/stream_reader_test.cc:
##########
@@ -947,5 +947,77 @@ TEST_F(TestReadingDataFiles, ByteArrayDecimal) {
EXPECT_EQ(i, 25);
}
+class TestMultiRowGroupStreamReader : public ::testing::Test {
+ public:
+ TestMultiRowGroupStreamReader() { createTestFile(); }
+
+ protected:
+ const char* GetDataFile() const { return
"stream_reader_multirowgroup_test.parquet"; }
+
+ void SetUp() {
+ PARQUET_ASSIGN_OR_THROW(auto infile,
::arrow::io::ReadableFile::Open(GetDataFile()));
+ auto file_reader = parquet::ParquetFileReader::Open(infile);
+ reader_ = StreamReader{std::move(file_reader)};
+ }
+
+ void TearDown() { reader_ = StreamReader{}; }
+
+ std::shared_ptr<schema::GroupNode> GetSchema() {
+ schema::NodeVector fields;
+ fields.push_back(schema::PrimitiveNode::Make("row_group_number",
Repetition::REQUIRED,
+ Type::INT32,
ConvertedType::UINT_16));
+
+ fields.push_back(schema::PrimitiveNode::Make("row_number",
Repetition::REQUIRED,
+ Type::INT64,
ConvertedType::UINT_64));
+
+ return std::static_pointer_cast<schema::GroupNode>(
+ schema::GroupNode::Make("schema", Repetition::REQUIRED, fields));
+ }
+
+ void createTestFile() {
+ PARQUET_ASSIGN_OR_THROW(auto outfile,
+
::arrow::io::FileOutputStream::Open(GetDataFile()));
+
+ auto file_writer = ParquetFileWriter::Open(outfile, GetSchema());
+
+ StreamWriter os{std::move(file_writer)};
+
+ int nrows = 0;
+ for (auto group = 0; group < num_row_groups; ++group) {
+ for (auto i = 0; i < num_rows_per_group; ++i) {
+ os << static_cast<uint16_t>(group);
+ os << static_cast<uint64_t>(nrows);
+ os << EndRow;
+ nrows++;
+ }
+ os.EndRowGroup();
+ }
+ }
+
+ StreamReader reader_;
+ static constexpr int num_row_groups = 5;
+ static constexpr int num_rows_per_group = 10;
+};
+
+TEST_F(TestMultiRowGroupStreamReader, SkipRows) {
+ // skip somewhere in the middle that is a few row groups in to the file
+ auto num_rows_to_skip = 33;
+
+ auto retval = reader_.SkipRows(num_rows_to_skip);
Review Comment:
Good idea -- done. Tried to cover the potential corner cases:
middle group -> middle group
middle group -> start of next group
middle group -> past end of file
--
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]
