Fixing NumPy/Pandas String fetching and add test case based on #448

duckdb · Mar 7, 2020 · bc8a11d · bc8a11d
1 parent f126192
commit bc8a11d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/tools/pythonpkg/cursor.cpp b/tools/pythonpkg/cursor.cpp
@@ -490,7 +490,7 @@ PyObject *duckdb_cursor_fetchnumpy(duckdb_Cursor *self) {
 					assert(!chunk->data[col_idx].sel_vector);
 					PyObject *str_obj;
 					if (!mask_data[chunk_idx + offset]) {
-						str_obj = PyUnicode_FromString(((const char **)chunk->data[col_idx].GetData())[chunk_idx]);
+						str_obj = PyUnicode_FromString(((duckdb::string_t*)chunk->data[col_idx].GetData())[chunk_idx].GetData());
 					} else {
 						assert(cols[col_idx].found_nil);
 						str_obj = Py_None;

diff --git a/tools/pythonpkg/tests/test_dbapi08.py b/tools/pythonpkg/tests/test_dbapi08.py
@@ -0,0 +1,17 @@
+# test fetchdf with various types
+import pandas
+import numpy
+
+class TestType(object):
+    def test_fetchdf(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE items(item VARCHAR)")
+        duckdb_cursor.execute("INSERT INTO items VALUES ('jeans'), (''), (NULL)")
+        res = duckdb_cursor.execute("SELECT item FROM items").fetchdf()
+        assert isinstance(res, pandas.DataFrame)
+
+        arr = numpy.ma.masked_array(['jeans', '', None])
+        arr.mask = [False, False, True]
+        arr = {'item': arr}
+        df = pandas.DataFrame.from_dict(arr)
+
+        pandas.testing.assert_frame_equal(res, df)