Repository: vaexio/vaex

Issue:
[BUG-REPORT] Groupby misbehaves when one of the grouping columns contains missing values
**Description**
If we try to group by one column that contains missing values, the result is erratic. 

Running
```
df1 = vaex.from_pandas(pd.DataFrame([1,2,3], columns=['idx']))
df2 = vaex.from_pandas(pd.DataFrame({'idx':[1,6,7], 'a': [0,0,0], 'b': [1,2,3]}))
df3 = df1.join(df2, how='left', on='idx')
df3.groupby(['idx', 'a']).agg({'b': 'sum'})
```
throws an exception:
`TypeError: int() argument must be a string, a bytes-like object or a number, not 'NoneType'`

However, if we save to Arrow and read back, the behavior is different. Sometimes a different exception is thrown:
`ArrowIndexError: Index -1 out of bounds`

Other times it does not fail and gives back a wrong result, which is really worrying:
```
df3.export('test.arrow')
z = vaex.open('test.arrow')
y = z.groupby(['idx', 'a']).agg({'b': 'sum'})
```
results in
```
z:
#	idx	a	b
0	1	0	1
1	2	--	--
2	3	--	--

and

y:
#	idx	a	b
0	1	0	1
1	2	0	0

```

**Software information**
 - Vaex version (`import vaex; vaex.__version__)`: {'vaex-core': '4.3.0.post1', 'vaex-hdf5': '0.8.0', 'vaex-ml': '0.12.0'}
 - Vaex was installed via: pip 
 - OS: Linux



Discussion:
Thanks for the report, I can confirm this is still an issue. I have an refactor planned that will also affect this part, I'll make sure it will also solve this.
I would love to contribute and solve this issue, if it still persists

Resolving patch:
diff --git a/packages/vaex-core/vaex/functions.py b/packages/vaex-core/vaex/functions.py
--- a/packages/vaex-core/vaex/functions.py
+++ b/packages/vaex-core/vaex/functions.py
@@ -148,7 +148,7 @@ def fillmissing(ar, value):
     See :`ismissing` for the definition of missing values.
     '''
     dtype = vaex.dtype_of(ar)
-    if dtype == str:
+    if dtype.is_arrow:
         return pc.fill_null(ar, value)
     ar = ar if not isinstance(ar, column.Column) else ar.to_numpy()
     mask = ismissing(ar)
@@ -2445,7 +2445,12 @@ def _ordinal_values(x, ordered_set):
     if not isinstance(x, vaex.array_types.supported_array_types) or isinstance(ordered_set, vaex.superutils.ordered_set_string):
         # sometimes the dtype can be object, but seen as an string array
         x = _to_string_sequence(x)
-    return ordered_set.map_ordinal(x)
+    else:
+        x = vaex.array_types.to_numpy(x)
+    indices = ordered_set.map_ordinal(x)
+    if np.ma.isMaskedArray(x):
+        indices[x.mask] = ordered_set.null_value
+    return indices
 
 @register_function()
 def _choose(ar, choices, default=None):
diff --git a/packages/vaex-core/vaex/groupby.py b/packages/vaex-core/vaex/groupby.py
--- a/packages/vaex-core/vaex/groupby.py
+++ b/packages/vaex-core/vaex/groupby.py
@@ -106,20 +106,31 @@ class BinnerInteger(BinnerBase):
 
     Useful for boolean, uint8, which only have a limited number of possibilities (2, 256)
     '''
-    def __init__(self, expression, label=None):
+    def __init__(self, expression, label=None, dropmissing=False):
         self.expression = expression
         self.df = expression.df
         self.dtype = self.expression.dtype
         self.label = label or self.expression._label
         if self.dtype.numpy == np.dtype('bool'):
-            self.bin_values = np.array([False, True])
-            self.N = 2
+            if dropmissing:
+                self.binby_expression = str(self.expression)
+                self.bin_values = np.array([False, True])
+                self.N = 2
+            else:
+                self.binby_expression = f'fillmissing(astype({str(self.expression)}, "uint8"), 2)'
+                self.bin_values = pa.array([False, True, None])
+                self.N = 3
         elif self.dtype.numpy == np.dtype('uint8'):
-            self.bin_values = np.arange(0, 256, dtype='uint8')
-            self.N = 256
+            if dropmissing:
+                self.binby_expression = str(self.expression)
+                self.bin_values = np.arange(0, 256, dtype="uint8")
+                self.N = 256
+            else:
+                self.binby_expression = f'fillmissing(astype({str(self.expression)}, "int16"), 256)'
+                self.bin_values = pa.array(list(range(256)) + [None])
+                self.N = 257
         else:
             raise TypeError(f'Only boolean and uint8 are supported, not {self.dtype}')
-        self.binby_expression = str(self.expression)
         self.sort_indices = None
         self._promise = vaex.promise.Promise.fulfilled(None)
 
@@ -133,6 +144,8 @@ class Grouper(BinnerBase):
     """Bins an expression to a set of unique bins, like an SQL like groupby."""
     def __init__(self, expression, df=None, sort=False, pre_sort=True, row_limit=None, df_original=None, materialize_experimental=False):
         self.df = df or expression.ds
+        self.sort = sort
+        self.pre_sort = pre_sort
         # we prefer to calculate the set the original dataframe to have better cache hits, and modify df
         if df_original is None:
             df_original = self.df
@@ -141,6 +154,7 @@ def __init__(self, expression, df=None, sort=False, pre_sort=True, row_limit=Non
         # make sure it's an expression
         self.expression = self.df[_ensure_string_from_expression(self.expression)]
         self.label = self.expression._label
+        dtype = self.expression.dtype
         if materialize_experimental:
             set, values = df_original._set(self.expression, unique_limit=row_limit, return_inverse=True)
             # TODO: add column should have a unique argument
@@ -169,9 +183,12 @@ def process(set):
                     self.bin_values = self.bin_values.astype(vaex.utils.required_dtype_for_max(max_value))
                 logger.debug('Constructed grouper for expression %s with %i values', str(expression), len(self.bin_values))
 
-                # since nan and null are at the start, we skip them with sorting
+                if set.has_null and (dtype.is_primitive or dtype.is_datetime):
+                    mask = np.zeros(shape=self.bin_values.shape, dtype="?")
+                    mask[set.null_value] = 1
+                    self.bin_values = np.ma.array(self.bin_values, mask=mask)
                 if self.sort:
-                    dtype = self.expression.dtype
+                    self.bin_values = vaex.array_types.to_arrow(self.bin_values)
                     indices = pa.compute.sort_indices(self.bin_values)#[offset:])
                     if pre_sort:
                         self.bin_values = pa.compute.take(self.bin_values, indices)
@@ -186,7 +203,6 @@ def process(set):
                             set = type(set)(self.bin_values, null_value, set.nan_count, set.null_count, fingerprint)
                         self.sort_indices = None
                     else:
-                        # TODO: skip first or first two values (null and/or nan)
                         self.sort_indices = vaex.array_types.to_numpy(indices)
                         # the bin_values will still be pre sorted, maybe that is confusing (implementation detail)
                         self.bin_values = pa.compute.take(self.bin_values, self.sort_indices)
@@ -197,7 +213,9 @@ def process(set):
                 self.basename = 'set_%s' % vaex.utils._python_save_name(str(self.expression) + "_" + set.fingerprint)
 
                 self.N = len(self.bin_values)
-                self.bin_values = self.expression.dtype.create_array(self.bin_values)
+                # for datetimes, we converted to int
+                if dtype.is_datetime:
+                    self.bin_values = dtype.create_array(self.bin_values)
             self._promise = process(df_original._set(self.expression, unique_limit=row_limit, delay=True))
 
     def _create_binner(self, df):
@@ -247,6 +265,9 @@ def compress(ar):
             bin_values = {}
             # NOTE: we can also use dict encoding instead of take
             for indices, parent in zip(indices_parents, parents):
+                if sort:
+                    assert parent.pre_sort, "cannot sort while parent not presorted"
+                    assert parent.sort_indices is None
                 dtype = vaex.dtype_of(parent.bin_values)
                 if dtype.is_struct:
                     # collapse parent struct into our flat struct
@@ -270,40 +291,66 @@ def set(bin_values):
 
 class GrouperCategory(BinnerBase):
     """Faster grouper that will use the fact that a column is categorical."""
-    def __init__(self, expression, df=None, sort=False, row_limit=None):
+
+    def __init__(self, expression, df=None, sort=False, row_limit=None, pre_sort=True):
         self.df = df or expression.ds
         self.sort = sort
+        self.pre_sort = pre_sort
         # make sure it's an expression
         expression = self.df[str(expression)]
         self.expression_original = expression
         self.label = expression._label
         self.expression = expression.index_values() if expression.dtype.is_encoded else expression
+        self.row_limit = row_limit
 
+        self.min_value = self.df.category_offset(self.expression_original)
         self.bin_values = self.df.category_labels(self.expression_original, aslist=False)
+        self.N = self.df.category_count(self.expression_original)
+        dtype = self.expression.dtype
         if self.sort:
-            self.sort_indices = pa.compute.sort_indices(self.bin_values)#[offset:])
-            self.bin_values = pa.compute.take(self.bin_values, self.sort_indices)
+            # not pre-sorting is faster
+            sort_indices = pa.compute.sort_indices(self.bin_values)
+            self.bin_values = pa.compute.take(self.bin_values, sort_indices)
+            if self.pre_sort:
+                sort_indices = vaex.array_types.to_numpy(sort_indices)
+                # TODO: this is kind of like expression.map
+                from .hash import ordered_set_type_from_dtype
+
+                ordered_set_type = ordered_set_type_from_dtype(dtype)
+                fingerprint = self.expression.fingerprint() + "-grouper-sort-mapper"
+                self.set = ordered_set_type(sort_indices + self.min_value, -1, 0, 0, fingerprint)
+                self.min_value = 0
+                self.sort_indices = None
+                self.basename = "set_%s" % vaex.utils._python_save_name(str(self.expression) + "_" + self.set.fingerprint)
+            else:
+                self.sort_indices = sort_indices
         else:
             self.sort_indices = None
         if isinstance(self.bin_values, list):
             self.bin_values = pa.array(self.bin_values)
 
-        self.N = self.df.category_count(self.expression_original)
         if row_limit is not None:
             if self.N > row_limit:
                 raise vaex.RowLimitException(f'Resulting grouper has {self.N:,} unique combinations, which is larger than the allowed row limit of {row_limit:,}')
-        self.min_value = self.df.category_offset(self.expression_original)
         # TODO: what do we do with null values for categories?
         # if self.set.has_null:
         #     self.N += 1
         #     keys += ['null']
-        self.binby_expression = str(self.expression)
         self._promise = vaex.promise.Promise.fulfilled(None)
 
     def _create_binner(self, df):
         assert df.dataset == self.df.dataset, "you passed a dataframe with a different dataset to the grouper/binned"
         self.df = df
-        self.binner = self.df._binner_ordinal(self.expression, self.N, self.min_value)
+        if self.sort and self.pre_sort:
+            if self.basename not in self.df.variables:
+                self.setname = df.add_variable(self.basename, self.set, unique=True)
+            else:
+                self.setname = self.basename
+            self.binby_expression = "_ordinal_values(%s, %s)" % (self.expression, self.setname)
+            self.binner = self.df._binner_ordinal(self.binby_expression, self.N, 0)
+        else:
+            self.binby_expression = str(self.expression)
+            self.binner = self.df._binner_ordinal(self.binby_expression, self.N, self.min_value)
 
 
 def _combine(df, groupers, sort, row_limit=None):
@@ -584,29 +631,40 @@ def agg(self, actions, delay=False):
         @vaex.delayed
         def aggregate(promise_by):
             arrays = super(GroupBy, self)._agg(actions)
-            has_non_existing_pairs = len(self.by) > 1
             # we don't want non-existing pairs (e.g. Amsterdam in France does not exist)
+            # but also, e.g. GrouperInteger will always expect missing values
+            # but they may not aways exist
             counts = self.counts
             # nobody wanted to know count*, but we need it if we included non-existing pairs
-            if has_non_existing_pairs and counts is None:
+            if counts is None:
                 # TODO: it seems this path is never tested
                 count_agg = vaex.agg.count(edges=True)
                 counts = self.df._agg(count_agg, self.binners, delay=_USE_DELAY)
             arrays = delayed_dict(arrays)
-            return counts, arrays, has_non_existing_pairs
+            return counts, arrays
 
         @vaex.delayed
         def process(args):
-            counts, arrays, has_non_existing_pairs = args
+            counts, arrays = args
             # arrays = {key: value.get() for key, value in arrays.items()}
             # take out the edges
             arrays = {key: vaex.utils.extract_central_part(value) for key, value in arrays.items()}
-            if has_non_existing_pairs:
-                counts = vaex.utils.extract_central_part(counts)
+            counts = vaex.utils.extract_central_part(counts)
 
             # make sure we respect the sorting
-            sorting = tuple(by.sort_indices if by.sort_indices is not None else slice(None) for by in self.by)
-            arrays = {key: value[sorting] for key, value in arrays.items()}
+            def sort(ar):
+                for i, by in list(enumerate(self.by))[::-1]:
+                    sort_indices = by.sort_indices
+                    if sort_indices is not None:
+                        # if sort_indices come from arrow, it will be uint64
+                        # which np.take does not like
+                        sort_indices = vaex.array_types.to_numpy(sort_indices)
+                        if sort_indices.dtype == np.dtype("uint64"):
+                            sort_indices = sort_indices.astype("int64")
+                        ar = np.take(ar, sort_indices, axis=i)
+                return ar
+
+            arrays = {key: sort(value) for key, value in arrays.items()}
 
             if self.combine and self.expand and isinstance(self.by[0], GrouperCombined):
                 assert len(self.by) == 1
@@ -616,18 +674,17 @@ def process(args):
                     assert value.ndim == 1
                     columns[key] = value
             else:
-                if has_non_existing_pairs:
-                    counts = counts[sorting]
-                    mask = counts > 0
-                    coords = [coord[mask] for coord in np.meshgrid(*self._coords1d, indexing='ij')]
-                    columns = {by.label: coord for by, coord in zip(self.by, coords)}
-                    for key, value in arrays.items():
-                        columns[key] = value[mask]
-                else:
-                    columns = {by.label: coord for by, coord in zip(self.by, self._coords1d)}
-                    for key, value in arrays.items():
-                        assert value.ndim == 1
-                        columns[key] = value
+                counts = sort(counts)
+                mask = counts > 0
+                columns = {}
+                for by, indices in zip(self.by, np.where(mask)):
+                    columns[by.label] = by.bin_values.take(indices)
+                if mask.sum() == mask.size:
+                    # if we want all, just take it all
+                    # should be faster
+                    mask = slice(None, None, None)
+                for key, value in arrays.items():
+                    columns[key] = value[mask]
             dataset_arrays = vaex.dataset.DatasetArrays(columns)
             dataset = DatasetGroupby(dataset_arrays, self.df, self.by_original, actions, combine=self.combine, expand=self.expand, sort=self.sort)
             df_grouped = vaex.from_dataset(dataset)


Test patch:
diff --git a/tests/groupby_test.py b/tests/groupby_test.py
--- a/tests/groupby_test.py
+++ b/tests/groupby_test.py
@@ -131,13 +131,15 @@ def test_groupby_sort_string(df_factory, as_category):
 
 
 @pytest.mark.parametrize("auto_encode", [False, True])
-def test_groupby_1d_cat(ds_local, auto_encode):
+@pytest.mark.parametrize("pre_sort", [False, True])
+def test_groupby_1d_cat(ds_local, auto_encode, pre_sort):
     df = ds_local.extract()
-    g = np.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2])
+    g = np.array([0, 0, 0, 0, 2, 2, 1, 1, 1, 1, ])
     df.add_column('g', g)
     df.categorize('g', labels=['cat', 'dog', 'snake'], inplace=True)
     df = df._future() if auto_encode else df
-    dfg = df.groupby(by=df.g, agg='count')
+    grouper = vaex.groupby.GrouperCategory(df.g, sort=True, pre_sort=pre_sort)
+    dfg = df.groupby(by=grouper, agg='count')
 
     assert dfg.g.tolist() == ['cat', 'dog', 'snake']
     assert dfg['count'].tolist() == [4, 4, 2]
@@ -294,14 +296,57 @@ def test_groupby_count():
     assert dfg._equals(dfg2)
 
 
-def test_groupby_boolean(df_factory):
-    # df = df_factory(g=[False, False, True, True, True, None, None], x=[3, None, 4, 5, None, 6, 7])
+@pytest.mark.parametrize("pre_sort", [False, True])
+def test_groupby_with_missing(df_factory, pre_sort):
+    df = df_factory(g=[0, 0, 1, 1, 1, None, None, 2])
+    grouper = vaex.groupby.Grouper(df.g, pre_sort=pre_sort, sort=True)
+    dfg = df.groupby(grouper, agg='count', sort=True, copy=False)
+    assert dfg['g'].tolist() == [0, 1, 2, None]
+    assert dfg['count'].tolist() == [2, 3, 1, 2]
+
+    df = df_factory(g=[0, 0, 1, 1, 1, None, None], x=[3, None, 4, 5, None, 6, 7])
+    grouper = vaex.groupby.Grouper(df.g, pre_sort=pre_sort, sort=True)
+    dfg = df.groupby(grouper, agg={'sum': vaex.agg.sum('x')})
+    assert dfg['g'].tolist() == [0, 1, None]
+    assert dfg['sum'].tolist() == [3, 4+5, 6+7]
+
+
+
+# we don't support pre_sort=False currently (not sure if possible)
+#@pytest.mark.parametrize("pre_sort", [False, True])
+@pytest.mark.parametrize("pre_sort", [True])
+@pytest.mark.parametrize("assume_sparse", [True, False])
+def test_groupby_with_missing_combine(df_factory, pre_sort, assume_sparse):
+    df = df_factory(g1=[0, 0, 1, 1, 1, None, None, 2],
+                    g2=[0, 1, 0, 1, 1, 0,    1,    0],
+    )
+    grouper1 = vaex.groupby.Grouper(df.g1, pre_sort=pre_sort, sort=True)
+    grouper2 = vaex.groupby.Grouper(df.g2, pre_sort=pre_sort, sort=True)
+    groupers = [grouper1, grouper2]
+    dfg = df.groupby(groupers, agg='count', sort=True, copy=False, assume_sparse=assume_sparse)
+    assert dfg['g1'].tolist() == [0, 0, 1, 1, 2, None, None]
+    assert dfg['count'].tolist() == [1, 1, 1, 2, 1, 1, 1]
+
+    df = df_factory(g=[0, 0, 1, 1, 1, None, None], x=[3, None, 4, 5, None, 6, 7])
+    grouper = vaex.groupby.Grouper(df.g, pre_sort=pre_sort, sort=True)
+    dfg = df.groupby(grouper, agg={'sum': vaex.agg.sum('x')})
+    assert dfg['g'].tolist() == [0, 1, None]
+    assert dfg['sum'].tolist() == [3, 4+5, 6+7]
+
+
+def test_groupby_boolean_without_null(df_factory):
     df = df_factory(g=[False, False, True, True, True], x=[3, None, 4, 5, None])
     dfg = df.groupby('g', agg={'sum': vaex.agg.sum('x')}, sort=True)
     assert dfg['g'].tolist() == [False, True]
     assert dfg['sum'].tolist() == [3, 4+5]
 
 
+def test_groupby_boolean_with_null(df_factory):
+    df = df_factory(g=[False, False, True, True, True, None, None], x=[3, None, 4, 5, None, 6, 7])
+    dfg = df.groupby('g', agg={'sum': vaex.agg.sum('x')}, sort=True)
+    assert dfg['g'].tolist() == [False, True, None]
+    assert dfg['sum'].tolist() == [3, 4+5, 6+7]
+
 def test_groupby_std():
     g = np.array([9, 2, 3, 4, 0, 1, 2, 3, 2, 5], dtype='int32')
     s = np.array(list(map(str, [0, 0, 0, 0, 1, 1, 1, 1, 2, 2])))