[C++] data corruption when using `group_by` and `aggregate` on large data sets

[adams-brian]

Describe the bug, including details regarding any error messages, version, and platform.

We recently found some data corruption issues when using group_by and aggregate with large data sets.

I was able to create a minimal reproducible example:

import pyarrow as pa
import pyarrow.compute as pc

COLUMN_COUNT = 5 # <- 4 works fine, 5 causes data corruption
LENGTH = 100_000_000
data = {}
# 'index' = [0, 1, 2, ... , LENGTH-2, LENGTH-1]
data['index'] = pc.indices_nonzero(pc.if_else(True, True, pa.nulls(LENGTH, pa.bool_())))
for i in range(COLUMN_COUNT):
 # fill 'i' with i (ex: '3' = [3, 3, 3, ... 3, 3])
 data[f'{i}'] = pa.nulls(LENGTH, pa.uint64()).fill_null(i)
t = pa.table(data) # <- create table from data
print('-------------------- ORIGINAL --------------------')
print(t)
a = t.group_by(t.column_names).aggregate([]) # <- should behave like a no-op
a = a.combine_chunks() # <- not necessary, just improves the print formatting
print('-------------- GROUP_BY / AGGREGATE --------------')
print(a)

In this example the group_by and aggregate are set up to behave like a no-op and everything works fine with COLUMN_COUNT <= 4:

-------------------- ORIGINAL --------------------
pyarrow.Table
index: uint64
0: uint64
1: uint64
2: uint64
3: uint64
----
index: [[0,1,2,3,4,...,99999995,99999996,99999997,99999998,99999999]]
0: [[0,0,0,0,0,...,0,0,0,0,0]]
1: [[1,1,1,1,1,...,1,1,1,1,1]]
2: [[2,2,2,2,2,...,2,2,2,2,2]]
3: [[3,3,3,3,3,...,3,3,3,3,3]]
-------------- GROUP_BY / AGGREGATE --------------
pyarrow.Table
index: uint64
0: uint64
1: uint64
2: uint64
3: uint64
----
index: [[0,1,2,3,4,...,99999995,99999996,99999997,99999998,99999999]]
0: [[0,0,0,0,0,...,0,0,0,0,0]]
1: [[1,1,1,1,1,...,1,1,1,1,1]]
2: [[2,2,2,2,2,...,2,2,2,2,2]]
3: [[3,3,3,3,3,...,3,3,3,3,3]]

...but results in data corruption if COLUMN_COUNT >= 5:

-------------------- ORIGINAL --------------------
pyarrow.Table
index: uint64
0: uint64
1: uint64
2: uint64
3: uint64
4: uint64
----
index: [[0,1,2,3,4,...,99999995,99999996,99999997,99999998,99999999]]
0: [[0,0,0,0,0,...,0,0,0,0,0]]
1: [[1,1,1,1,1,...,1,1,1,1,1]]
2: [[2,2,2,2,2,...,2,2,2,2,2]]
3: [[3,3,3,3,3,...,3,3,3,3,3]]
4: [[4,4,4,4,4,...,4,4,4,4,4]]
-------------- GROUP_BY / AGGREGATE --------------
pyarrow.Table
index: uint64
0: uint64
1: uint64
2: uint64
3: uint64
4: uint64
----
index: [[0,1,2,3,4,...,3,3,3,3,3]]
0: [[0,0,0,0,0,...,4,4,4,4,4]]
1: [[1,1,1,1,1,...,10521510,10521511,10521512,10521513,10521514]]
2: [[2,2,2,2,2,...,0,0,0,0,0]]
3: [[3,3,3,3,3,...,1,1,1,1,1]]
4: [[4,4,4,4,4,...,2,2,2,2,2]]

Component(s)

Python

Version

pyarrow==12.0.1. I downgraded to pyarrow==11, pyarrow==10, pyarrow==9, and pyarrow==8 and observed the above behavior in those versions as well.

Platform

Linux (Ubuntu 22.04 LTS)

[westonpace]

Hmm...given how stable this is I wonder if we're overflowing an int32 somewhere.

[westonpace]

Indeed, this was the case. In the internals of the "group by" code is a hash map. This hash map stores keys in row-major order using the row-format (row-major order is commonly used with hash-maps since it's expected they will be accessed non-sequentially). This structure is called the RowTableImpl:

arrow/cpp/src/arrow/compute/row/row_internal.h

Line 162 in 9d92ed4

class ARROW_EXPORT RowTableImpl {

In the example we are using 6 columns (if we include the index) of int8 as the hash keys. This means each row in our row table is 48 bytes large. The test uses 100,000,000 rows. So we end up with 4,800,000,000 bytes which is over the int32 limit of 4,294,967,296. With 5 columns we only have 4,000,000,000 bytes and so it fits.

This should be ok. The RowTableImpl appears to have been designed to store 2^64 bytes of data. In practice we do not have any unit tests covering this case and so, unsurprisingly, we have mistakes. The following patch fixes the reproduction from the OP: c901d1e

Unfortunately, I still have pretty low confidence in the robustness of this code when working with tables this large. There are different paths for boolean fields, dictionary fields, string-like fields, and all of these paths are futher parameterized by AVX2 vs. not-AVX2 and we no longer have CI for AVX2.

We are missing a set of unit tests to cover these scenarios and I don't think I'll have time to add them for a while. So I'm uncertain at the moment whether it would be better to put in this short-term fix or to simply reject and fail a plan once the keys table exceeds 4GB. If anyone is interested in taking on the effort of adding these unit tests I'd be happy to explain further. In the meantime I'm not sure when I'll get to it. I'm marking this blocker to force myself to make some kind of decision on this in the 13.0.0 timeframe.

[raulcd]

@westonpace do you think we should prioritize this to have a fix for 13.0.0? or is this not a blocker for the release?

[raulcd]

@jorisvandenbossche Should this be a release blocker?

[jorisvandenbossche]

@westonpace I assume this is not a new issue? In which case I don't think it would be a blocker

[westonpace]

It would be good to at least change this to return an error in time for 13.0.0. It's not a new issue but it is "invalid data" which is pretty bad. I agree it is probably not a blocker.