Repository: jorgecarleitao/arrow2
Pull request #302 (2021-08-19 14:32:26)

Fixed error in reading a non-finished IPC stream.
When the stream reader gets called `.next()` and the stream has not finished, we should not error nor return `None` (end of stream), and instead should offer the state `Waiting`, so that the user can decide what to do if no new batches have been observed, but no finished state has also been found.

This PR changes the return state of the stream reader returns `Option<Result<State>>` where `State`:

```rust
pub enum State {
    Waiting,
    Some(RecordBatch),
}
```

* `None` describes the end of the stream
* `Some(Err)` describes an error
* `Some(Ok(State::Waiting))` describes that the stream has not finished, but that no new data is available to read
* `Some(Ok(State::Some(_))` describes a new batch

Thanks a lot to @HagaiHargil for clarifying the limitations of the stream reader.

# Backwards incompatible changes:

* `StreamReader` no longer implements `RecordBatchReader`
* `Iterator` implementation of `StreamReader` no longer returns `Option<Result<RecordBatch>>` and instead returns `Option<Result<State>>`. Use `State::unwrap` if you are certain that the stream contains a batch (and is not waiting).


codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#302](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (9ba69ba) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/bbe607ccbf3d7079c1fac6ad53fc6e5d4368e165?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (bbe607c) will **increase** coverage by `0.00%`.
> The diff coverage is `68.75%`.

> :exclamation: Current head 9ba69ba differs from pull request most recent head f6775a6. Consider uploading reports for the commit f6775a6 to get more accurate results
[![Impacted file tree graph](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302/graphs/tree.svg?width=650&height=150&src=pr&token=AgyTF60R3D&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)

```diff
@@           Coverage Diff           @@
##             main     #302   +/-   ##
=======================================
  Coverage   80.52%   80.52%           
=======================================
  Files         323      323           
  Lines       21135    21140    +5     
=======================================
+ Hits        17019    17024    +5     
  Misses       4116     4116           
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302?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/stream.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL3N0cmVhbS5ycw==) | `82.92% <54.54%> (+0.64%)` | :arrow_up: |
| [tests/it/io/ipc/common.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-dGVzdHMvaXQvaW8vaXBjL2NvbW1vbi5ycw==) | `100.00% <100.00%> (ø)` | |
| [tests/it/io/ipc/read/stream.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-dGVzdHMvaXQvaW8vaXBjL3JlYWQvc3RyZWFtLnJz) | `100.00% <100.00%> (ø)` | |
| [tests/it/io/ipc/write/stream.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-dGVzdHMvaXQvaW8vaXBjL3dyaXRlL3N0cmVhbS5ycw==) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Last update [bbe607c...f6775a6](https://codecov.io/gh/jorgecarleitao/arrow2/pull/302?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).

HagaiHargil: Hey,

I can confirm that this works as intended for me! TBH, I still had core dumps due to memory allocations, but once I switched over to sockets this resolved itself.

The nice thing about the new `StreamState` enum is that my app is now naturally "locked in" to the timing of the original stream - when its event rate is high my app responds faster, and when it's lower I'm not needlessly sampling the socket.

With respect to docs - it's always hard to convince the Arrow mailing list that something in unclear in their docs for some reason. Perhaps you could add a warning here that data should not be simultaneously written and read from simple files, but only from sockets? It's kinda obvious, but the Arrow docs hint that it is possible using their (Pythonic) file format.

Thanks again!
jorgecarleitao: Awesome! Could you share or PR the socket solution? I added an example in this PR demonstrating files, but if we change the example to use sockets like you are doing, maybe it is more obvious what users should use?

HagaiHargil: Not sure what's the best way to "merge" my proposed changes, so I'll just add them here:

 examples/ipc_pyarrow/main.py:

```python
import pyarrow as pa
from time import sleep
import socket

# Set up the data exchange socket
HOST = 127.0.0.1
PORT = 12989
sk = socket.socket()
sk.bind((HOST, PORT))
sk.listen()

data = [
    pa.array([1, 2, 3, 4]),
    pa.array(["foo", "bar", "baz", None]),
    pa.array([True, None, False, True]),
]

batch = pa.record_batch(data, names=["f0", "f1", "f2"])

# Accept incoming connection and stream the data away
connection, address = sk.accept()
dummy_socket_file = connection.makefile('wb')
writer = pa.ipc.new_stream(dummy_socket_file, batch.schema)
while True:
    for _ in range(10):
        writer.write(batch)
    sleep(1)
```

examples/ipc_pyarrow/src/main.rs:

```rust
use std::net::TcpStream;
use std::thread;
use std::time::Duration;

use arrow2::array::{Array, Int64Array};
use arrow2::datatypes::DataType;
use arrow2::error::Result;
use arrow2::io::ipc::read;

fn main() -> Result<()> {
    const ADDRESS: &str = "127.0.0.1:12989";

    let mut reader = TcpStream::connect(ADDRESS)?;
    let metadata = read::read_stream_metadata(&mut reader)?;
    let mut stream = read::StreamReader::new(&mut reader, metadata);

    let mut idx = 0;
    loop {
        match stream.next() {
            Some(x) => match x {
                Ok(read::StreamState::Some(b)) => {
                    idx += 1;
                    println!("batch: {:?}", idx)
                }
                Ok(read::StreamState::Waiting) => thread::sleep(Duration::from_millis(4000)),
                Err(l) => println!("{:?} ({})", l, idx),
            },
            None => break,
        };
    }

    Ok(())
}
```
jorgecarleitao: Awesome, I incorporated it in the example; thank you for taking the time. I aim to release this on the v0.4, so that you do not need to depend on the github version.

Patch:
diff --git a/examples/ipc_pyarrow/.gitignore b/examples/ipc_pyarrow/.gitignore
new file mode 100644
index 0000000000..7d92770a26
--- /dev/null
+++ b/examples/ipc_pyarrow/.gitignore
@@ -0,0 +1,1 @@
+data.arrows
diff --git a/examples/ipc_pyarrow/Cargo.toml b/examples/ipc_pyarrow/Cargo.toml
new file mode 100644
index 0000000000..47bc82181b
--- /dev/null
+++ b/examples/ipc_pyarrow/Cargo.toml
@@ -0,0 +1,7 @@
+[package]
+name = "ipc_stream"
+version = "0.1.0"
+edition = "2018"
+
+[dependencies]
+arrow2 = { path = "../../", default-features = false, features = ["io_ipc"] }
diff --git a/examples/ipc_pyarrow/main.py b/examples/ipc_pyarrow/main.py
new file mode 100644
index 0000000000..c3b1468dea
--- /dev/null
+++ b/examples/ipc_pyarrow/main.py
@@ -0,0 +1,24 @@
+import pyarrow as pa
+from time import sleep
+import socket
+
+# Set up the data exchange socket
+sk = socket.socket()
+sk.bind(("127.0.0.1", 12989))
+sk.listen()
+
+data = [
+    pa.array([1, 2, 3, 4]),
+    pa.array(["foo", "bar", "baz", None]),
+    pa.array([True, None, False, True]),
+]
+
+batch = pa.record_batch(data, names=["f0", "f1", "f2"])
+
+# Accept incoming connection and stream the data away
+connection, address = sk.accept()
+dummy_socket_file = connection.makefile("wb")
+with pa.RecordBatchStreamWriter(dummy_socket_file, batch.schema) as writer:
+    for i in range(50):
+        writer.write_batch(batch)
+        sleep(1)
diff --git a/examples/ipc_pyarrow/run.sh b/examples/ipc_pyarrow/run.sh
new file mode 100755
index 0000000000..193fc1c418
--- /dev/null
+++ b/examples/ipc_pyarrow/run.sh
@@ -0,0 +1,7 @@
+python main.py &
+PRODUCER_PID=$!
+
+sleep 1 # wait for metadata to be available.
+cargo run
+
+kill $PRODUCER_PID
diff --git a/examples/ipc_pyarrow/src/main.rs b/examples/ipc_pyarrow/src/main.rs
new file mode 100644
index 0000000000..ce92e4e1b2
--- /dev/null
+++ b/examples/ipc_pyarrow/src/main.rs
@@ -0,0 +1,33 @@
+use std::net::TcpStream;
+use std::thread;
+use std::time::Duration;
+
+use arrow2::array::{Array, Int64Array};
+use arrow2::datatypes::DataType;
+use arrow2::error::Result;
+use arrow2::io::ipc::read;
+
+fn main() -> Result<()> {
+    const ADDRESS: &str = "127.0.0.1:12989";
+
+    let mut reader = TcpStream::connect(ADDRESS)?;
+    let metadata = read::read_stream_metadata(&mut reader)?;
+    let mut stream = read::StreamReader::new(&mut reader, metadata);
+
+    let mut idx = 0;
+    loop {
+        match stream.next() {
+            Some(x) => match x {
+                Ok(read::StreamState::Some(b)) => {
+                    idx += 1;
+                    println!("batch: {:?}", idx)
+                }
+                Ok(read::StreamState::Waiting) => thread::sleep(Duration::from_millis(2000)),
+                Err(l) => println!("{:?} ({})", l, idx),
+            },
+            None => break,
+        };
+    }
+
+    Ok(())
+}
diff --git a/guide/src/SUMMARY.md b/guide/src/SUMMARY.md
index cdf5aec37e..6ba760bec7 100644
--- a/guide/src/SUMMARY.md
+++ b/guide/src/SUMMARY.md
@@ -13,4 +13,5 @@
     - [Read Parquet](./io/parquet_read.md)
     - [Write Parquet](./io/parquet_write.md)
     - [Read Arrow](./io/ipc_read.md)
+    - [Read Arrow stream](./io/ipc_stream_read.md)
     - [Write Arrow](./io/ipc_write.md)
diff --git a/guide/src/io/ipc_stream_read.md b/guide/src/io/ipc_stream_read.md
new file mode 100644
index 0000000000..0ab872ac98
--- /dev/null
+++ b/guide/src/io/ipc_stream_read.md
@@ -0,0 +1,21 @@
+# Read Arrow streams
+
+When compiled with feature `io_ipc`, this crate can be used to read Arrow streams.
+
+The example below shows how to read from a stream:
+
+```rust
+{{#include ../../../examples/ipc_pyarrow/src/main.rs}}
+```
+
+e.g. written by pyarrow:
+
+```python,ignore
+{{#include ../../../examples/ipc_pyarrow/main.py}}
+```
+
+via
+
+```bash,ignore
+{{#include ../../../examples/ipc_pyarrow/run.sh}}
+```
diff --git a/src/io/ipc/read/mod.rs b/src/io/ipc/read/mod.rs
index 8ed127e269..76b139a589 100644
--- a/src/io/ipc/read/mod.rs
+++ b/src/io/ipc/read/mod.rs
@@ -24,4 +24,4 @@ mod stream;
 
 pub use common::{read_dictionary, read_record_batch};
 pub use reader::{read_file_metadata, FileMetadata, FileReader};
-pub use stream::{read_stream_metadata, StreamMetadata, StreamReader};
+pub use stream::{read_stream_metadata, StreamMetadata, StreamReader, StreamState};
diff --git a/src/io/ipc/read/stream.rs b/src/io/ipc/read/stream.rs
index 5018f7953b..ed7f5955df 100644
--- a/src/io/ipc/read/stream.rs
+++ b/src/io/ipc/read/stream.rs
@@ -23,7 +23,7 @@ use gen::Schema::MetadataVersion;
 use crate::array::*;
 use crate::datatypes::Schema;
 use crate::error::{ArrowError, Result};
-use crate::record_batch::{RecordBatch, RecordBatchReader};
+use crate::record_batch::RecordBatch;
 
 use super::super::CONTINUATION_MARKER;
 use super::super::{convert, gen};
@@ -76,12 +76,27 @@ pub fn read_stream_metadata<R: Read>(reader: &mut R) -> Result<StreamMetadata> {
     })
 }
 
+pub enum StreamState {
+    Waiting,
+    Some(RecordBatch),
+}
+
+impl StreamState {
+    pub fn unwrap(self) -> RecordBatch {
+        if let StreamState::Some(batch) = self {
+            batch
+        } else {
+            panic!("The batch is not available")
+        }
+    }
+}
+
 /// Reads the next item
 pub fn read_next<R: Read>(
     reader: &mut R,
     metadata: &StreamMetadata,
     dictionaries_by_field: &mut Vec<Option<ArrayRef>>,
-) -> Result<Option<RecordBatch>> {
+) -> Result<Option<StreamState>> {
     // determine metadata length
     let mut meta_size: [u8; 4] = [0; 4];
 
@@ -92,7 +107,7 @@ pub fn read_next<R: Read>(
                 // Handle EOF without the "0xFFFFFFFF 0x00000000"
                 // valid according to:
                 // https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format
-                Ok(None)
+                Ok(Some(StreamState::Waiting))
             } else {
                 Err(ArrowError::from(e))
             };
@@ -144,7 +159,7 @@ pub fn read_next<R: Read>(
                 &mut reader,
                 0,
             )
-            .map(Some)
+            .map(|x| Some(StreamState::Some(x)))
         }
         gen::Message::MessageHeader::DictionaryBatch => {
             let batch = message.header_as_dictionary_batch().ok_or_else(|| {
@@ -168,7 +183,7 @@ pub fn read_next<R: Read>(
             // read the next message until we encounter a RecordBatch
             read_next(reader, metadata, dictionaries_by_field)
         }
-        gen::Message::MessageHeader::NONE => Ok(None),
+        gen::Message::MessageHeader::NONE => Ok(Some(StreamState::Waiting)),
         t => Err(ArrowError::Ipc(format!(
             "Reading types other than record batches not yet supported, unable to read {:?} ",
             t
@@ -210,32 +225,26 @@ impl<R: Read> StreamReader<R> {
         self.finished
     }
 
-    fn maybe_next(&mut self) -> Result<Option<RecordBatch>> {
+    fn maybe_next(&mut self) -> Result<Option<StreamState>> {
+        if self.finished {
+            return Ok(None);
+        }
         let batch = read_next(
             &mut self.reader,
             &self.metadata,
             &mut self.dictionaries_by_field,
         )?;
         if batch.is_none() {
-            self.finished = false;
-        }
-        if self.finished {
-            return Ok(None);
+            self.finished = true;
         }
         Ok(batch)
     }
 }
 
 impl<R: Read> Iterator for StreamReader<R> {
-    type Item = Result<RecordBatch>;
+    type Item = Result<StreamState>;
 
     fn next(&mut self) -> Option<Self::Item> {
         self.maybe_next().transpose()
     }
 }
-
-impl<R: Read> RecordBatchReader for StreamReader<R> {
-    fn schema(&self) -> &Schema {
-        self.metadata.schema.as_ref()
-    }
-}


Test patch:
diff --git a/integration-testing/src/bin/arrow-file-to-stream.rs b/integration-testing/src/bin/arrow-file-to-stream.rs
index 9b1ae23a0b..34da2bdf07 100644
--- a/integration-testing/src/bin/arrow-file-to-stream.rs
+++ b/integration-testing/src/bin/arrow-file-to-stream.rs
@@ -30,7 +30,7 @@ fn main() -> Result<()> {
     let mut reader = read::FileReader::new(&mut f, metadata, None);
     let schema = reader.schema();
 
-    let mut writer = StreamWriter::try_new(std::io::stdout(), &schema)?;
+    let mut writer = StreamWriter::try_new(std::io::stdout(), schema)?;
 
     reader.try_for_each(|batch| {
         let batch = batch?;
diff --git a/integration-testing/src/bin/arrow-stream-to-file.rs b/integration-testing/src/bin/arrow-stream-to-file.rs
index c9db624d64..c17ff7d9e9 100644
--- a/integration-testing/src/bin/arrow-stream-to-file.rs
+++ b/integration-testing/src/bin/arrow-stream-to-file.rs
@@ -29,9 +29,9 @@ fn main() -> Result<()> {
 
     let mut writer = io::stdout();
 
-    let mut writer = FileWriter::try_new(&mut writer, &schema)?;
+    let mut writer = FileWriter::try_new(&mut writer, schema)?;
 
-    arrow_stream_reader.try_for_each(|batch| writer.write(&batch?))?;
+    arrow_stream_reader.try_for_each(|batch| writer.write(&batch?.unwrap()))?;
     writer.finish()?;
 
     Ok(())
diff --git a/integration-testing/src/flight_client_scenarios/integration_test.rs b/integration-testing/src/flight_client_scenarios/integration_test.rs
index 7cb128cdf6..820c82114b 100644
--- a/integration-testing/src/flight_client_scenarios/integration_test.rs
+++ b/integration-testing/src/flight_client_scenarios/integration_test.rs
@@ -131,7 +131,7 @@ async fn send_batch(
     options: &write::IpcWriteOptions,
 ) -> Result {
     let (dictionary_flight_data, mut batch_flight_data) =
-        arrow_flight::utils::flight_data_from_arrow_batch(batch, &options);
+        arrow_flight::utils::flight_data_from_arrow_batch(batch, options);
 
     upload_tx
         .send_all(&mut stream::iter(dictionary_flight_data).map(Ok))
@@ -169,7 +169,7 @@ async fn verify_data(
             consume_flight_location(
                 location,
                 ticket.clone(),
-                &expected_data,
+                expected_data,
                 expected_schema.clone(),
             )
             .await?;
diff --git a/integration-testing/src/flight_server_scenarios/integration_test.rs b/integration-testing/src/flight_server_scenarios/integration_test.rs
index de6951fee4..1695464709 100644
--- a/integration-testing/src/flight_server_scenarios/integration_test.rs
+++ b/integration-testing/src/flight_server_scenarios/integration_test.rs
@@ -295,7 +295,7 @@ async fn record_batch_from_message(
         schema_ref,
         None,
         true,
-        &dictionaries_by_field,
+        dictionaries_by_field,
         MetadataVersion::V5,
         &mut reader,
         0,
diff --git a/tests/it/io/ipc/common.rs b/tests/it/io/ipc/common.rs
index f9f2be91fe..6a46c03393 100644
--- a/tests/it/io/ipc/common.rs
+++ b/tests/it/io/ipc/common.rs
@@ -62,6 +62,9 @@ pub fn read_arrow_stream(version: &str, file_name: &str) -> (Schema, Vec<RecordB
 
     (
         schema.as_ref().clone(),
-        reader.collect::<Result<_>>().unwrap(),
+        reader
+            .map(|x| x.map(|x| x.unwrap()))
+            .collect::<Result<_>>()
+            .unwrap(),
     )
 }
diff --git a/tests/it/io/ipc/read/stream.rs b/tests/it/io/ipc/read/stream.rs
index 9aa5e31a1b..d6ad005e78 100644
--- a/tests/it/io/ipc/read/stream.rs
+++ b/tests/it/io/ipc/read/stream.rs
@@ -22,7 +22,7 @@ fn test_file(version: &str, file_name: &str) -> Result<()> {
 
     batches
         .iter()
-        .zip(reader.map(|x| x.unwrap()))
+        .zip(reader.map(|x| x.unwrap().unwrap()))
         .for_each(|(lhs, rhs)| {
             assert_eq!(lhs, &rhs);
         });
diff --git a/tests/it/io/ipc/write/stream.rs b/tests/it/io/ipc/write/stream.rs
index c13831839b..e9019d01f9 100644
--- a/tests/it/io/ipc/write/stream.rs
+++ b/tests/it/io/ipc/write/stream.rs
@@ -34,7 +34,10 @@ fn test_file(version: &str, file_name: &str) {
 
     assert_eq!(schema.as_ref(), &expected_schema);
 
-    let batches = reader.collect::<Result<Vec<_>>>().unwrap();
+    let batches = reader
+        .map(|x| x.map(|x| x.unwrap()))
+        .collect::<Result<Vec<_>>>()
+        .unwrap();
 
     assert_eq!(batches, expected_batches);
 }