ENH: making some more progress. string factorization conversion

Author: wesm

pandas/io/tests/test_cparser.py

@@ -54,11 +54,26 @@ def test_file_handle(self):
         finally:
             f.close()
 
-    # def test_StringIO(self):
-    #     text = open(self.csv1, 'rb').read()
+    def test_file_handle_mmap(self):
+        try:
+            f = open(self.csv1, 'rb')
+            reader = parser.TextReader(f, memory_map=True)
+            result = reader.read()
+        finally:
+            f.close()
+
+    def test_StringIO(self):
+        text = open(self.csv1, 'rb').read()
+        reader = parser.TextReader(BytesIO(text))
+        result = reader.read()
+
+    def test_string_factorize(self):
+        # should this be optional?
+        data = 'a\nb\na\nb\na'
+        reader = parser.TextReader(StringIO(data))
+        result = reader.read()
+        self.assert_(len(set(map(id, result[0]))) == 2)
 
-    #     reader = parser.TextReader(BytesIO(text))
-    #     result = reader.read()
 
 if __name__ == '__main__':
     nose.runmodule(argv=[__file__,'-vvs','-x','--pdb', '--pdb-failure'],

pandas/src/khash.h

@@ -613,6 +613,8 @@ KHASH_SET_INIT_PYOBJECT(pyset)
 
 KHASH_MAP_INIT_STR(str, Py_ssize_t)
 
+KHASH_MAP_INIT_STR(strbox, kh_pyobject_t)
+
 KHASH_MAP_INIT_INT(int32, Py_ssize_t)
 KHASH_MAP_INIT_INT64(int64, Py_ssize_t)
 KHASH_MAP_INIT_FLOAT64(float64, Py_ssize_t)

pandas/src/khash.pxd

@@ -55,6 +55,7 @@ cdef extern from "khash.h":
 
     bint kh_exist_str(kh_str_t*, khiter_t)
 
+
     ctypedef struct kh_int64_t:
         khint_t n_buckets, size, n_occupied, upper_bound
         uint32_t *flags
@@ -102,3 +103,22 @@ cdef extern from "khash.h":
     inline void kh_del_int32(kh_int32_t*, khint_t)
 
     bint kh_exist_int32(kh_int32_t*, khiter_t)
+
+    # sweep factorize
+
+    ctypedef struct kh_strbox_t:
+        khint_t n_buckets, size, n_occupied, upper_bound
+        uint32_t *flags
+        kh_cstr_t *keys
+        PyObject **vals
+
+    inline kh_strbox_t* kh_init_strbox()
+    inline void kh_destroy_strbox(kh_strbox_t*)
+    inline void kh_clear_strbox(kh_strbox_t*)
+    inline khint_t kh_get_strbox(kh_strbox_t*, kh_cstr_t)
+    inline void kh_resize_strbox(kh_strbox_t*, khint_t)
+    inline khint_t kh_put_strbox(kh_strbox_t*, kh_cstr_t, int*)
+    inline void kh_del_strbox(kh_strbox_t*, khint_t)
+
+    bint kh_exist_strbox(kh_strbox_t*, khiter_t)
+

pandas/src/parser.pyx

@@ -3,6 +3,10 @@ import numpy as np
 
 cnp.import_array()
 
+from khash cimport *
+
+from cpython cimport PyString_FromString, Py_INCREF, PyString_AsString
+
 cdef extern from "Python.h":
     ctypedef struct FILE
     FILE* PyFile_AsFile(object)
@@ -143,8 +147,10 @@ cdef class TextReader:
 
     cdef public:
         object delimiter, na_values, converters, thousands, delim_whitespace
+        object memory_map
 
-    def __cinit__(self, source, delimiter=',', header=0, memory_map=False,
+    def __cinit__(self, source, delimiter=',', header=0,
+                  memory_map=False,
                   chunksize=DEFAULT_CHUNKSIZE,
                   delim_whitespace=False,
                   na_values=None,
@@ -173,6 +179,7 @@ cdef class TextReader:
         self.delimiter = delimiter
         self.delim_whitespace = delim_whitespace
 
+        self.memory_map = memory_map
         self.na_values = na_values
         self.converters = converters
         self.thousands = thousands
@@ -185,16 +192,22 @@ cdef class TextReader:
             self.file_handle.close()
 
     cdef _setup_parser_source(self, source):
-        cdef int status
+        cdef:
+            int status
 
         if isinstance(source, (basestring, file)):
             if isinstance(source, basestring):
                 source = open(source, 'rb')
                 self.should_close = True
 
             self.file_handle = source
-            status = parser_file_source_init(self.parser,
-                                             PyFile_AsFile(source))
+
+            if self.memory_map:
+                status = parser_mmap_init(self.parser,
+                                          PyFile_AsFile(source))
+            else:
+                status = parser_file_source_init(self.parser,
+                                                 PyFile_AsFile(source))
 
             if status != 0:
                 raise Exception('Initializing from file failed')
@@ -209,7 +222,8 @@ cdef class TextReader:
                 raise ValueError('Only ascii/bytes supported at the moment')
 
             status = parser_array_source_init(self.parser,
-                                              <char*> bytes, len(bytes))
+                                              PyString_AsString(bytes),
+                                              len(bytes))
             if status != 0:
                 raise Exception('Initializing parser from file-like '
                                 'object failed')
@@ -248,6 +262,10 @@ cdef class TextReader:
         for i in range(ncols):
             col_res = _try_double(self.parser, i, 0, self.parser.lines)
 
+            if col_res is None:
+                col_res = _string_box_factorize(self.parser, i,
+                                                0, self.parser.lines)
+
             results[i] = col_res
 
         return results
@@ -256,6 +274,54 @@ class CParserError(Exception):
     pass
 
 
+# ----------------------------------------------------------------------
+# Type conversions / inference support code
+
+cdef _string_box_factorize(parser_t *parser, int col,
+                           int line_start, int line_end):
+    cdef:
+        int error
+        Py_ssize_t i
+        size_t lines
+        coliter_t it
+        char *word
+        cnp.ndarray[object] result
+
+        int ret = 0
+        kh_strbox_t *table
+        kh_iter_t
+
+        object pyval
+
+
+    table = kh_init_strbox()
+
+    lines = line_end - line_start
+    result = np.empty(lines, dtype=np.object_)
+
+    coliter_setup(&it, parser, col)
+    for i in range(lines):
+        word = COLITER_NEXT(it)
+
+        k = kh_get_strbox(table, word)
+
+        # in the hash table
+        if k != table.n_buckets:
+            # this increments the refcount, but need to test
+            pyval = <object> table.vals[k]
+        else:
+            # box it. new ref?
+            pyval = PyString_FromString(word)
+
+            k = kh_put_strbox(table, word, &ret)
+            table.vals[k] = <PyObject*> pyval
+
+        result[i] = pyval
+
+    return result
+
+
+
 cdef _try_double(parser_t *parser, int col, int line_start, int line_end):
     cdef:
         int error