Repository: vaexio/vaex

Issue:
Can't open Arrow files
See https://github.com/vaexio/vaex/blob/2eaba353cc2f330933d6e9a4392b792bfad98183/packages/vaex-arrow/vaex_arrow/dataset.py#L26-L31.

It seems that `arrow` now throws an `OSError` when trying to open a non-stream file file with `open_stream`:

```
>>> pyarrow.ipc.open_stream("feat.arrow")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/me/.local/lib/python3.8/site-packages/pyarrow/ipc.py", line 146, in open_stream
    return RecordBatchStreamReader(source)
  File "/home/me/.local/lib/python3.8/site-packages/pyarrow/ipc.py", line 62, in __init__
    self._open(source)
  File "pyarrow/ipc.pxi", line 360, in pyarrow.lib._RecordBatchStreamReader._open
  File "pyarrow/error.pxi", line 123, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 100, in pyarrow.lib.check_status
OSError: Invalid flatbuffers message.
```


Discussion:
Hi Laurențiu,

thanks, great feedback, I thought this was checked in CI, but we should test and fix that!

cheers,

Maarten
Would you dare to do a PR?

Resolving patch:
diff --git a/packages/vaex-core/vaex/arrow/dataset.py b/packages/vaex-core/vaex/arrow/dataset.py
--- a/packages/vaex-core/vaex/arrow/dataset.py
+++ b/packages/vaex-core/vaex/arrow/dataset.py
@@ -154,16 +154,16 @@ def from_table(table):
 
 def open(path, fs_options):
     source = vaex.file.open_for_arrow(path=path, mode='rb', fs_options=fs_options, mmap=True)
-    try:
-        # first we try if it opens as stream
-        reader = pa.ipc.open_stream(source)
-    except pa.lib.ArrowInvalid:
-        # if not, we open as file
+    file_signature = source.read(6)
+    is_arrow_file = file_signature == b'ARROW1'
+    source.seek(0)
+
+    if is_arrow_file:
         reader = pa.ipc.open_file(source)
         # for some reason this reader is not iterable
         batches = [reader.get_batch(i) for i in range(reader.num_record_batches)]
     else:
-        # if a stream, we're good
+        reader = pa.ipc.open_stream(source)
         batches = reader  # this reader is iterable
     table = pa.Table.from_batches(batches)
     return from_table(table)
diff --git a/packages/vaex-core/vaex/dataframe.py b/packages/vaex-core/vaex/dataframe.py
--- a/packages/vaex-core/vaex/dataframe.py
+++ b/packages/vaex-core/vaex/dataframe.py
@@ -5927,7 +5927,7 @@ def export(self, path, progress=None, chunk_size=default_chunk_size, parallel=Tr
             raise ValueError('''Unrecognized file extension. Please use .arrow, .hdf5, .parquet, .fits, or .csv to export to the particular file format.''')
 
     @docsubst
-    def export_arrow(self, to, progress=None, chunk_size=default_chunk_size, parallel=True, reduce_large=True, fs_options=None):
+    def export_arrow(self, to, progress=None, chunk_size=default_chunk_size, parallel=True, reduce_large=True, fs_options=None, as_stream=True):
         """Exports the DataFrame to a file of stream written with arrow
 
         :param to: filename, file object, or :py:data:`pyarrow.RecordBatchStreamWriter`, py:data:`pyarrow.RecordBatchFileWriter` or :py:data:`pyarrow.parquet.ParquetWriter`
@@ -5936,6 +5936,8 @@ def export_arrow(self, to, progress=None, chunk_size=default_chunk_size, paralle
         :param bool parallel: {evaluate_parallel}
         :param bool reduce_large: If True, convert arrow large_string type to string type
         :param dict fs_options: Coming soon...
+        :param bool as_stream: Write as an Arrow stream if true, else a file.
+            see also https://arrow.apache.org/docs/format/Columnar.html?highlight=arrow1#ipc-file-format
         :return:
         """
         progressbar = vaex.utils.progressbars(progress)
@@ -5954,8 +5956,12 @@ def write(writer):
             schema = self[0:1].to_arrow_table(parallel=False, reduce_large=reduce_large).schema
             fs_options = fs_options or {}
             with vaex.file.open_for_arrow(path=to, mode='wb', fs_options=fs_options) as sink:
-                writer = pa.RecordBatchStreamWriter(sink, schema)
-                write(writer)
+                if as_stream:
+                    with pa.RecordBatchStreamWriter(sink, schema) as writer:
+                        write(writer)
+                else:
+                    with pa.RecordBatchFileWriter(sink, schema) as writer:
+                        write(writer)
         else:
             write(to)
 


Test patch:
diff --git a/tests/arrow/io_test.py b/tests/arrow/io_test.py
--- a/tests/arrow/io_test.py
+++ b/tests/arrow/io_test.py
@@ -1,3 +1,5 @@
+import pytest
+
 import pandas as pd
 import pyarrow as pa
 import vaex
@@ -6,22 +8,12 @@
 table = pa.Table.from_pandas(df)
 
 
-def test_arrow_write_table(tmpdir):
+@pytest.mark.parametrize("as_stream", [True, False])
+def test_arrow_write_table(tmpdir, as_stream):
     path = str(tmpdir.join('test.arrow'))
-    with pa.OSFile(path, 'wb') as sink:
-        with pa.RecordBatchFileWriter(sink, table.schema) as writer:
-            writer.write_table(table)
-
-    df = vaex.open(path)
-
-
-def test_arrow_write_stream(tmpdir):
-    path = str(tmpdir.join('test.arrow'))
-    with pa.OSFile(path, 'wb') as sink:
-        with pa.RecordBatchStreamWriter(sink, table.schema) as writer:
-            writer.write_table(table)
-
+    vaex.from_arrow_table(table).export_arrow(path, as_stream=as_stream)
     df = vaex.open(path)
+    assert 'col1' in df
 
 
 def test_chunks(df_trimmed, tmpdir):