Better tests

apache · Jul 8, 2022 · bc45b9e · bc45b9e
1 parent b9322fc
commit bc45b9e
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 14 deletions.
diff --git a/parquet/src/arrow/arrow_reader.rs b/parquet/src/arrow/arrow_reader.rs
@@ -1532,7 +1532,6 @@ mod tests {
     #[test]
     fn test_row_group_exact_multiple() {
         let schema = Arc::new(Schema::new(vec![
-            Field::new("int", ArrowDataType::Int32, false),
             Field::new(
                 "list",
                 ArrowDataType::List(Box::new(Field::new(
@@ -1557,16 +1556,13 @@ mod tests {
         )
         .unwrap();
         for _ in 0..2 {
-            let mut int_builder = Int32Builder::new(10);
             let mut list_builder = ListBuilder::new(Int32Builder::new(10));
-            for i in 0..10 {
-                int_builder.append_value(i).unwrap();
+            for _ in 0..10 {
                 list_builder.append(true).unwrap();
             }
             let batch = RecordBatch::try_new(
                 schema.clone(),
                 vec![
-                    Arc::new(int_builder.finish()),
                     Arc::new(list_builder.finish()),
                 ],
             )

diff --git a/parquet/src/arrow/record_reader/mod.rs b/parquet/src/arrow/record_reader/mod.rs
@@ -361,7 +361,7 @@ where
                 // If reached end of column chunk => end of a record
                 if records_read != records_to_read
                     && end_of_column
-                    && self.values_written != 0
+                    && self.values_written != self.num_values
                 {
                     records_read += 1;
                     end_of_last_record = self.values_written;
@@ -772,15 +772,21 @@ mod tests {
         assert_eq!(record_reader.num_records(), 4);
         assert_eq!(record_reader.num_values(), 7);
 
-        let mut record_reader = RecordReader::<Int32Type>::new(desc);
-        let page_reader = Box::new(InMemoryPageReader::new(vec![page]));
-        record_reader.set_page_reader(page_reader).unwrap();
-        assert_eq!(record_reader.read_records(3).unwrap(), 3);
-        assert_eq!(record_reader.num_records(), 3);
-        assert_eq!(record_reader.num_values(), 5);
-
-        assert_eq!(record_reader.read_records(3).unwrap(), 1);
+        assert_eq!(record_reader.read_records(4).unwrap(), 0);
         assert_eq!(record_reader.num_records(), 4);
         assert_eq!(record_reader.num_values(), 7);
+
+        record_reader.read_records(4).unwrap();
+
+        let page_reader = Box::new(InMemoryPageReader::new(vec![page.clone()]));
+        record_reader.set_page_reader(page_reader).unwrap();
+
+        assert_eq!(record_reader.read_records(4).unwrap(), 4);
+        assert_eq!(record_reader.num_records(), 8);
+        assert_eq!(record_reader.num_values(), 14);
+
+        assert_eq!(record_reader.read_records(4).unwrap(), 0);
+        assert_eq!(record_reader.num_records(), 8);
+        assert_eq!(record_reader.num_values(), 14);
     }
 }