Added non-ordered projections in IPC

src/io/ipc/read/common.rs

@@ -1,4 +1,4 @@
-use std::collections::VecDeque;
+use std::collections::{HashMap, HashSet, VecDeque};
 use std::io::{Read, Seek};
 use std::sync::Arc;
 
@@ -277,3 +277,44 @@ mod tests {
         )
     }
 }
+
+pub fn prepare_projection(
+    fields: &[Field],
+    mut projection: Vec<usize>,
+) -> (Vec<usize>, HashMap<usize, usize>, Vec<Field>) {
+    assert_eq!(
+        projection.iter().collect::<HashSet<_>>().len(),
+        projection.len(),
+        "The projection on IPC must not contain duplicates"
+    );
+
+    let fields = projection.iter().map(|x| fields[*x].clone()).collect();
+
+    // selected index; index in
+    let sorted_projection = projection
+        .iter()
+        .copied()
+        .enumerate()
+        .map(|x| (x.1, x.0))
+        .collect::<HashMap<_, _>>(); // e.g. [2, 1] -> {2: 0, 1: 1}
+    projection.sort_unstable(); // e.g. [2, 1] -> [1, 2]
+
+    (projection, sorted_projection, fields)
+}
+
+pub fn apply_projection(
+    chunk: Chunk<Arc<dyn Array>>,
+    projection: &[usize],
+    map: &HashMap<usize, usize>,
+) -> Chunk<Arc<dyn Array>> {
+    // re-order according to projection
+    let arrays = chunk.into_arrays();
+    let arrays = projection
+        .iter()
+        .map(|x| {
+            let index = map.get(x).unwrap();
+            arrays[*index].clone()
+        })
+        .collect();
+    Chunk::new(arrays)
+}

src/io/ipc/read/file_async.rs

@@ -1,4 +1,5 @@
 //! Async reader for Arrow IPC files
+use std::collections::HashMap;
 use std::io::SeekFrom;
 use std::sync::Arc;
 
@@ -16,7 +17,7 @@ use crate::datatypes::{Field, Schema};
 use crate::error::{ArrowError, Result};
 use crate::io::ipc::{IpcSchema, ARROW_MAGIC, CONTINUATION_MARKER};
 
-use super::common::{read_dictionary, read_record_batch};
+use super::common::{apply_projection, prepare_projection, read_dictionary, read_record_batch};
 use super::reader::get_serialized_batch;
 use super::schema::fb_to_schema;
 use super::Dictionaries;
@@ -25,8 +26,8 @@ use super::FileMetadata;
 /// Async reader for Arrow IPC files
 pub struct FileStream<'a> {
     stream: BoxStream<'a, Result<Chunk<Arc<dyn Array>>>>,
+    schema: Option<Schema>,
     metadata: FileMetadata,
-    schema: Schema,
 }
 
 impl<'a> FileStream<'a> {
@@ -38,23 +39,15 @@ impl<'a> FileStream<'a> {
     where
         R: AsyncRead + AsyncSeek + Unpin + Send + 'a,
     {
-        let schema = if let Some(projection) = projection.as_ref() {
-            projection.windows(2).for_each(|x| {
-                assert!(
-                    x[0] < x[1],
-                    "IPC projection must be ordered and non-overlapping",
-                )
-            });
-            let fields = projection
-                .iter()
-                .map(|&x| metadata.schema.fields[x].clone())
-                .collect::<Vec<_>>();
-            Schema {
+        let (projection, schema) = if let Some(projection) = projection {
+            let (p, h, fields) = prepare_projection(&metadata.schema.fields, projection);
+            let schema = Schema {
                 fields,
                 metadata: metadata.schema.metadata.clone(),
-            }
+            };
+            (Some((p, h)), Some(schema))
         } else {
-            metadata.schema.clone()
+            (None, None)
         };
 
         let stream = Self::stream(reader, metadata.clone(), projection);
@@ -72,13 +65,13 @@ impl<'a> FileStream<'a> {
 
     /// Get the projected schema from the IPC file.
     pub fn schema(&self) -> &Schema {
-        &self.schema
+        self.schema.as_ref().unwrap_or(&self.metadata.schema)
     }
 
     fn stream<R>(
         mut reader: R,
         metadata: FileMetadata,
-        projection: Option<Vec<usize>>,
+        projection: Option<(Vec<usize>, HashMap<usize, usize>)>,
     ) -> BoxStream<'a, Result<Chunk<Arc<dyn Array>>>>
     where
         R: AsyncRead + AsyncSeek + Unpin + Send + 'a,
@@ -90,11 +83,19 @@ impl<'a> FileStream<'a> {
                 let chunk = read_batch(
                     &mut reader,
                     &metadata,
-                    projection.as_deref(),
+                    projection.as_ref().map(|x| x.0.as_ref()),
                     block,
                     &mut meta_buffer,
                     &mut block_buffer,
                 ).await?;
+
+                let chunk = if let Some((projection, map)) = &projection {
+                    // re-order according to projection
+                    apply_projection(chunk, projection, map)
+                } else {
+                    chunk
+                };
+
                 yield chunk;
             }
         }

src/io/ipc/read/reader.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::convert::TryInto;
 use std::io::{Read, Seek, SeekFrom};
 use std::sync::Arc;
@@ -37,7 +38,7 @@ pub struct FileReader<R: Read + Seek> {
     reader: R,
     metadata: FileMetadata,
     current_block: usize,
-    projection: Option<(Vec<usize>, Schema)>,
+    projection: Option<(Vec<usize>, HashMap<usize, usize>, Schema)>,
     buffer: Vec<u8>,
 }
 
@@ -230,24 +231,13 @@ impl<R: Read + Seek> FileReader<R> {
     /// # Panic
     /// Panics iff the projection is not in increasing order (e.g. `[1, 0]` nor `[0, 1, 1]` are valid)
     pub fn new(reader: R, metadata: FileMetadata, projection: Option<Vec<usize>>) -> Self {
-        if let Some(projection) = projection.as_ref() {
-            projection.windows(2).for_each(|x| {
-                assert!(
-                    x[0] < x[1],
-                    "The projection on IPC must be ordered and non-overlapping"
-                );
-            });
-        }
         let projection = projection.map(|projection| {
-            let fields = projection
-                .iter()
-                .map(|x| metadata.schema.fields[*x].clone())
-                .collect();
+            let (p, h, fields) = prepare_projection(&metadata.schema.fields, projection);
             let schema = Schema {
                 fields,
                 metadata: metadata.schema.metadata.clone(),
             };
-            (projection, schema)
+            (p, h, schema)
         });
         Self {
             reader,
@@ -262,7 +252,7 @@ impl<R: Read + Seek> FileReader<R> {
     pub fn schema(&self) -> &Schema {
         self.projection
             .as_ref()
-            .map(|x| &x.1)
+            .map(|x| &x.2)
             .unwrap_or(&self.metadata.schema)
     }
 
@@ -285,13 +275,21 @@ impl<R: Read + Seek> Iterator for FileReader<R> {
         if self.current_block < self.metadata.blocks.len() {
             let block = self.current_block;
             self.current_block += 1;
-            Some(read_batch(
+            let chunk = read_batch(
                 &mut self.reader,
                 &self.metadata,
                 self.projection.as_ref().map(|x| x.0.as_ref()),
                 block,
                 &mut self.buffer,
-            ))
+            );
+
+            let chunk = if let Some((projection, map, _)) = &self.projection {
+                // re-order according to projection
+                chunk.map(|chunk| apply_projection(chunk, projection, map))
+            } else {
+                chunk
+            };
+            Some(chunk)
         } else {
             None
         }

tests/it/io/ipc/read/file.rs

@@ -157,28 +157,37 @@ fn read_generated_200_compression_zstd() -> Result<()> {
     test_file("2.0.0-compression", "generated_zstd")
 }
 
-fn test_projection(version: &str, file_name: &str, column: usize) -> Result<()> {
+fn test_projection(version: &str, file_name: &str, columns: Vec<usize>) -> Result<()> {
     let testdata = crate::test_util::arrow_test_data();
     let mut file = File::open(format!(
         "{}/arrow-ipc-stream/integration/{}/{}.arrow_file",
         testdata, version, file_name
     ))?;
 
     let metadata = read_file_metadata(&mut file)?;
-    let mut reader = FileReader::new(&mut file, metadata, Some(vec![column]));
 
-    assert_eq!(reader.schema().fields.len(), 1);
+    let expected = columns
+        .iter()
+        .copied()
+        .map(|x| metadata.schema.fields[x].clone())
+        .collect::<Vec<_>>();
+
+    let mut reader = FileReader::new(&mut file, metadata, Some(columns));
+
+    assert_eq!(reader.schema().fields, expected);
 
     reader.try_for_each(|rhs| {
-        assert_eq!(rhs?.arrays().len(), 1);
+        assert_eq!(rhs?.arrays().len(), expected.len());
         Result::Ok(())
     })?;
     Ok(())
 }
 
 #[test]
 fn read_projected() -> Result<()> {
-    test_projection("1.0.0-littleendian", "generated_primitive", 1)?;
-    test_projection("1.0.0-littleendian", "generated_dictionary", 2)?;
-    test_projection("1.0.0-littleendian", "generated_nested", 0)
+    test_projection("1.0.0-littleendian", "generated_primitive", vec![1])?;
+    test_projection("1.0.0-littleendian", "generated_dictionary", vec![2])?;
+    test_projection("1.0.0-littleendian", "generated_nested", vec![0])?;
+
+    test_projection("1.0.0-littleendian", "generated_primitive", vec![2, 1])
 }