Draft: test non exact batch size

Author: zhuqi-lucas

Cargo.lock

@@ -229,3 +229,21 @@ unexpected_cfgs = { level = "warn", check-cfg = [
     "cfg(tarpaulin_include)",
 ] }
 unused_qualifications = "deny"
+
+
+# patch branch to use testing version of arrow
+# pin to zhuqi-lucas/support_exact_size_config branch
+[patch.crates-io]
+arrow = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+arrow-array = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+arrow-buffer = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+arrow-cast = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+arrow-data = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+arrow-ipc = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+arrow-schema = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+arrow-select = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+arrow-string = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+arrow-ord = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+arrow-flight = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+parquet = { git = "https://github.com/zhuqi-lucas/arrow-rs.git", rev = "60a239ae614568bacff819494b9332c98052badc" }
+

Cargo.toml

@@ -60,7 +60,7 @@ impl LimitedBatchCoalescer {
         fetch: Option<usize>,
     ) -> Self {
         Self {
-            inner: BatchCoalescer::new(schema, target_batch_size),
+            inner: BatchCoalescer::new(schema, target_batch_size).with_exact_size(false),
             total_rows: 0,
             fetch,
             finished: false,
@@ -156,7 +156,7 @@ mod tests {
             .with_batches(std::iter::repeat_n(batch, 10))
             // expected output is batches of exactly 21 rows (except for the final batch)
             .with_target_batch_size(21)
-            .with_expected_output_sizes(vec![21, 21, 21, 17])
+            .with_expected_output_sizes(vec![24, 24, 24, 8])
             .run()
     }
 
@@ -169,7 +169,7 @@ mod tests {
             // expected to behave the same as `test_concat_batches`
             .with_target_batch_size(21)
             .with_fetch(Some(100))
-            .with_expected_output_sizes(vec![21, 21, 21, 17])
+            .with_expected_output_sizes(vec![24, 24, 24, 8])
             .run();
     }
 
@@ -181,7 +181,7 @@ mod tests {
             // input is 10 batches x 8 rows (80 rows) with fetch limit of 50
             .with_target_batch_size(21)
             .with_fetch(Some(50))
-            .with_expected_output_sizes(vec![21, 21, 8])
+            .with_expected_output_sizes(vec![24, 24, 2])
             .run();
     }
 
@@ -191,7 +191,7 @@ mod tests {
         Test::new()
             .with_batches(std::iter::repeat_n(batch, 10))
             // input is 10 batches x 8 rows (80 rows) with fetch limit of 48
-            .with_target_batch_size(24)
+            .with_target_batch_size(21)
             .with_fetch(Some(48))
             .with_expected_output_sizes(vec![24, 24])
             .run();