extending khash-maps with tracemalloc ability (works for >=Py3.6), similar to the way numpy does

Author: realead

pandas/_libs/hashtable.pyx

@@ -13,10 +13,14 @@ cnp.import_array()
 
 
 from pandas._libs cimport util
-from pandas._libs.khash cimport kh_str_t, khiter_t
+from pandas._libs.khash cimport KHASH_TRACE_DOMAIN, kh_str_t, khiter_t
 from pandas._libs.missing cimport checknull
 
 
+def get_hashtable_trace_domain():
+    return KHASH_TRACE_DOMAIN
+
+
 cdef int64_t NPY_NAT = util.get_nat()
 SIZE_HINT_LIMIT = (1 << 20) + 7
 

pandas/_libs/khash.pxd

@@ -14,6 +14,8 @@ from numpy cimport (
 
 
 cdef extern from "khash_python.h":
+    const int KHASH_TRACE_DOMAIN
+
     ctypedef uint32_t khint_t
     ctypedef khint_t khiter_t
 

pandas/_libs/src/klib/khash.h

@@ -115,6 +115,24 @@ int main() {
 #include "../inline_helper.h"
 
 
+// hooks for memory allocator, C-runtime allocator used per default
+#ifndef KHASH_MALLOC
+#define KHASH_MALLOC malloc
+#endif
+
+#ifndef KHASH_REALLOC
+#define KHASH_REALLOC realloc
+#endif
+
+#ifndef KHASH_CALLOC
+#define KHASH_CALLOC calloc
+#endif
+
+#ifndef KHASH_FREE
+#define KHASH_FREE free
+#endif
+
+
 #if UINT_MAX == 0xffffffffu
 typedef unsigned int khint32_t;
 #elif ULONG_MAX == 0xffffffffu
@@ -265,14 +283,14 @@ static const double __ac_HASH_UPPER = 0.77;
 		khval_t *vals;													\
 	} kh_##name##_t;													\
 	SCOPE kh_##name##_t *kh_init_##name(void) {								\
-		return (kh_##name##_t*)calloc(1, sizeof(kh_##name##_t));		\
+		return (kh_##name##_t*)KHASH_CALLOC(1, sizeof(kh_##name##_t));		\
 	}																	\
 	SCOPE void kh_destroy_##name(kh_##name##_t *h)						\
 	{																	\
 		if (h) {														\
-			free(h->keys); free(h->flags);								\
-			free(h->vals);												\
-			free(h);													\
+			KHASH_FREE(h->keys); KHASH_FREE(h->flags);								\
+			KHASH_FREE(h->vals);												\
+			KHASH_FREE(h);													\
 		}																\
 	}																	\
 	SCOPE void kh_clear_##name(kh_##name##_t *h)						\
@@ -305,11 +323,11 @@ static const double __ac_HASH_UPPER = 0.77;
 			if (new_n_buckets < 4) new_n_buckets = 4;					\
 			if (h->size >= (khint_t)(new_n_buckets * __ac_HASH_UPPER + 0.5)) j = 0;	/* requested size is too small */ \
 			else { /* hash table size to be changed (shrink or expand); rehash */ \
-				new_flags = (khint32_t*)malloc(__ac_fsize(new_n_buckets) * sizeof(khint32_t));	\
+				new_flags = (khint32_t*)KHASH_MALLOC(__ac_fsize(new_n_buckets) * sizeof(khint32_t));	\
 				memset(new_flags, 0xff, __ac_fsize(new_n_buckets) * sizeof(khint32_t)); \
 				if (h->n_buckets < new_n_buckets) {	/* expand */		\
-					h->keys = (khkey_t*)realloc(h->keys, new_n_buckets * sizeof(khkey_t)); \
-					if (kh_is_map) h->vals = (khval_t*)realloc(h->vals, new_n_buckets * sizeof(khval_t)); \
+					h->keys = (khkey_t*)KHASH_REALLOC(h->keys, new_n_buckets * sizeof(khkey_t)); \
+					if (kh_is_map) h->vals = (khval_t*)KHASH_REALLOC(h->vals, new_n_buckets * sizeof(khval_t)); \
 				} /* otherwise shrink */								\
 			}															\
 		}																\
@@ -342,10 +360,10 @@ static const double __ac_HASH_UPPER = 0.77;
 				}														\
 			}															\
 			if (h->n_buckets > new_n_buckets) { /* shrink the hash table */ \
-				h->keys = (khkey_t*)realloc(h->keys, new_n_buckets * sizeof(khkey_t)); \
-				if (kh_is_map) h->vals = (khval_t*)realloc(h->vals, new_n_buckets * sizeof(khval_t)); \
+				h->keys = (khkey_t*)KHASH_REALLOC(h->keys, new_n_buckets * sizeof(khkey_t)); \
+				if (kh_is_map) h->vals = (khval_t*)KHASH_REALLOC(h->vals, new_n_buckets * sizeof(khval_t)); \
 			}															\
-			free(h->flags); /* free the working space */				\
+			KHASH_FREE(h->flags); /* free the working space */				\
 			h->flags = new_flags;										\
 			h->n_buckets = new_n_buckets;								\
 			h->n_occupied = h->size;									\

pandas/_libs/src/klib/khash_python.h

@@ -1,6 +1,59 @@
 #include <string.h>
 #include <Python.h>
 
+// khash should report usage to tracemalloc
+#if PY_VERSION_HEX >= 0x03060000
+#include <pymem.h>
+#if PY_VERSION_HEX < 0x03070000
+#define PyTraceMalloc_Track _PyTraceMalloc_Track
+#define PyTraceMalloc_Untrack _PyTraceMalloc_Untrack
+#endif
+#else
+#define PyTraceMalloc_Track(...)
+#define PyTraceMalloc_Untrack(...)
+#endif
+
+
+static const int KHASH_TRACE_DOMAIN = 424242;
+void *traced_malloc(size_t size){
+    void * ptr = malloc(size);
+    if(ptr!=NULL){
+        PyTraceMalloc_Track(KHASH_TRACE_DOMAIN, (uintptr_t)ptr, size);
+    }
+    return ptr;
+}
+
+void *traced_calloc(size_t num, size_t size){
+    void * ptr = calloc(num, size);
+    if(ptr!=NULL){
+        PyTraceMalloc_Track(KHASH_TRACE_DOMAIN, (uintptr_t)ptr, num*size);
+    }
+    return ptr;
+}
+
+void *traced_realloc(void* old_ptr, size_t size){
+    void * ptr = realloc(old_ptr, size);
+    if(ptr!=NULL){
+        if(old_ptr != ptr){
+            PyTraceMalloc_Untrack(KHASH_TRACE_DOMAIN, (uintptr_t)old_ptr);
+        }
+        PyTraceMalloc_Track(KHASH_TRACE_DOMAIN, (uintptr_t)ptr, size);
+    }
+    return ptr;
+}
+
+void traced_free(void* ptr){
+    if(ptr!=NULL){
+        PyTraceMalloc_Untrack(KHASH_TRACE_DOMAIN, (uintptr_t)ptr);
+    }
+    free(ptr);
+}
+
+
+#define KHASH_MALLOC traced_malloc
+#define KHASH_REALLOC traced_realloc
+#define KHASH_CALLOC traced_calloc
+#define KHASH_FREE traced_free
 #include "khash.h"
 
 // Previously we were using the built in cpython hash function for doubles
@@ -128,7 +181,7 @@ typedef struct {
 typedef kh_str_starts_t* p_kh_str_starts_t;
 
 p_kh_str_starts_t PANDAS_INLINE kh_init_str_starts(void) {
-	kh_str_starts_t *result = (kh_str_starts_t*)calloc(1, sizeof(kh_str_starts_t));
+	kh_str_starts_t *result = (kh_str_starts_t*)KHASH_CALLOC(1, sizeof(kh_str_starts_t));
 	result->table = kh_init_str();
 	return result;
 }
@@ -151,7 +204,7 @@ khint_t PANDAS_INLINE kh_get_str_starts_item(const kh_str_starts_t* table, const
 
 void PANDAS_INLINE kh_destroy_str_starts(kh_str_starts_t* table) {
 	kh_destroy_str(table->table);
-	free(table);
+	KHASH_FREE(table);
 }
 
 void PANDAS_INLINE kh_resize_str_starts(kh_str_starts_t* table, khint_t val) {

pandas/tests/libs/test_hashtable.py

@@ -1,3 +1,6 @@
+from contextlib import contextmanager
+import tracemalloc
+
 import numpy as np
 import pytest
 
@@ -6,6 +9,23 @@
 import pandas._testing as tm
 
 
+@contextmanager
+def activated_tracemalloc():
+    tracemalloc.start()
+    try:
+        yield
+    finally:
+        tracemalloc.stop()
+
+
+def get_allocated_khash_memory():
+    snapshot = tracemalloc.take_snapshot()
+    snapshot = snapshot.filter_traces(
+        (tracemalloc.DomainFilter(True, ht.get_hashtable_trace_domain()),)
+    )
+    return sum(map(lambda x: x.size, snapshot.traces))
+
+
 @pytest.mark.parametrize(
     "table_type, dtype",
     [
@@ -101,6 +121,21 @@ def test_unique(self, table_type, dtype):
         unique = table.unique(keys)
         tm.assert_numpy_array_equal(unique, expected)
 
+    def test_tracemalloc_works(self, table_type, dtype):
+        if dtype in (np.int8, np.uint8):
+            N = 256
+        else:
+            N = 30000
+        keys = np.arange(N).astype(dtype)
+        with activated_tracemalloc():
+            table = table_type()
+            table.map_locations(keys)
+            used = get_allocated_khash_memory()
+            my_size = table.sizeof()
+            assert used == my_size
+            del table
+            assert get_allocated_khash_memory() == 0
+
 
 @pytest.mark.parametrize(
     "table_type, dtype",