mirror of
https://github.com/wassname/ray.git
synced 2026-07-23 13:10:11 +08:00
Tuples as keys in dicts can be serialized
This commit is contained in:
@@ -5,12 +5,13 @@ using namespace arrow;
|
||||
namespace numbuf {
|
||||
|
||||
std::shared_ptr<arrow::StructArray> DictBuilder::Finish(
|
||||
std::shared_ptr<Array> key_tuple_data,
|
||||
std::shared_ptr<Array> list_data,
|
||||
std::shared_ptr<Array> tuple_data,
|
||||
std::shared_ptr<Array> 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 keys = keys_.Finish(nullptr, key_tuple_data, nullptr);
|
||||
auto vals = vals_.Finish(list_data, tuple_data, dict_data);
|
||||
auto keys_field = std::make_shared<Field>("keys", keys->type());
|
||||
auto vals_field = std::make_shared<Field>("vals", vals->type());
|
||||
|
||||
@@ -34,6 +34,7 @@ public:
|
||||
value list of the dictionary
|
||||
*/
|
||||
std::shared_ptr<arrow::StructArray> Finish(
|
||||
std::shared_ptr<arrow::Array> key_tuple_data,
|
||||
std::shared_ptr<arrow::Array> list_data,
|
||||
std::shared_ptr<arrow::Array> tuple_data,
|
||||
std::shared_ptr<arrow::Array> dict_data);
|
||||
|
||||
@@ -171,16 +171,20 @@ 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*> sublists, subtuples, subdicts, keytuples, 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, keytuples, dummy));
|
||||
DCHECK(dummy.size() == 0);
|
||||
RETURN_NOT_OK(append(value, result.vals(), sublists, subtuples, subdicts));
|
||||
}
|
||||
}
|
||||
std::shared_ptr<Array> key_val_tuples;
|
||||
if (keytuples.size() > 0) {
|
||||
RETURN_NOT_OK(SerializeSequences(keytuples, &key_val_tuples));
|
||||
}
|
||||
std::shared_ptr<Array> val_list;
|
||||
if (sublists.size() > 0) {
|
||||
RETURN_NOT_OK(SerializeSequences(sublists, &val_list));
|
||||
@@ -193,7 +197,7 @@ Status SerializeDict(std::vector<PyObject*> dicts, std::shared_ptr<Array>* out)
|
||||
if (subdicts.size() > 0) {
|
||||
RETURN_NOT_OK(SerializeDict(subdicts, &val_dict));
|
||||
}
|
||||
*out = result.Finish(val_list, val_tuples, val_dict);
|
||||
*out = result.Finish(key_val_tuples, val_list, val_tuples, val_dict);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
|
||||
@@ -3,8 +3,9 @@ import libnumbuf
|
||||
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"},
|
||||
|
||||
Reference in New Issue
Block a user