Repository: vaexio/vaex

Issue:
[BUG-REPORT] `vaex.from_arrow_dataset` throws`'pyarrow._dataset.FileFragment' object has no attribute 'row_groups'`
**Description**
I'm trying to read `feather` files as a dataset into one dataframe. Works with `pyarrow`, exception in `vaex`:
```python
import pyarrow
import pyarrow.dataset
import pyarrow.feather
import pandas
import vaex

# pandas version
df1 = pandas.DataFrame({"a": [1,2,3], "b": [4,5,6]})
df2 = pandas.DataFrame({"a": [10,20,30], "b": [40,50,60]})

# write independent Feather files
path_base = "/home/toaster/Downloads/"
pyarrow.feather.write_feather(df1, path_base + "123.feather")
pyarrow.feather.write_feather(df2, path_base + "456.feather")

# read Feather files as Arrow dataset
ds = pyarrow.dataset.dataset(path_base, format="feather")
ds.files  # lists both "123.feather" and "456.feather"
ds.to_table().to_pandas()  # outputs expected dataframe


# trying to import Arrow dataset in Vaex
vaex.from_arrow_dataset(ds)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In [48], line 1
----> 1 df_vaex = vaex.from_arrow_dataset(ds)

File ~/mambaforge-pypy3/envs/puma-lab/lib/python3.8/site-packages/vaex/__init__.py:384, in from_arrow_dataset(arrow_dataset)
    382 '''Create a DataFrame from an Apache Arrow dataset.'''
    383 import vaex.arrow.dataset
--> 384 return from_dataset(vaex.arrow.dataset.DatasetArrow(arrow_dataset))

File ~/mambaforge-pypy3/envs/puma-lab/lib/python3.8/site-packages/vaex/arrow/dataset.py:258, in DatasetArrow.__init__(self, ds, max_rows_read)
    256 def __init__(self, ds, max_rows_read=1024**2*10):
    257     self._arrow_ds = ds
--> 258     super().__init__(max_rows_read=max_rows_read)

File ~/mambaforge-pypy3/envs/puma-lab/lib/python3.8/site-packages/vaex/arrow/dataset.py:26, in DatasetArrowBase.__init__(self, max_rows_read)
     24 super().__init__()
     25 self.max_rows_read = max_rows_read
---> 26 self._create_columns()

File ~/mambaforge-pypy3/envs/puma-lab/lib/python3.8/site-packages/vaex/arrow/dataset.py:268, in DatasetArrow._create_columns(self)
    267 def _create_columns(self):
--> 268     super()._create_columns()
    269     # self._ids = frozendict({name: vaex.cache.fingerprint(self._fingerprint, name) for name in self._columns})
    270     self._ids = frozendict()

File ~/mambaforge-pypy3/envs/puma-lab/lib/python3.8/site-packages/vaex/arrow/dataset.py:35, in DatasetArrowBase._create_columns(self)
     33     if hasattr(fragment, "ensure_complete_metadata"):
     34         fragment.ensure_complete_metadata()
---> 35     for rg in fragment.row_groups:
     36         row_count += rg.num_rows
     37 self._row_count = row_count

AttributeError: 'pyarrow._dataset.FileFragment' object has no attribute 'row_groups'
---------------------------------------------------------------------------


# check if we have to write using `pyarrow.dataset.write_dataset` instead of independend Feather files
table_pa = ds.to_table(filter=pyarrow.dataset.field("a") > 2) 
pyarrow.dataset.write_dataset(table_pa, "test", format="feather")

ds_test = pyarrow.dataset.dataset("test", format="feather")
vaex.from_arrow_dataset(ds_test)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In [30], line 1
----> 1 vaex.from_arrow_dataset(ds_test)

File ~/mambaforge-pypy3/envs/puma-lab/lib/python3.8/site-packages/vaex/__init__.py:384, in from_arrow_dataset(arrow_dataset)
    382 '''Create a DataFrame from an Apache Arrow dataset.'''
    383 import vaex.arrow.dataset
--> 384 return from_dataset(vaex.arrow.dataset.DatasetArrow(arrow_dataset))

File ~/mambaforge-pypy3/envs/puma-lab/lib/python3.8/site-packages/vaex/arrow/dataset.py:258, in DatasetArrow.__init__(self, ds, max_rows_read)
    256 def __init__(self, ds, max_rows_read=1024**2*10):
    257     self._arrow_ds = ds
--> 258     super().__init__(max_rows_read=max_rows_read)

File ~/mambaforge-pypy3/envs/puma-lab/lib/python3.8/site-packages/vaex/arrow/dataset.py:26, in DatasetArrowBase.__init__(self, max_rows_read)
     24 super().__init__()
     25 self.max_rows_read = max_rows_read
---> 26 self._create_columns()

File ~/mambaforge-pypy3/envs/puma-lab/lib/python3.8/site-packages/vaex/arrow/dataset.py:268, in DatasetArrow._create_columns(self)
    267 def _create_columns(self):
--> 268     super()._create_columns()
    269     # self._ids = frozendict({name: vaex.cache.fingerprint(self._fingerprint, name) for name in self._columns})
    270     self._ids = frozendict()

File ~/mambaforge-pypy3/envs/puma-lab/lib/python3.8/site-packages/vaex/arrow/dataset.py:35, in DatasetArrowBase._create_columns(self)
     33     if hasattr(fragment, "ensure_complete_metadata"):
     34         fragment.ensure_complete_metadata()
---> 35     for rg in fragment.row_groups:
     36         row_count += rg.num_rows
     37 self._row_count = row_count

AttributeError: 'pyarrow._dataset.FileFragment' object has no attribute 'row_groups'
---------------------------------------------------------------------------
​

#Workarounds
# as a workaround this seems to work, but it loads all the data into memory!
table_pa = ds.to_table()
df_vaex = vaex.from_arrow_table(table_pa)

# this also works
df_vaex = vaex.open(path_base + "*.feather")
```



**Software information**
 - Vaex version (`import vaex; vaex.__version__)`: 4.13.0
 - pyarrow version: 9.0.0
 - Vaex was installed via: pip
 - OS: Ubuntu 22.04.01 LTS


Discussion:

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
@@ -32,8 +32,11 @@ def _create_columns(self):
         for fragment in self._arrow_ds.get_fragments():
             if hasattr(fragment, "ensure_complete_metadata"):
                 fragment.ensure_complete_metadata()
-            for rg in fragment.row_groups:
-                row_count += rg.num_rows
+            if hasattr(fragment, "count_rows"):
+                row_count += fragment.count_rows()
+            else:
+                for rg in fragment.row_groups:
+                    row_count += rg.num_rows
         self._row_count = row_count
         self._columns = {name: vaex.dataset.ColumnProxy(self, name, type) for name, type in
                           zip(self._arrow_ds.schema.names, self._arrow_ds.schema.types)}
@@ -82,13 +85,20 @@ def _chunk_producer(self, columns, chunk_size=None, reverse=False, start=0, end=
         columns_partition = tuple(columns_partition)
 
         for fragment_large in self._arrow_ds.get_fragments():
-            fragment_large_rows = sum([rg.num_rows for rg in fragment_large.row_groups])
-            fragments = [fragment_large]
-            # when do we want to split up? File size? max chunk size?
-            if fragment_large_rows > self.max_rows_read:
-                fragments = fragment_large.split_by_row_group()
+            if hasattr(fragment_large, "count_rows"):
+                fragment_large_rows = fragment_large.count_rows()
+                fragments = [fragment_large]
+            else:
+                fragment_large_rows = sum([rg.num_rows for rg in fragment_large.row_groups])
+                fragments = [fragment_large]
+                # when do we want to split up? File size? max chunk size?
+                if fragment_large_rows > self.max_rows_read:
+                    fragments = fragment_large.split_by_row_group()
             for fragment in fragments:
-                rows = sum([rg.num_rows for rg in fragment.row_groups])
+                if hasattr(fragment_large, "count_rows"):
+                    rows = fragment_large.count_rows()
+                else:
+                    rows = sum([rg.num_rows for rg in fragment.row_groups])
                 chunk_start = offset
                 chunk_end = offset + rows
 


Test patch:
diff --git a/tests/from_test.py b/tests/from_test.py
--- a/tests/from_test.py
+++ b/tests/from_test.py
@@ -40,3 +40,18 @@ def test_from_arrow_dataset():
     df = vaex.from_arrow_dataset(ds)
     assert df.col1.sum() == 45
     assert df.fingerprint() == df.fingerprint()
+
+
+
+def test_from_arrow_dataset(tmpdir):
+    import numpy as np
+    x = np.arange(10)
+    df = vaex.from_arrays(x=x)
+    import pyarrow.dataset
+    path1 = tmpdir / 'sample1.feather'
+    path2 = tmpdir / 'sample2.feather'
+    df[:5].export(path1)
+    df[5:].export(path2)
+    ds = pyarrow.dataset.dataset(tmpdir, format='feather')
+    df = vaex.from_arrow_dataset(ds)
+    assert df.x.sum() == 45