Initial implementation of read-only buffer access to raw tables

ChangeLog

@@ -1,3 +1,9 @@
+2015-03-10  Camille Scott  <camille.scott.w@gmail.com>
+
+   * lib/counting.hh, khmer/_khmermodule.cc: Expose the raw tables of
+   count-min sketches to the world of python using a buffer interface.
+   * tests/test_counting_hash.py: Tests of the above functionality.
+
 2015-03-08  Michael R. Crusoe  <mcrusoe@msu.edu>
 
    * Makefile: make 'pep8' target be more verbose

khmer/_khmermodule.cc

@@ -682,6 +682,27 @@ PyObject *
 hash_abundance_distribution_with_reads_parser(khmer_KCountingHash_Object * me,
         PyObject * args);
 
+static
+PyObject *
+hash_get_raw_tables(khmer_KCountingHash_Object * self, PyObject * args)
+{
+    CountingHash * counting = self->counting;
+
+    Byte ** table_ptrs = counting->get_raw_tables();
+    std::vector<HashIntoType> sizes = counting->get_tablesizes();
+
+    PyObject * raw_tables = PyList_New(sizes.size());
+    for (unsigned int i=0; i<sizes.size(); ++i) {
+        PyObject * buf = PyBuffer_FromMemory(table_ptrs[i], sizes[i]);
+        if(!PyBuffer_Check(buf)) {
+            return NULL;
+        }
+        PyList_SET_ITEM(raw_tables, i, buf);
+    }
+
+    return raw_tables;
+}
+
 static
 PyObject *
 hash_set_use_bigcount(khmer_KCountingHash_Object * me, PyObject * args)
@@ -1532,6 +1553,10 @@ static PyMethodDef khmer_counting_methods[] = {
     },
     { "output_fasta_kmer_pos_freq", (PyCFunction)hash_output_fasta_kmer_pos_freq, METH_VARARGS, "" },
     { "get", (PyCFunction)hash_get, METH_VARARGS, "Get the count for the given k-mer" },
+    {
+        "get_raw_tables", (PyCFunction)hash_get_raw_tables,
+        METH_VARARGS, "Get a list of the raw tables as memoryview objects"
+    },
     { "get_min_count", (PyCFunction)hash_get_min_count, METH_VARARGS, "Get the smallest count of all the k-mers in the string" },
     { "get_max_count", (PyCFunction)hash_get_max_count, METH_VARARGS, "Get the largest count of all the k-mers in the string" },
     { "get_median_count", (PyCFunction)hash_get_median_count, METH_VARARGS, "Get the median, average, and stddev of the k-mer counts in the string" },

lib/counting.hh

@@ -87,6 +87,13 @@ public:
         }
     }
 
+    // Writing to the tables outside of defined methods has undefined behavior!
+    // As such, this should only be used to return read-only interfaces
+    Byte ** get_raw_tables()
+    {
+        return _counts;
+    }
+
     virtual BoundedCounterType test_and_set_bits(const char * kmer)
     {
         BoundedCounterType x = get_count(kmer); // @CTB just hash it, yo.

tests/test_counting_hash.py

@@ -101,6 +101,27 @@ def test_collision_3(self):
         assert hi.get(GG) == 2
 
 
+def test_get_raw_tables():
+    ht = khmer.new_counting_hash(20, 1e5, 4)
+    tables = ht.get_raw_tables()
+
+    for size, table in zip(ht.hashsizes(), tables):
+        assert isinstance(table, buffer)
+        assert size == len(table)
+
+
+def test_get_raw_tables_view():
+    ht = khmer.new_counting_hash(20, 1e5, 4)
+    tables = ht.get_raw_tables()
+    for tab in tables:
+        memv = memoryview(tab)
+        assert sum(memv.tolist()) == 0
+    ht.consume('AAAATTTTCCCCGGGGAAAA')
+    for tab in tables:
+        memv = memoryview(tab)
+        assert sum(memv.tolist()) == 1
+
+
 @attr('linux')
 def test_toobig():
     try: