[arrow-rs] branch master updated: Use IPC row count info in IPC reader (#1796)

[vi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new b5fbd119a Use IPC row count info in IPC reader (#1796)
b5fbd119a is described below

commit b5fbd119a53dfd90d370fd119f463ad196a3df5b
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Mon Jun 6 09:59:53 2022 -0700

    Use IPC row count info in IPC reader (#1796)
    
    * Use IPC row count info
    
    * Add test
    
    * Update arrow/src/ipc/reader.rs
    
    Co-authored-by: Raphael Taylor-Davies <17...@users.noreply.github.com>
    
    Co-authored-by: Raphael Taylor-Davies <17...@users.noreply.github.com>
---
 arrow/src/ipc/reader.rs | 28 +++++++++++++++++++++++++---
 1 file changed, 25 insertions(+), 3 deletions(-)

diff --git a/arrow/src/ipc/reader.rs b/arrow/src/ipc/reader.rs
index 868098327..1ac519382 100644
--- a/arrow/src/ipc/reader.rs
+++ b/arrow/src/ipc/reader.rs
@@ -30,7 +30,7 @@ use crate::compute::cast;
 use crate::datatypes::{DataType, Field, IntervalUnit, Schema, SchemaRef, UnionMode};
 use crate::error::{ArrowError, Result};
 use crate::ipc;
-use crate::record_batch::{RecordBatch, RecordBatchReader};
+use crate::record_batch::{RecordBatch, RecordBatchOptions, RecordBatchReader};
 
 use ipc::CONTINUATION_MARKER;
 use DataType::*;
@@ -608,6 +608,11 @@ pub fn read_record_batch(
     let mut node_index = 0;
     let mut arrays = vec![];
 
+    let options = RecordBatchOptions {
+        row_count: Some(batch.length() as usize),
+        ..Default::default()
+    };
+
     if let Some(projection) = projection {
         // project fields
         for (idx, field) in schema.fields().iter().enumerate() {
@@ -643,7 +648,11 @@ pub fn read_record_batch(
             }
         }
 
-        RecordBatch::try_new(Arc::new(schema.project(projection)?), arrays)
+        RecordBatch::try_new_with_options(
+            Arc::new(schema.project(projection)?),
+            arrays,
+            &options,
+        )
     } else {
         // keep track of index as lists require more than one node
         for field in schema.fields() {
@@ -661,7 +670,7 @@ pub fn read_record_batch(
             buffer_index = triple.2;
             arrays.push(triple.0);
         }
-        RecordBatch::try_new(schema, arrays)
+        RecordBatch::try_new_with_options(schema, arrays, &options)
     }
 }
 
@@ -1933,4 +1942,17 @@ mod tests {
         let output_batch = roundtrip_ipc_stream(&input_batch);
         assert_eq!(input_batch, output_batch);
     }
+
+    #[test]
+    fn test_no_columns_batch() {
+        let schema = Arc::new(Schema::new(vec![]));
+        let options = RecordBatchOptions {
+            match_field_names: true,
+            row_count: Some(10),
+        };
+        let input_batch =
+            RecordBatch::try_new_with_options(schema, vec![], &options).unwrap();
+        let output_batch = roundtrip_ipc_stream(&input_batch);
+        assert_eq!(input_batch, output_batch);
+    }
 }