Repository: jorgecarleitao/arrow2
Pull request #1205 (2022-08-04 07:33:16)

Improved API of `mmap`
Two changes.

### 1. expose more information on the `FileMetaData`.
I wanted to know how many `Chunks` there are in the IPC file. This also exposes the `size` attribute as that can also be valuable information.

If they are an invariant for unsafe code, I can expose them via functions?


### 2. use Arc<T>

There was an assumption that `Arc<Mmap` implements `Clone + AsRef<[u8]`. It does not. Calling `as_ref` on an `Arc<T> ` gives us `T`, which was opaque to us.

Changed the function signature so that we can pass an `Arc<Mmap>`. 

codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1205?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#1205](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1205?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (8562dec) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/5711cfcbb785450697eb38e221b5dc1825550ee6?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (5711cfc) will **increase** coverage by `0.01%`.
> The diff coverage is `n/a`.

```diff
@@            Coverage Diff             @@
##             main    #1205      +/-   ##
==========================================
+ Coverage   83.53%   83.55%   +0.01%     
==========================================
  Files         368      368              
  Lines       37448    37448              
==========================================
+ Hits        31283    31288       +5     
+ Misses       6165     6160       -5     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1205?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/io/ipc/read/reader.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1205/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL3JlYWRlci5ycw==) | `96.32% <ø> (ø)` | |
| [src/array/utf8/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1205/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3V0ZjgvbW9kLnJz) | `83.64% <0.00%> (+0.92%)` | :arrow_up: |
| [src/io/ipc/read/stream\_async.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1205/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL3N0cmVhbV9hc3luYy5ycw==) | `76.02% <0.00%> (+1.36%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).

ritchie46: I will update the tests tonight.

Patch:
diff --git a/examples/ipc_file_mmap.rs b/examples/ipc_file_mmap.rs
index 038a8eb822..51dd2a2747 100644
--- a/examples/ipc_file_mmap.rs
+++ b/examples/ipc_file_mmap.rs
@@ -8,7 +8,7 @@ use arrow2::mmap::{mmap_dictionaries_unchecked, mmap_unchecked};
 // Arrow2 requires a struct that implements `Clone + AsRef<[u8]>`, which
 // usually `Arc<Mmap>` supports. Here we mock it
 #[derive(Clone)]
-struct Mmap(Arc<Vec<u8>>);
+struct Mmap(Vec<u8>);
 
 impl AsRef<[u8]> for Mmap {
     #[inline]
@@ -19,7 +19,7 @@ impl AsRef<[u8]> for Mmap {
 
 fn main() -> Result<()> {
     // given a mmap
-    let mmap = Mmap(Arc::new(vec![]));
+    let mmap = Arc::new(Mmap(vec![]));
 
     // read the metadata
     let metadata = read::read_file_metadata(&mut std::io::Cursor::new(mmap.as_ref()))?;
diff --git a/src/ffi/mmap.rs b/src/ffi/mmap.rs
index 4a84d9804b..4bd006393f 100644
--- a/src/ffi/mmap.rs
+++ b/src/ffi/mmap.rs
@@ -1,4 +1,5 @@
 use std::collections::VecDeque;
+use std::sync::Arc;
 
 use crate::array::{Array, DictionaryKey, FixedSizeListArray, ListArray, Offset, StructArray};
 use crate::datatypes::DataType;
@@ -85,11 +86,11 @@ fn get_validity<'a>(
 }
 
 fn create_array<
-    T: Clone + AsRef<[u8]>,
+    T: AsRef<[u8]>,
     I: Iterator<Item = Option<*const u8>>,
     II: Iterator<Item = ArrowArray>,
 >(
-    data: T,
+    data: Arc<T>,
     num_rows: usize,
     null_count: usize,
     buffers: I,
@@ -111,7 +112,7 @@ fn create_array<
 
     let dictionary_ptr = dictionary.map(|array| Box::into_raw(Box::new(array)));
 
-    let mut private_data = Box::new(PrivateData::<T> {
+    let mut private_data = Box::new(PrivateData::<Arc<T>> {
         data,
         buffers_ptr,
         children_ptr,
@@ -127,7 +128,7 @@ fn create_array<
         buffers: private_data.buffers_ptr.as_mut_ptr(),
         children: private_data.children_ptr.as_mut_ptr(),
         dictionary: private_data.dictionary_ptr.unwrap_or(std::ptr::null_mut()),
-        release: Some(release::<T>),
+        release: Some(release::<Arc<T>>),
         private_data: Box::into_raw(private_data) as *mut ::std::os::raw::c_void,
     }
 }
@@ -152,8 +153,8 @@ unsafe extern "C" fn release<T>(array: *mut ArrowArray) {
     array.release = None;
 }
 
-fn mmap_binary<O: Offset, T: Clone + AsRef<[u8]>>(
-    data: T,
+fn mmap_binary<O: Offset, T: AsRef<[u8]>>(
+    data: Arc<T>,
     node: &Node,
     block_offset: usize,
     buffers: &mut VecDeque<IpcBuffer>,
@@ -168,7 +169,7 @@ fn mmap_binary<O: Offset, T: Clone + AsRef<[u8]>>(
         .try_into()
         .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?;
 
-    let data_ref = data.as_ref();
+    let data_ref = data.as_ref().as_ref();
 
     let validity = get_validity(data_ref, block_offset, buffers, null_count)?.map(|x| x.as_ptr());
 
@@ -186,8 +187,8 @@ fn mmap_binary<O: Offset, T: Clone + AsRef<[u8]>>(
     ))
 }
 
-fn mmap_fixed_size_binary<T: Clone + AsRef<[u8]>>(
-    data: T,
+fn mmap_fixed_size_binary<T: AsRef<[u8]>>(
+    data: Arc<T>,
     node: &Node,
     block_offset: usize,
     buffers: &mut VecDeque<IpcBuffer>,
@@ -202,7 +203,7 @@ fn mmap_fixed_size_binary<T: Clone + AsRef<[u8]>>(
         .try_into()
         .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?;
 
-    let data_ref = data.as_ref();
+    let data_ref = data.as_ref().as_ref();
 
     let validity = get_validity(data_ref, block_offset, buffers, null_count)?.map(|x| x.as_ptr());
 
@@ -218,8 +219,8 @@ fn mmap_fixed_size_binary<T: Clone + AsRef<[u8]>>(
     ))
 }
 
-fn mmap_null<T: Clone + AsRef<[u8]>>(
-    data: T,
+fn mmap_null<T: AsRef<[u8]>>(
+    data: Arc<T>,
     node: &Node,
     _block_offset: usize,
     _buffers: &mut VecDeque<IpcBuffer>,
@@ -244,8 +245,8 @@ fn mmap_null<T: Clone + AsRef<[u8]>>(
     ))
 }
 
-fn mmap_boolean<T: Clone + AsRef<[u8]>>(
-    data: T,
+fn mmap_boolean<T: AsRef<[u8]>>(
+    data: Arc<T>,
     node: &Node,
     block_offset: usize,
     buffers: &mut VecDeque<IpcBuffer>,
@@ -260,7 +261,7 @@ fn mmap_boolean<T: Clone + AsRef<[u8]>>(
         .try_into()
         .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?;
 
-    let data_ref = data.as_ref();
+    let data_ref = data.as_ref().as_ref();
 
     let validity = get_validity(data_ref, block_offset, buffers, null_count)?.map(|x| x.as_ptr());
 
@@ -280,13 +281,13 @@ fn mmap_boolean<T: Clone + AsRef<[u8]>>(
     ))
 }
 
-fn mmap_primitive<P: NativeType, T: Clone + AsRef<[u8]>>(
-    data: T,
+fn mmap_primitive<P: NativeType, T: AsRef<[u8]>>(
+    data: Arc<T>,
     node: &Node,
     block_offset: usize,
     buffers: &mut VecDeque<IpcBuffer>,
 ) -> Result<ArrowArray, Error> {
-    let data_ref = data.as_ref();
+    let data_ref = data.as_ref().as_ref();
 
     let num_rows: usize = node
         .length()
@@ -313,8 +314,8 @@ fn mmap_primitive<P: NativeType, T: Clone + AsRef<[u8]>>(
 }
 
 #[allow(clippy::too_many_arguments)]
-fn mmap_list<O: Offset, T: Clone + AsRef<[u8]>>(
-    data: T,
+fn mmap_list<O: Offset, T: AsRef<[u8]>>(
+    data: Arc<T>,
     node: &Node,
     block_offset: usize,
     data_type: &DataType,
@@ -335,7 +336,7 @@ fn mmap_list<O: Offset, T: Clone + AsRef<[u8]>>(
         .try_into()
         .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?;
 
-    let data_ref = data.as_ref();
+    let data_ref = data.as_ref().as_ref();
 
     let validity = get_validity(data_ref, block_offset, buffers, null_count)?.map(|x| x.as_ptr());
 
@@ -363,8 +364,8 @@ fn mmap_list<O: Offset, T: Clone + AsRef<[u8]>>(
 }
 
 #[allow(clippy::too_many_arguments)]
-fn mmap_fixed_size_list<T: Clone + AsRef<[u8]>>(
-    data: T,
+fn mmap_fixed_size_list<T: AsRef<[u8]>>(
+    data: Arc<T>,
     node: &Node,
     block_offset: usize,
     data_type: &DataType,
@@ -387,7 +388,7 @@ fn mmap_fixed_size_list<T: Clone + AsRef<[u8]>>(
         .try_into()
         .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?;
 
-    let data_ref = data.as_ref();
+    let data_ref = data.as_ref().as_ref();
 
     let validity = get_validity(data_ref, block_offset, buffers, null_count)?.map(|x| x.as_ptr());
 
@@ -412,8 +413,8 @@ fn mmap_fixed_size_list<T: Clone + AsRef<[u8]>>(
 }
 
 #[allow(clippy::too_many_arguments)]
-fn mmap_struct<T: Clone + AsRef<[u8]>>(
-    data: T,
+fn mmap_struct<T: AsRef<[u8]>>(
+    data: Arc<T>,
     node: &Node,
     block_offset: usize,
     data_type: &DataType,
@@ -434,7 +435,7 @@ fn mmap_struct<T: Clone + AsRef<[u8]>>(
         .try_into()
         .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?;
 
-    let data_ref = data.as_ref();
+    let data_ref = data.as_ref().as_ref();
 
     let validity = get_validity(data_ref, block_offset, buffers, null_count)?.map(|x| x.as_ptr());
 
@@ -466,8 +467,8 @@ fn mmap_struct<T: Clone + AsRef<[u8]>>(
 }
 
 #[allow(clippy::too_many_arguments)]
-fn mmap_dict<K: DictionaryKey, T: Clone + AsRef<[u8]>>(
-    data: T,
+fn mmap_dict<K: DictionaryKey, T: AsRef<[u8]>>(
+    data: Arc<T>,
     node: &Node,
     block_offset: usize,
     _: &DataType,
@@ -486,7 +487,7 @@ fn mmap_dict<K: DictionaryKey, T: Clone + AsRef<[u8]>>(
         .try_into()
         .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?;
 
-    let data_ref = data.as_ref();
+    let data_ref = data.as_ref().as_ref();
 
     let dictionary = dictionaries
         .get(&ipc_field.dictionary_id.unwrap())
@@ -507,8 +508,8 @@ fn mmap_dict<K: DictionaryKey, T: Clone + AsRef<[u8]>>(
     ))
 }
 
-fn get_array<T: Clone + AsRef<[u8]>>(
-    data: T,
+fn get_array<T: AsRef<[u8]>>(
+    data: Arc<T>,
     block_offset: usize,
     data_type: &DataType,
     ipc_field: &IpcField,
@@ -587,8 +588,8 @@ fn get_array<T: Clone + AsRef<[u8]>>(
 }
 
 /// Maps a memory region to an [`Array`].
-pub(crate) unsafe fn mmap<T: Clone + AsRef<[u8]>>(
-    data: T,
+pub(crate) unsafe fn mmap<T: AsRef<[u8]>>(
+    data: Arc<T>,
     block_offset: usize,
     data_type: DataType,
     ipc_field: &IpcField,
diff --git a/src/io/ipc/read/reader.rs b/src/io/ipc/read/reader.rs
index 9fb428eced..37b8437c27 100644
--- a/src/io/ipc/read/reader.rs
+++ b/src/io/ipc/read/reader.rs
@@ -27,13 +27,13 @@ pub struct FileMetadata {
     /// The blocks in the file
     ///
     /// A block indicates the regions in the file to read to get data
-    pub(crate) blocks: Vec<arrow_format::ipc::Block>,
+    pub blocks: Vec<arrow_format::ipc::Block>,
 
     /// Dictionaries associated to each dict_id
     pub(crate) dictionaries: Option<Vec<arrow_format::ipc::Block>>,
 
     /// The total size of the file in bytes
-    pub(crate) size: u64,
+    pub size: u64,
 }
 
 fn read_dictionary_message<R: Read + Seek>(
diff --git a/src/mmap/mod.rs b/src/mmap/mod.rs
index b7822b5f4c..55bd7d01f8 100644
--- a/src/mmap/mod.rs
+++ b/src/mmap/mod.rs
@@ -1,5 +1,6 @@
 //! Memory maps regions defined on the IPC format into [`Array`].
 use std::collections::VecDeque;
+use std::sync::Arc;
 
 use crate::array::Array;
 use crate::chunk::Chunk;
@@ -74,10 +75,10 @@ fn get_buffers_nodes(
     Ok((buffers, field_nodes))
 }
 
-unsafe fn _mmap_record<T: Clone + AsRef<[u8]>>(
+unsafe fn _mmap_record<T: AsRef<[u8]>>(
     fields: &[Field],
     ipc_fields: &[IpcField],
-    data: T,
+    data: Arc<T>,
     batch: RecordBatchRef,
     offset: usize,
     dictionaries: &Dictionaries,
@@ -104,14 +105,14 @@ unsafe fn _mmap_record<T: Clone + AsRef<[u8]>>(
         .and_then(Chunk::try_new)
 }
 
-unsafe fn _mmap_unchecked<T: Clone + AsRef<[u8]>>(
+unsafe fn _mmap_unchecked<T: AsRef<[u8]>>(
     fields: &[Field],
     ipc_fields: &[IpcField],
-    data: T,
+    data: Arc<T>,
     block: Block,
     dictionaries: &Dictionaries,
 ) -> Result<Chunk<Box<dyn Array>>, Error> {
-    let (message, offset) = read_message(data.as_ref(), block)?;
+    let (message, offset) = read_message(data.as_ref().as_ref(), block)?;
     let batch = get_record_batch(message)?;
     _mmap_record(
         fields,
@@ -134,15 +135,15 @@ unsafe fn _mmap_unchecked<T: Clone + AsRef<[u8]>>(
 /// The caller must ensure that `data` contains a valid buffers, for example:
 /// * Offsets in variable-sized containers must be in-bounds and increasing
 /// * Utf8 data is valid
-pub unsafe fn mmap_unchecked<T: Clone + AsRef<[u8]>>(
+pub unsafe fn mmap_unchecked<T: AsRef<[u8]>>(
     metadata: &FileMetadata,
     dictionaries: &Dictionaries,
-    data: T,
+    data: Arc<T>,
     chunk: usize,
 ) -> Result<Chunk<Box<dyn Array>>, Error> {
     let block = metadata.blocks[chunk];
 
-    let (message, offset) = read_message(data.as_ref(), block)?;
+    let (message, offset) = read_message(data.as_ref().as_ref(), block)?;
     let batch = get_record_batch(message)?;
     _mmap_record(
         &metadata.schema.fields,
@@ -154,13 +155,13 @@ pub unsafe fn mmap_unchecked<T: Clone + AsRef<[u8]>>(
     )
 }
 
-unsafe fn mmap_dictionary<T: Clone + AsRef<[u8]>>(
+unsafe fn mmap_dictionary<T: AsRef<[u8]>>(
     metadata: &FileMetadata,
-    data: T,
+    data: Arc<T>,
     block: Block,
     dictionaries: &mut Dictionaries,
 ) -> Result<(), Error> {
-    let (message, offset) = read_message(data.as_ref(), block)?;
+    let (message, offset) = read_message(data.as_ref().as_ref(), block)?;
     let batch = get_dictionary_batch(&message)?;
 
     let id = batch
@@ -205,9 +206,9 @@ unsafe fn mmap_dictionary<T: Clone + AsRef<[u8]>>(
 /// The caller must ensure that `data` contains a valid buffers, for example:
 /// * Offsets in variable-sized containers must be in-bounds and increasing
 /// * Utf8 data is valid
-pub unsafe fn mmap_dictionaries_unchecked<T: Clone + AsRef<[u8]>>(
+pub unsafe fn mmap_dictionaries_unchecked<T: AsRef<[u8]>>(
     metadata: &FileMetadata,
-    data: T,
+    data: Arc<T>,
 ) -> Result<Dictionaries, Error> {
     let blocks = if let Some(blocks) = &metadata.dictionaries {
         blocks


Test patch:
diff --git a/tests/it/io/ipc/mmap.rs b/tests/it/io/ipc/mmap.rs
index 9142227b1f..ef317b9a69 100644
--- a/tests/it/io/ipc/mmap.rs
+++ b/tests/it/io/ipc/mmap.rs
@@ -3,6 +3,7 @@ use arrow2::chunk::Chunk;
 use arrow2::datatypes::{DataType, Field, Schema};
 use arrow2::error::Result;
 use arrow2::io::ipc::read::read_file_metadata;
+use std::sync::Arc;
 
 use super::write::file::write;
 
@@ -10,9 +11,9 @@ fn round_trip(array: Box<dyn Array>) -> Result<()> {
     let schema = Schema::from(vec![Field::new("a", array.data_type().clone(), true)]);
     let columns = Chunk::try_new(vec![array.clone()])?;
 
-    let data = write(&[columns], &schema, None, None)?;
+    let data = Arc::new(write(&[columns], &schema, None, None)?);
 
-    let metadata = read_file_metadata(&mut std::io::Cursor::new(&data))?;
+    let metadata = read_file_metadata(&mut std::io::Cursor::new(data.as_ref()))?;
 
     let dictionaries =
         unsafe { arrow2::mmap::mmap_dictionaries_unchecked(&metadata, data.clone())? };