Merge pull request #4 from pcmoritz/key_tuples

Tuple keys in dicts can be serialized

Author: robertnishihara

cpp/src/numbuf/dict.cc

@@ -5,13 +5,14 @@ using namespace arrow;
 namespace numbuf {
 
 std::shared_ptr<arrow::StructArray> DictBuilder::Finish(
-    std::shared_ptr<Array> list_data,
-    std::shared_ptr<Array> tuple_data,
-    std::shared_ptr<Array> dict_data) {
+    std::shared_ptr<Array> key_tuple_data,
+    std::shared_ptr<Array> val_list_data,
+    std::shared_ptr<Array> val_tuple_data,
+    std::shared_ptr<Array> val_dict_data) {
   // lists and dicts can't be keys of dicts in Python, that is why for
   // the keys we do not need to collect sublists
-  auto keys = keys_.Finish(nullptr, nullptr, nullptr);
-  auto vals = vals_.Finish(list_data, tuple_data, dict_data);
+  auto keys = keys_.Finish(nullptr, key_tuple_data, nullptr);
+  auto vals = vals_.Finish(val_list_data, val_tuple_data, val_dict_data);
   auto keys_field = std::make_shared<Field>("keys", keys->type());
   auto vals_field = std::make_shared<Field>("vals", vals->type());
   auto type = std::make_shared<StructType>(std::vector<FieldPtr>({keys_field, vals_field}));

cpp/src/numbuf/dict.h

@@ -34,9 +34,10 @@ class DictBuilder {
         value list of the dictionary
   */
   std::shared_ptr<arrow::StructArray> Finish(
-    std::shared_ptr<arrow::Array> list_data,
-    std::shared_ptr<arrow::Array> tuple_data,
-    std::shared_ptr<arrow::Array> dict_data);
+    std::shared_ptr<arrow::Array> key_tuple_data,
+    std::shared_ptr<arrow::Array> val_list_data,
+    std::shared_ptr<arrow::Array> val_tuple_data,
+    std::shared_ptr<arrow::Array> val_dict_data);
 
 private:
   SequenceBuilder keys_;

python/src/pynumbuf/adapters/python.cc

@@ -171,29 +171,33 @@ Status DeserializeTuple(std::shared_ptr<Array> array, int32_t start_idx, int32_t
 
 Status SerializeDict(std::vector<PyObject*> dicts, std::shared_ptr<Array>* out) {
   DictBuilder result;
-  std::vector<PyObject*> sublists, subtuples, subdicts, dummy;
+  std::vector<PyObject*> key_tuples, val_lists, val_tuples, val_dicts, dummy;
   for (const auto& dict : dicts) {
     PyObject *key, *value;
     Py_ssize_t pos = 0;
     while (PyDict_Next(dict, &pos, &key, &value)) {
-      RETURN_NOT_OK(append(key, result.keys(), dummy, dummy, dummy));
+      RETURN_NOT_OK(append(key, result.keys(), dummy, key_tuples, dummy));
       DCHECK(dummy.size() == 0);
-      RETURN_NOT_OK(append(value, result.vals(), sublists, subtuples, subdicts));
+      RETURN_NOT_OK(append(value, result.vals(), val_lists, val_tuples, val_dicts));
     }
   }
-  std::shared_ptr<Array> val_list;
-  if (sublists.size() > 0) {
-    RETURN_NOT_OK(SerializeSequences(sublists, &val_list));
+  std::shared_ptr<Array> key_tuples_arr;
+  if (key_tuples.size() > 0) {
+    RETURN_NOT_OK(SerializeSequences(key_tuples, &key_tuples_arr));
   }
-  std::shared_ptr<Array> val_tuples;
-  if (subtuples.size() > 0) {
-    RETURN_NOT_OK(SerializeSequences(subtuples, &val_tuples));
+  std::shared_ptr<Array> val_list_arr;
+  if (val_lists.size() > 0) {
+    RETURN_NOT_OK(SerializeSequences(val_lists, &val_list_arr));
   }
-  std::shared_ptr<Array> val_dict;
-  if (subdicts.size() > 0) {
-    RETURN_NOT_OK(SerializeDict(subdicts, &val_dict));
+  std::shared_ptr<Array> val_tuples_arr;
+  if (val_tuples.size() > 0) {
+    RETURN_NOT_OK(SerializeSequences(val_tuples, &val_tuples_arr));
+  }
+  std::shared_ptr<Array> val_dict_arr;
+  if (val_dicts.size() > 0) {
+    RETURN_NOT_OK(SerializeDict(val_dicts, &val_dict_arr));
   }
-  *out = result.Finish(val_list, val_tuples, val_dict);
+  *out = result.Finish(key_tuples_arr, val_list_arr, val_tuples_arr, val_dict_arr);
   return Status::OK();
 }
 

python/test/runtest.py

@@ -3,8 +3,9 @@
 import numpy as np
 from numpy.testing import assert_equal
 
-TEST_OBJECTS = [[1, "hello", 3.0], 42, 43L, "hello world", u"x", u"\u262F",
-                42.0, 1L << 62, (1.0, "hi"),
+TEST_OBJECTS = [{(1,2) : 1}, {() : 2}, [1, "hello", 3.0], 42, 43L, "hello world", 
+                u"x", u"\u262F", 42.0, 
+                1L << 62, (1.0, "hi"),
                 None, (None, None), ("hello", None),
                 True, False, (True, False), "hello",
                 {True: "hello", False: "world"},