You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2022/07/15 16:18:51 UTC

[GitHub] [arrow-rs] alamb commented on a diff in pull request #2032: Add more tests of RecordReader Batch Size Edge Cases (#2025)

alamb commented on code in PR #2032:
URL: https://github.com/apache/arrow-rs/pull/2032#discussion_r922316988


##########
parquet/src/arrow/arrow_reader.rs:
##########
@@ -1564,9 +1563,25 @@ mod tests {
         writer.close().unwrap();
 
         let mut file_reader = ParquetFileArrowReader::try_new(Bytes::from(buf)).unwrap();
-        let mut record_reader = file_reader.get_record_reader(8).unwrap();
-        assert_eq!(8, record_reader.next().unwrap().unwrap().num_rows());
-        assert_eq!(8, record_reader.next().unwrap().unwrap().num_rows());
-        assert_eq!(4, record_reader.next().unwrap().unwrap().num_rows());
+        let mut record_reader = file_reader.get_record_reader(batch_size).unwrap();
+        assert_eq!(
+            batch_size,
+            record_reader.next().unwrap().unwrap().num_rows()
+        );
+        assert_eq!(
+            batch_size,
+            record_reader.next().unwrap().unwrap().num_rows()
+        );
+    }
+
+    #[test]
+    fn test_row_group_exact_multiple() {
+        use crate::arrow::record_reader::MIN_BATCH_SIZE;
+        test_row_group_batch(8, 8);

Review Comment:
   👍 love it



##########
parquet/src/arrow/arrow_reader.rs:
##########
@@ -1564,9 +1563,25 @@ mod tests {
         writer.close().unwrap();
 
         let mut file_reader = ParquetFileArrowReader::try_new(Bytes::from(buf)).unwrap();
-        let mut record_reader = file_reader.get_record_reader(8).unwrap();
-        assert_eq!(8, record_reader.next().unwrap().unwrap().num_rows());
-        assert_eq!(8, record_reader.next().unwrap().unwrap().num_rows());
-        assert_eq!(4, record_reader.next().unwrap().unwrap().num_rows());
+        let mut record_reader = file_reader.get_record_reader(batch_size).unwrap();
+        assert_eq!(
+            batch_size,
+            record_reader.next().unwrap().unwrap().num_rows()
+        );
+        assert_eq!(
+            batch_size,
+            record_reader.next().unwrap().unwrap().num_rows()
+        );
+    }
+
+    #[test]
+    fn test_row_group_exact_multiple() {
+        use crate::arrow::record_reader::MIN_BATCH_SIZE;
+        test_row_group_batch(8, 8);

Review Comment:
   ```suggestion
           test_row_group_batch(8, 8);
           test_row_group_batch(10, 8);
           test_row_group_batch(8, 10);
   ```
   
   Would it be worth keeping the `10` and `8` in the current test? As well as 8 and 10 for good measure?



-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org