COMPAT/PERF add checks, adapt test, fix data.m after realloc

Author: mattip

pandas/_libs/hashtable_class_helper.pxi.in

@@ -85,11 +85,15 @@ cdef class {{name}}Vector:
  self.data.n = 0
  self.data.m = _INIT_VEC_CAP
  self.data.data = <{{arg}}*> malloc(self.data.m * sizeof({{arg}}))
+ if not self.data.data:
+ raise MemoryError()
 
  cdef void resize(self) nogil:
- # TODO: handle failure to allocate
  self.data.m = max(self.data.m * 4, _INIT_VEC_CAP)
  self.data.data = <{{arg}}*> realloc(self.data.data, self.data.m * sizeof({{arg}}))
+ if not self.data.data:
+ with gil:
+ raise MemoryError()
 
  def __dealloc__(self):
  if self.data is not NULL:
@@ -110,6 +114,7 @@ cdef class {{name}}Vector:
  raise ValueError("Vector.to_array() can only be called once")
 
  self.data.data = <{{arg}}*> realloc(self.data.data, self.data.n * sizeof({{arg}}))
+ self.data.m = self.data.n
  self.external_view_exists = True
  shape[0] = self.data.n
  ao = cnp.PyArray_SimpleNewFromData(1, shape, {{idtype}}, <void*>self.data.data)
@@ -120,6 +125,8 @@ cdef class {{name}}Vector:
  cdef inline void append(self, {{arg}} x):
 
  if needs_resize(self.data):
+ if self.external_view_exists:
+ raise ValueError("external reference but Vector.resize() needed")
  self.resize()
 
  append_data_{{dtype}}(self.data, x)
@@ -145,12 +152,17 @@ cdef class StringVector:
  self.data.n = 0
  self.data.m = _INIT_VEC_CAP
  self.data.data = <char **> malloc(self.data.m * sizeof(char *))
+ if not self.data.data:
+ raise MemoryError()
 
  cdef void resize(self) nogil:
  self.data.m = max(self.data.m * 4, _INIT_VEC_CAP)
 
  self.data.data = <char **> realloc(self.data.data,
  self.data.m * sizeof(char *))
+ if not self.data.data:
+ with gil:
+ raise MemoryError()
 
  def __dealloc__(self):
  if self.data is not NULL:
@@ -183,6 +195,8 @@ cdef class StringVector:
  cdef inline void append(self, char * x):
 
  if needs_resize(self.data):
+ if self.external_view_exists:
+ raise ValueError("external reference but Vector.resize() needed")
  self.resize()
 
  append_data_string(self.data, x)
@@ -208,6 +222,8 @@ cdef class ObjectVector:
 
  cdef inline append(self, object o):
  if self.n == self.m:
+ if self.external_view_exists:
+ raise ValueError("external reference but Vector.resize() needed")
  self.m = max(self.m * 2, _INIT_VEC_CAP)
  self.ao.resize(self.m, refcheck=False)
  self.data = <PyObject**> self.ao.data
@@ -219,7 +235,7 @@ cdef class ObjectVector:
  def to_array(self):
  if self.external_view_exists:
  raise ValueError("Vector.to_array() can only be called once")
- self.ao.resize(self.n)
+ self.ao.resize(self.n, refcheck=False)
  self.m = self.n
  self.external_view_exists = True
  return self.ao
@@ -380,6 +396,7 @@ cdef class {{name}}HashTable(HashTable):
 
  if needs_resize(ud):
  uniques.resize()
+ ud = uniques.data
  append_data_{{dtype}}(ud, val)
  labels[i] = count
  count += 1
@@ -454,18 +471,21 @@ cdef class {{name}}HashTable(HashTable):
  kh_put_{{dtype}}(self.table, val, &ret)
  if needs_resize(ud):
  uniques.resize()
+ ud = uniques.data
  append_data_{{dtype}}(ud, val)
  elif not seen_na:
  seen_na = 1
  if needs_resize(ud):
  uniques.resize()
+ ud = uniques.data
  append_data_{{dtype}}(ud, NAN)
  {{else}}
  k = kh_get_{{dtype}}(self.table, val)
  if k == self.table.n_buckets:
  kh_put_{{dtype}}(self.table, val, &ret)
  if needs_resize(ud):
  uniques.resize()
+ ud = uniques.data
  append_data_{{dtype}}(ud, val)
  {{endif}}
 

pandas/tests/test_algos.py

@@ -1070,25 +1070,37 @@ def test_vector_resize(self):
  # Test for memory errors after internal vector
  # reallocations (pull request #7157)
 
- def _test_vector_resize(htable, uniques, dtype, nvals):
+ def _test_vector_resize(htable, uniques, dtype, nvals, safely_resizes):
  vals = np.array(np.random.randn(1000), dtype=dtype)
  # get_labels appends to the vector
  htable.get_labels(vals[:nvals], uniques, 0, -1)
- # to_array resizes the vector
- uniques.to_array()
- htable.get_labels(vals, uniques, 0, -1)
+ # to_array may resize the vector
+ tmp = uniques.to_array()
+ oldshape = tmp.shape
+ if safely_resizes:
+ htable.get_labels(vals, uniques, 0, -1)
+ else:
+ with pytest.raises(ValueError) as excinfo:
+ htable.get_labels(vals, uniques, 0, -1)
+ assert excinfo.value.message == 'external reference but Vector.resize() needed'
+ with pytest.raises(ValueError) as excinfo:
+ # cannot resize, tmp is holding a reference
+ tmp2 = uniques.to_array()
+ assert excinfo.value.message == 'Vector.to_array() can only be called once'
+ assert tmp.shape == oldshape
 
  test_cases = [
- (hashtable.PyObjectHashTable, hashtable.ObjectVector, 'object'),
- (hashtable.StringHashTable, hashtable.ObjectVector, 'object'),
- (hashtable.Float64HashTable, hashtable.Float64Vector, 'float64'),
- (hashtable.Int64HashTable, hashtable.Int64Vector, 'int64'),
- (hashtable.UInt64HashTable, hashtable.UInt64Vector, 'uint64')]
+ (hashtable.PyObjectHashTable, hashtable.ObjectVector, 'object', False),
+ (hashtable.StringHashTable, hashtable.ObjectVector, 'object', True),
+ (hashtable.Float64HashTable, hashtable.Float64Vector, 'float64', True),
+ (hashtable.Int64HashTable, hashtable.Int64Vector, 'int64', True),
+ (hashtable.UInt64HashTable, hashtable.UInt64Vector, 'uint64', True)]
 
- for (tbl, vect, dtype) in test_cases:
+ for (tbl, vect, dtype, safely_resizes) in test_cases:
  # resizing to empty is a special case
- _test_vector_resize(tbl(), vect(), dtype, 0)
- _test_vector_resize(tbl(), vect(), dtype, 10)
+ print dtype, safely_resizes
+ _test_vector_resize(tbl(), vect(), dtype, 0, safely_resizes)
+ _test_vector_resize(tbl(), vect(), dtype, 10, safely_resizes)
 
 
 def test_quantile():