Repository: vaexio/vaex

Issue:
[BUG-REPORT] AttributeError: 'pyarrow.lib.ChunkedArray' object has no attribute 'offset'
Hello, I am seeing the following error calling `mean` on a filtered dataset:

```
AttributeError                            Traceback (most recent call last)
<ipython-input-16-cc5013060237> in <module>
----> 1 df.column_a.mean()

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/expression.py in mean(self, binby, limits, shape, selection, delay, progress)
    550         del kwargs['self']
    551         kwargs['expression'] = self.expression
--> 552         return self.ds.mean(**kwargs)
    553 
    554     def std(self, binby=[], limits=None, shape=default_shape, selection=False, delay=False, progress=None):

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/dataframe.py in mean(self, expression, binby, limits, shape, selection, delay, progress, edges, array_type)
    851         :return: {return_stat_scalar}
    852         """
--> 853         return self._compute_agg('mean', expression, binby, limits, shape, selection, delay, edges, progress, array_type=array_type)
    854         logger.debug("mean of %r, with binby=%r, limits=%r, shape=%r, selection=%r, delay=%r", expression, binby, limits, shape, selection, delay)
    855         expression = _ensure_strings_from_expressions(expression)

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/dataframe.py in _compute_agg(self, name, expression, binby, limits, shape, selection, delay, edges, progress, extra_expressions, array_type)
    683             # give errors when evaluated with NaN's present.
    684             # TODO: GET RID OF THIS
--> 685             len(self) # fill caches and masks
    686             # pass
    687         grid = self._create_grid(binby, limits, shape, selection=selection, delay=True)

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/dataframe.py in __len__(self)
   3763         else:
   3764             if self._cached_filtered_length is None:
-> 3765                self. _cached_filtered_length = int(self.count())
   3766             return self._cached_filtered_length
   3767 

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/dataframe.py in count(self, expression, binby, limits, shape, selection, delay, edges, progress, array_type)
    768         :return: {return_stat_scalar}
    769         """
--> 770         return self._compute_agg('count', expression, binby, limits, shape, selection, delay, edges, progress, array_type=array_type)
    771 
    772     @delayed

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/dataframe.py in _compute_agg(self, name, expression, binby, limits, shape, selection, delay, edges, progress, extra_expressions, array_type)
    742         stats = [compute(expression, grid, selection=selection, edges=edges, progressbar=progressbar) for expression in expressions]
    743         var = finish(grid, *stats)
--> 744         return self._delay(delay, var)
    745 
    746     @docsubst

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/dataframe.py in _delay(self, delay, task, progressbar)
   1506             return task
   1507         else:
-> 1508             self.execute()
   1509             return task.get()
   1510 

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/dataframe.py in execute(self)
    343         '''Execute all delayed jobs.'''
    344         from .asyncio import just_run
--> 345         just_run(self.execute_async())
    346 
    347     async def execute_async(self):

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/asyncio.py in just_run(coro)
     33             nest_asyncio.apply()
     34             check_patch_tornado()
---> 35         return loop.run_until_complete(coro)
     36     finally:
     37         if not had_loop:  # remove loop if we did not have one

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/nest_asyncio.py in run_until_complete(self, future)
     96                 raise RuntimeError(
     97                     'Event loop stopped before Future completed.')
---> 98             return f.result()
     99         finally:
    100             self._thread_id = old_thread_id

~/.conda-envs/datalakeutils/lib/python3.6/asyncio/tasks.py in _step(***failed resolving arguments***)
    180                 result = coro.send(None)
    181             else:
--> 182                 result = coro.throw(exc)
    183         except StopIteration as exc:
    184             if self._must_cancel:

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/dataframe.py in execute_async(self)
    348         '''Async version of execute'''
    349         # no need to clear _task_aggs anymore, since they will be removed for the executors' task list
--> 350         await self.executor.execute_async()
    351 
    352     @property

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/execution.py in execute_async(self)
    182                                                     all([all(task.signal_progress.emit(p)) for task in tasks]) and
    183                                                     all([not task.cancelled for task in tasks]),
--> 184                                                     cancel=lambda: self._cancel(run), unpack=True, run=run):
    185                     pass  # just eat all element
    186                 logger.debug("executing took %r seconds" % (time.time() - t0))

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/multithreading.py in map_async(self, callable, iterator, count, on_error, progress, cancel, unpack, **kwargs_extra)
     88         total = 0
     89         for i, value in buffer(enumerate(iterator), self._max_workers + 3):
---> 90             value = await value
     91             if value != None:
     92                 total += value

~/.conda-envs/datalakeutils/lib/python3.6/asyncio/futures.py in __iter__(self)
    325         if not self.done():
    326             self._asyncio_future_blocking = True
--> 327             yield self  # This tells Task to wait for completion.
    328         assert self.done(), "yield from wasn't used with future"
    329         return self.result()  # May raise too.

~/.conda-envs/datalakeutils/lib/python3.6/asyncio/tasks.py in _wakeup(self, future)
    248     def _wakeup(self, future):
    249         try:
--> 250             future.result()
    251         except Exception as exc:
    252             # This may also be a cancellation.

~/.conda-envs/datalakeutils/lib/python3.6/asyncio/futures.py in result(self)
    241         self._log_traceback = False
    242         if self._exception is not None:
--> 243             raise self._exception
    244         return self._result
    245 

~/.conda-envs/datalakeutils/lib/python3.6/concurrent/futures/thread.py in run(self)
     54 
     55         try:
---> 56             result = self.fn(*self.args, **self.kwargs)
     57         except BaseException as exc:
     58             self.future.set_exception(exc)

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/multithreading.py in <lambda>(value)
     84         else:
     85             loop = asyncio.get_event_loop()
---> 86             iterator = (loop.run_in_executor(self, lambda value=value: wrapped(value)) for value in iterator)
     87 
     88         total = 0

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/multithreading.py in wrapped(*args, **kwargs)
     76                     # print("SLEEP", self._debug_sleep)
     77                     time.sleep(self._debug_sleep)
---> 78                 return callable(self.local.index, *args, **kwargs, **kwargs_extra)
     79         # convert to list so we can count
     80         time_last = time.time() - 100

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/execution.py in process_part(self, thread_index, i1, i2, chunks, run)
    232                 assert len(chunk) == N, f'Oops, got a chunk ({name}) of length {len(chunk)} while it is expected to be of length {N} (at {i1}-{i2}'
    233             if run.pre_filter:
--> 234                 filter_mask = df.evaluate_selection_mask(None, i1=i1, i2=i2, cache=True)
    235                 chunks = {k:vaex.array_types.filter(v, filter_mask) for k, v, in chunks.items()}
    236             else:

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/dataframe.py in evaluate_selection_mask(self, name, i1, i2, selection, cache, filtered, pre_filtered)
   2578         if name in [None, False] and self.filtered and filtered:
   2579             scope_global = scopes._BlockScopeSelection(self, i1, i2, None, cache=cache)
-> 2580             mask_global = scope_global.evaluate(FILTER_SELECTION_NAME)
   2581             return vaex.utils.unmask_selection_mask(mask_global)
   2582         elif self.filtered and filtered and name != FILTER_SELECTION_NAME:

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/scopes.py in evaluate(self, expression)
    181         try:
    182             expression = _ensure_string_from_expression(expression)
--> 183             result = eval(expression, expression_namespace, self)
    184         except:
    185             import traceback as tb

<string> in <module>

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/scopes.py in __getitem__(self, variable)
    217                 if variable in self.df.variables:
    218                     return self.df.variables[variable]
--> 219                 mask_values = selection.evaluate(self.df, variable, self.i1, self.i2, self.filter_mask)
    220 
    221                 # get a view on a subset of the mask

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/selections.py in evaluate(self, df, name, i1, i2, filter_mask)
    128         else:
    129             previous_mask = None
--> 130         result = df._evaluate_selection_mask(self.boolean_expression, i1, i2, filter_mask=filter_mask)
    131         if isinstance(result, bool):
    132             N = i2 - i1

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/dataframe.py in _evaluate_selection_mask(self, name, i1, i2, selection, cache, filter_mask)
   2564         i2 = i2 or len(self)
   2565         scope = scopes._BlockScopeSelection(self, i1, i2, selection, cache=cache, filter_mask=filter_mask)
-> 2566         mask = scope.evaluate(name)
   2567         # TODO: can we do without arrow->numpy conversion?
   2568         mask = vaex.array_types.to_numpy(mask)

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/scopes.py in evaluate(self, expression)
    181         try:
    182             expression = _ensure_string_from_expression(expression)
--> 183             result = eval(expression, expression_namespace, self)
    184         except:
    185             import traceback as tb

<string> in <module>

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/arrow/numpy_dispatch.py in operator(a, b)
     85                 result_data = op['op'](a_data, b_data)
     86             if isinstance(a, NumpyDispatch):
---> 87                 result_data = a.add_missing(result_data)
     88             if isinstance(b, NumpyDispatch):
     89                 result_data = b.add_missing(result_data)

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/arrow/numpy_dispatch.py in add_missing(self, ar)
     39             if isinstance(self._array, vaex.array_types.supported_arrow_array_types):
     40                 ar = vaex.array_types.to_arrow(ar)
---> 41                 ar = combine_missing(ar, self._array)
     42             # else: both numpy, handled by numpy
     43         else:

~/.conda-envs/datalakeutils/lib/python3.6/site-packages/vaex/arrow/numpy_dispatch.py in combine_missing(a, b)
     11         # not optimal
     12         nulls = pc.invert(pc.or_(a.is_null(), b.is_null()))
---> 13         assert nulls.offset == 0
     14         nulls_buffer = nulls.buffers()[1]
     15         # this is not the case: no reason why it should be (TODO: open arrow issue)

AttributeError: 'pyarrow.lib.ChunkedArray' object has no attribute 'offset'
```

PyArrow's documentation on ChunkedArray: https://arrow.apache.org/docs/python/generated/pyarrow.ChunkedArray.html

Vaex Version installed via pip:

```
vaex-core 4.0.0a8
```
OS: Ubuntu 20.04.1 LTS (Focal Fossa)


Discussion:
Hi, 

Which pyarrow version do you have? Can you please upgrade to 2.0? 
If that does not help, can you please provide some code (a simple minimal reproducible example).

Thanks
Sure thing. I have `pyarrow 2.0.0` installed already. For context, the dataframe I'm working with has one million rows and 380 columns. I filter the dataframe using some column value and I call `mean` on a different column.

```
df_subset = df[df.column_a < 100]
avg_b = df_subset.column_b.mean()
```
These are the only operations I am running on the dataframe.
I know enough 👍 

Resolving patch:
diff --git a/packages/vaex-core/vaex/arrow/convert.py b/packages/vaex-core/vaex/arrow/convert.py
--- a/packages/vaex-core/vaex/arrow/convert.py
+++ b/packages/vaex-core/vaex/arrow/convert.py
@@ -10,7 +10,7 @@
 
 def ensure_not_chunked(arrow_array):
     if isinstance(arrow_array, pa.ChunkedArray):
-        if len(arrow_array.chunks) == 0:
+        if len(arrow_array.chunks) == 1:
             return arrow_array.chunks[0]
         table = pa.Table.from_arrays([arrow_array], ["single"])
         table_concat = table.combine_chunks()
@@ -243,3 +243,32 @@ def trim_buffers_ipc(ar):
 def trim_buffers_for_pickle(ar):
     # future version of pyarrow might fix this, so we have a single entry point for this
     return trim_buffers_ipc(ar)
+
+
+def align(a, b):
+    '''Align two arrays/chunked arrays such that they have the same chunk lengths'''
+    if len(a) != len(b):
+        raise ValueError(f'Length of arrays should be equals ({len(a)} != {len(b)})')
+    if isinstance(a, pa.ChunkedArray) and len(a.chunks) == 1:
+        a = a.chunks[0]
+    if isinstance(b, pa.ChunkedArray) and len(b.chunks) == 1:
+        b = b.chunks[0]
+    if isinstance(a, pa.ChunkedArray) and isinstance(b, pa.ChunkedArray):
+        lengths_a = [len(c) for c in a.chunks]
+        lengths_b = [len(c) for c in b.chunks]
+        if lengths_a == lengths_b:
+            return a, b
+        else:
+            return ensure_not_chunked(a), ensure_not_chunked(b)
+    elif isinstance(a, pa.ChunkedArray) and not isinstance(b, pa.ChunkedArray):
+        lengths = [len(c) for c in a.chunks]
+        # numpy cannot do an exclusive cumsum
+        offsets = np.cumsum([0] + lengths)
+        offsets = offsets[:-1]
+        return a, pa.chunked_array([b.slice(offset, length) for offset, length in zip(offsets, lengths)])
+    elif not isinstance(a, pa.ChunkedArray) and isinstance(b, pa.ChunkedArray):
+        b, a = align(b, a)
+        return a, b
+    else:
+        return a, b
+    return
diff --git a/packages/vaex-core/vaex/arrow/numpy_dispatch.py b/packages/vaex-core/vaex/arrow/numpy_dispatch.py
--- a/packages/vaex-core/vaex/arrow/numpy_dispatch.py
+++ b/packages/vaex-core/vaex/arrow/numpy_dispatch.py
@@ -6,18 +6,31 @@
 
 
 def combine_missing(a, b):
-    assert a.offset == 0
+    # return a copy of a with missing values of a and b combined
     if a.null_count > 0 or b.null_count > 0:
+        a, b = vaex.arrow.convert.align(a, b)
+        if isinstance(a, pa.ChunkedArray):
+            # divide and conquer
+            assert isinstance(b, pa.ChunkedArray)
+            assert len(a.chunks) == len(b.chunks)
+            return pa.chunked_array([combine_missing(ca, cb) for ca, cb in zip(a.chunks, b.chunks)])
+        if a.offset != 0:
+            a = vaex.arrow.convert.trim_buffers_ipc(a)
+        if b.offset != 0:
+            b = vaex.arrow.convert.trim_buffers_ipc(b)
+        assert a.offset == 0
+        assert b.offset == 0
         # not optimal
         nulls = pc.invert(pc.or_(a.is_null(), b.is_null()))
         assert nulls.offset == 0
         nulls_buffer = nulls.buffers()[1]
         # this is not the case: no reason why it should be (TODO: open arrow issue)
         # assert nulls.buffers()[0] is None
+        buffers = a.buffers()
+        return pa.Array.from_buffers(a.type, len(a), [nulls_buffer] + buffers[1:])
     else:
-        nulls_buffer = None
-    buffers = a.buffers()
-    return pa.Array.from_buffers(a.type, len(a), [nulls_buffer] + buffers[1:])
+        return a
+
 
 class NumpyDispatch:
     def __init__(self, ar):


Test patch:
diff --git a/tests/arrow/convert_test.py b/tests/arrow/convert_test.py
--- a/tests/arrow/convert_test.py
+++ b/tests/arrow/convert_test.py
@@ -40,3 +40,32 @@ def test_large_string_to_string(offset, chunked):
     if not chunked:
         assert ns.offset < 8
     assert s.to_pylist() == ns.to_pylist()
+
+
+def test_align():
+    ac = pa.chunked_array([[1, 2], [3], [4, 5, 6]])
+    ac2 = pa.chunked_array([[1, 2], [3, 4, 5], [6]])
+    af = pa.array([1, 2, 3, 4, 5, 6])
+    x1, x2 = convert.align(ac, af)
+    assert isinstance(x1, pa.ChunkedArray)
+    assert isinstance(x2, pa.ChunkedArray)
+    assert x1.to_pylist() == af.to_pylist()
+    assert x2.to_pylist() == af.to_pylist()
+
+    x1, x2 = convert.align(af, ac)
+    assert isinstance(x1, pa.ChunkedArray)
+    assert isinstance(x2, pa.ChunkedArray)
+    assert x1.to_pylist() == af.to_pylist()
+    assert x2.to_pylist() == af.to_pylist()
+
+    x1, x2 = convert.align(ac, ac2)
+    assert isinstance(x1, pa.Array)
+    assert isinstance(x2, pa.Array)
+    assert x1.to_pylist() == af.to_pylist()
+    assert x2.to_pylist() == af.to_pylist()
+
+    x1, x2 = convert.align(af, af)
+    assert isinstance(x1, pa.Array)
+    assert isinstance(x2, pa.Array)
+    assert x1.to_pylist() == af.to_pylist()
+    assert x2.to_pylist() == af.to_pylist()
diff --git a/tests/common.py b/tests/common.py
--- a/tests/common.py
+++ b/tests/common.py
@@ -363,13 +363,29 @@ def create(x):
 
 
 @pytest.fixture(scope='session')
-def array_factory_arrow():
+def array_factory_arrow_normal():
     return pa.array
 
 
-@pytest.fixture(scope='session', params=['array_factory_numpy', 'array_factory_arrow'])
-def array_factory(request, array_factory_numpy, array_factory_arrow):
-    named = dict(array_factory_numpy=array_factory_numpy, array_factory_arrow=array_factory_arrow)
+@pytest.fixture(scope='session')
+def array_factory_arrow_chunked():
+    def create(x):
+        x = pa.array(x)
+        mid = len(x) // 3
+        c1 = x.slice(0, mid)
+        c2 = x.slice(mid)
+        return pa.chunked_array([c1, c2])
+    return create
+
+
+@pytest.fixture(scope='session', params=['array_factory_numpy', 'array_factory_arrow_normal', 'array_factory_arrow_chunked'])
+def array_factory(request, array_factory_numpy, array_factory_arrow_normal, array_factory_arrow_chunked):
+    named = dict(array_factory_numpy=array_factory_numpy, array_factory_arrow_normal=array_factory_arrow_normal, array_factory_arrow_chunked=array_factory_arrow_chunked)
+    return named[request.param]
+
+@pytest.fixture(scope='session', params=['array_factory_arrow_normal', 'array_factory_arrow_chunked'])
+def array_factory_arrow(request, array_factory_arrow_normal, array_factory_arrow_chunked):
+    named = dict(array_factory_arrow_normal=array_factory_arrow_normal, array_factory_arrow_chunked=array_factory_arrow_chunked)
     return named[request.param]
 
 array_factory1 = array_factory
diff --git a/tests/compute_test.py b/tests/compute_test.py
--- a/tests/compute_test.py
+++ b/tests/compute_test.py
@@ -18,16 +18,16 @@ def y(array_factory2):
 
 
 @pytest.fixture(scope='session')
-def s():
-    return pa.array(['a', 'b', None, 'd'])
+def s(array_factory_arrow):
+    return array_factory_arrow(['a', 'b', None, 'd'])
 
 
 def test_add(x, y):
     df = vaex.from_arrays(x=x, y=y)
     df['z'] = df.x + df.y
     assert df['z'].tolist() == [1, 3, None, None]
-    assert isinstance(df['z'].values.tolist()[0], numbers.Number)
-    assert df['z'].values.tolist() == [1, 3, None, None]
+    assert isinstance(df['z'].tolist()[0], numbers.Number)
+    assert df['z'].tolist() == [1, 3, None, None]
     assert df['z'].dtype == 'numeric'
     # import pdb; pdb.set_trace()
 
diff --git a/tests/internal/strings_module_test.py b/tests/internal/strings_module_test.py
--- a/tests/internal/strings_module_test.py
+++ b/tests/internal/strings_module_test.py
@@ -192,12 +192,14 @@ def test_arrow_basics():
     # ref counts start at 2
     assert sys.getrefcount(bytes_strings) == 2
     assert sys.getrefcount(indices) == 2
+    assert sys.getrefcount(null_bitmap) == 2
 
     sl = vaex.strings.StringList32(bytes_strings, indices, len(ar), offset, null_bitmap, 0)
     assert sys.getrefcount(sl) == 2
     # we should keep a reference only
     assert sys.getrefcount(bytes_strings) == 3
     assert sys.getrefcount(indices) == 3
+    assert sys.getrefcount(null_bitmap) == 3
 
     assert sl.get(0) == "aap"
     assert sl.get(1) == "noot"