Skip to content

Commit 3d7a5a6

Browse files
pitrouxhochy
pitrou
authored and
xhochy
committed
ARROW-2074: [Python] Infer lists of dicts as struct arrays
Also refactor the type inference visitor and remove the superfluous separate SeqVisitor; improve inference visitor performance by 30%; and add a struct type inference benchmark. Author: Antoine Pitrou <antoine@python.org> Closes apache#1935 from pitrou/ARROW-2074-infer-dict-lists and squashes the following commits: 13ed6c3 <Antoine Pitrou> Fix tests on 2.7 3baa2ea <Antoine Pitrou> ARROW-2074: Infer lists of dicts as struct arrays
1 parent 2d278ab commit 3d7a5a6

5 files changed

Lines changed: 213 additions & 127 deletions

File tree

‎cpp/src/arrow/python/builtin_convert.cc‎

Lines changed: 122 additions & 120 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
#include <algorithm>
2323
#include <limits>
24+
#include <map>
2425
#include <sstream>
2526
#include <string>
2627
#include <utility>
@@ -49,9 +50,11 @@ Status InvalidConversion(PyObject* obj, const std::string& expected_types,
4950
return Status::OK();
5051
}
5152

52-
class ScalarVisitor {
53+
class TypeInferrer {
54+
// A type inference visitor for Python values
55+
5356
public:
54-
ScalarVisitor()
57+
TypeInferrer()
5558
: total_count_(0),
5659
none_count_(0),
5760
bool_count_(0),
@@ -62,14 +65,46 @@ class ScalarVisitor {
6265
binary_count_(0),
6366
unicode_count_(0),
6467
decimal_count_(0),
68+
list_count_(0),
69+
struct_count_(0),
6570
max_decimal_metadata_(std::numeric_limits<int32_t>::min(),
6671
std::numeric_limits<int32_t>::min()),
6772
decimal_type_() {
68-
PyAcquireGIL lock;
6973
Status status = internal::ImportDecimalType(&decimal_type_);
7074
DCHECK_OK(status);
7175
}
7276

77+
// Infer value type from a sequence of values
78+
Status VisitSequence(PyObject* obj) {
79+
// Loop through a sequence
80+
if (PyArray_Check(obj)) {
81+
Py_ssize_t size = PySequence_Size(obj);
82+
OwnedRef value_ref;
83+
84+
for (Py_ssize_t i = 0; i < size; ++i) {
85+
auto array = reinterpret_cast<PyArrayObject*>(obj);
86+
auto ptr = reinterpret_cast<const char*>(PyArray_GETPTR1(array, i));
87+
88+
value_ref.reset(PyArray_GETITEM(array, ptr));
89+
RETURN_IF_PYERROR();
90+
RETURN_NOT_OK(Visit(value_ref.obj()));
91+
}
92+
} else if (PySequence_Check(obj)) {
93+
OwnedRef seq_ref(PySequence_Fast(obj, "Object is not a sequence or iterable"));
94+
RETURN_IF_PYERROR();
95+
PyObject* seq = seq_ref.obj();
96+
97+
Py_ssize_t size = PySequence_Fast_GET_SIZE(seq);
98+
for (Py_ssize_t i = 0; i < size; ++i) {
99+
PyObject* value = PySequence_Fast_GET_ITEM(seq, i);
100+
RETURN_NOT_OK(Visit(value));
101+
}
102+
} else {
103+
return Status::TypeError("Object is not a sequence or iterable");
104+
}
105+
return Status::OK();
106+
}
107+
73108
Status Visit(PyObject* obj) {
74109
++total_count_;
75110
if (obj == Py_None || internal::PyFloat_IsNaN(obj)) {
@@ -103,6 +138,10 @@ class ScalarVisitor {
103138
ss << type->ToString();
104139
return Status::Invalid(ss.str());
105140
}
141+
} else if (PyList_Check(obj) || PyArray_Check(obj)) {
142+
return VisitList(obj);
143+
} else if (PyDict_Check(obj)) {
144+
return VisitDict(obj);
106145
} else if (PyObject_IsInstance(obj, decimal_type_.obj())) {
107146
RETURN_NOT_OK(max_decimal_metadata_.Update(obj));
108147
++decimal_count_;
@@ -118,14 +157,36 @@ class ScalarVisitor {
118157
return Status::OK();
119158
}
120159

121-
std::shared_ptr<DataType> GetType() {
160+
Status Validate() const {
161+
if (list_count_ > 0) {
162+
if (list_count_ + none_count_ != total_count_) {
163+
return Status::Invalid("cannot mix list and non-list, non-null values");
164+
}
165+
RETURN_NOT_OK(list_inferrer_->Validate());
166+
} else if (struct_count_ > 0) {
167+
if (struct_count_ + none_count_ != total_count_) {
168+
return Status::Invalid("cannot mix struct and non-struct, non-null values");
169+
}
170+
for (const auto& it : struct_inferrers_) {
171+
RETURN_NOT_OK(it.second.Validate());
172+
}
173+
}
174+
return Status::OK();
175+
}
176+
177+
std::shared_ptr<DataType> GetType() const {
122178
// TODO(wesm): handling mixed-type cases
123-
if (decimal_count_) {
179+
if (list_count_) {
180+
auto value_type = list_inferrer_->GetType();
181+
DCHECK(value_type != nullptr);
182+
return list(value_type);
183+
} else if (struct_count_) {
184+
return GetStructType();
185+
} else if (decimal_count_) {
124186
return decimal(max_decimal_metadata_.precision(), max_decimal_metadata_.scale());
125187
} else if (float_count_) {
126188
return float64();
127189
} else if (int_count_) {
128-
// TODO(wesm): tighter type later
129190
return int64();
130191
} else if (date_count_) {
131192
return date64();
@@ -144,6 +205,53 @@ class ScalarVisitor {
144205

145206
int64_t total_count() const { return total_count_; }
146207

208+
protected:
209+
Status VisitList(PyObject* obj) {
210+
if (!list_inferrer_) {
211+
list_inferrer_.reset(new TypeInferrer);
212+
}
213+
++list_count_;
214+
return list_inferrer_->VisitSequence(obj);
215+
}
216+
217+
Status VisitDict(PyObject* obj) {
218+
PyObject* key_obj;
219+
PyObject* value_obj;
220+
Py_ssize_t pos = 0;
221+
222+
while (PyDict_Next(obj, &pos, &key_obj, &value_obj)) {
223+
std::string key;
224+
if (PyUnicode_Check(key_obj)) {
225+
RETURN_NOT_OK(internal::PyUnicode_AsStdString(key_obj, &key));
226+
} else if (PyBytes_Check(key_obj)) {
227+
key = internal::PyBytes_AsStdString(key_obj);
228+
} else {
229+
std::stringstream ss;
230+
ss << "Expected dict key of type str or bytes, got '" << Py_TYPE(key_obj)->tp_name
231+
<< "'";
232+
return Status::TypeError(ss.str());
233+
}
234+
// Get or create visitor for this key
235+
auto it = struct_inferrers_.find(key);
236+
if (it == struct_inferrers_.end()) {
237+
it = struct_inferrers_.insert(std::make_pair(key, TypeInferrer())).first;
238+
}
239+
TypeInferrer* visitor = &it->second;
240+
RETURN_NOT_OK(visitor->Visit(value_obj));
241+
}
242+
++struct_count_;
243+
return Status::OK();
244+
}
245+
246+
std::shared_ptr<DataType> GetStructType() const {
247+
std::vector<std::shared_ptr<Field>> fields;
248+
for (const auto& it : struct_inferrers_) {
249+
const auto struct_field = field(it.first, it.second.GetType());
250+
fields.emplace_back(struct_field);
251+
}
252+
return struct_(fields);
253+
}
254+
147255
private:
148256
int64_t total_count_;
149257
int64_t none_count_;
@@ -155,6 +263,10 @@ class ScalarVisitor {
155263
int64_t binary_count_;
156264
int64_t unicode_count_;
157265
int64_t decimal_count_;
266+
int64_t list_count_;
267+
std::unique_ptr<TypeInferrer> list_inferrer_;
268+
int64_t struct_count_;
269+
std::map<std::string, TypeInferrer> struct_inferrers_;
158270

159271
internal::DecimalMetadata max_decimal_metadata_;
160272

@@ -163,116 +275,6 @@ class ScalarVisitor {
163275
OwnedRefNoGIL decimal_type_;
164276
};
165277

166-
static constexpr int MAX_NESTING_LEVELS = 32;
167-
168-
// SeqVisitor is used to infer the type.
169-
class SeqVisitor {
170-
public:
171-
SeqVisitor() : max_nesting_level_(0), max_observed_level_(0), nesting_histogram_() {
172-
std::fill(nesting_histogram_, nesting_histogram_ + MAX_NESTING_LEVELS, 0);
173-
}
174-
175-
// co-recursive with VisitElem
176-
Status Visit(PyObject* obj, int level = 0) {
177-
max_nesting_level_ = std::max(max_nesting_level_, level);
178-
179-
// Loop through a sequence
180-
if (!PySequence_Check(obj))
181-
return Status::TypeError("Object is not a sequence or iterable");
182-
183-
Py_ssize_t size = PySequence_Size(obj);
184-
for (int64_t i = 0; i < size; ++i) {
185-
OwnedRef ref;
186-
if (PyArray_Check(obj)) {
187-
auto array = reinterpret_cast<PyArrayObject*>(obj);
188-
auto ptr = reinterpret_cast<const char*>(PyArray_GETPTR1(array, i));
189-
190-
ref.reset(PyArray_GETITEM(array, ptr));
191-
RETURN_IF_PYERROR();
192-
193-
RETURN_NOT_OK(VisitElem(ref, level));
194-
} else {
195-
ref.reset(PySequence_GetItem(obj, i));
196-
RETURN_IF_PYERROR();
197-
RETURN_NOT_OK(VisitElem(ref, level));
198-
}
199-
}
200-
return Status::OK();
201-
}
202-
203-
std::shared_ptr<DataType> GetType() {
204-
// If all the non-list inputs were null (or there were no inputs)
205-
std::shared_ptr<DataType> result;
206-
if (scalars_.total_count() == 0) {
207-
// Lists of Lists of NULL
208-
result = null();
209-
} else {
210-
// Lists of Lists of [X]
211-
result = scalars_.GetType();
212-
}
213-
for (int i = 0; i < max_nesting_level_; ++i) {
214-
result = std::make_shared<ListType>(result);
215-
}
216-
return result;
217-
}
218-
219-
Status Validate() const {
220-
if (scalars_.total_count() > 0) {
221-
if (num_nesting_levels() > 1) {
222-
return Status::Invalid("Mixed nesting levels not supported");
223-
// If the nesting goes deeper than the deepest scalar
224-
} else if (max_observed_level_ < max_nesting_level_) {
225-
return Status::Invalid("Mixed nesting levels not supported");
226-
}
227-
}
228-
return Status::OK();
229-
}
230-
231-
// Returns the number of nesting levels which have scalar elements.
232-
int num_nesting_levels() const {
233-
int result = 0;
234-
for (int i = 0; i < MAX_NESTING_LEVELS; ++i) {
235-
if (nesting_histogram_[i] > 0) {
236-
++result;
237-
}
238-
}
239-
return result;
240-
}
241-
242-
private:
243-
ScalarVisitor scalars_;
244-
245-
// Track observed
246-
// Deapest nesting level (irregardless of scalars)
247-
int max_nesting_level_;
248-
int max_observed_level_;
249-
250-
// Number of scalar elements at each nesting level.
251-
// (TOOD: We really only need to know if a scalar is present, not the count).
252-
int nesting_histogram_[MAX_NESTING_LEVELS];
253-
254-
// Visits a specific element (inner part of the loop).
255-
Status VisitElem(const OwnedRef& item_ref, int level) {
256-
DCHECK_NE(item_ref.obj(), NULLPTR);
257-
if (PyList_Check(item_ref.obj()) || PyArray_Check(item_ref.obj())) {
258-
RETURN_NOT_OK(Visit(item_ref.obj(), level + 1));
259-
} else if (PyDict_Check(item_ref.obj())) {
260-
return Status::NotImplemented("No type inference for dicts");
261-
} else {
262-
// We permit nulls at any level of nesting, but they aren't treated like
263-
// other scalar values as far as the checking for mixed nesting structure
264-
if (item_ref.obj() != Py_None) {
265-
++nesting_histogram_[level];
266-
}
267-
if (level > max_observed_level_) {
268-
max_observed_level_ = level;
269-
}
270-
return scalars_.Visit(item_ref.obj());
271-
}
272-
return Status::OK();
273-
}
274-
};
275-
276278
// Convert *obj* to a sequence if necessary
277279
// Fill *size* to its length. If >= 0 on entry, *size* is an upper size
278280
// bound that may lead to truncation.
@@ -319,11 +321,11 @@ Status ConvertToSequenceAndInferSize(PyObject* obj, PyObject** seq, int64_t* siz
319321
// Non-exhaustive type inference
320322
Status InferArrowType(PyObject* obj, std::shared_ptr<DataType>* out_type) {
321323
PyDateTime_IMPORT;
322-
SeqVisitor seq_visitor;
323-
RETURN_NOT_OK(seq_visitor.Visit(obj));
324-
RETURN_NOT_OK(seq_visitor.Validate());
324+
TypeInferrer inferrer;
325+
RETURN_NOT_OK(inferrer.VisitSequence(obj));
326+
RETURN_NOT_OK(inferrer.Validate());
325327

326-
*out_type = seq_visitor.GetType();
328+
*out_type = inferrer.GetType();
327329
if (*out_type == nullptr) {
328330
return Status::TypeError("Unable to determine data type");
329331
}

‎python/benchmarks/convert_builtins.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,7 @@ class InferPyListToArray(object):
5151
"""
5252
size = 10 ** 5
5353
types = ('int64', 'float64', 'bool', 'decimal', 'binary', 'ascii',
54-
'unicode', 'int64 list')
55-
# TODO add 'struct' when supported
54+
'unicode', 'int64 list', 'struct')
5655

5756
param_names = ['type']
5857
params = [types]

0 commit comments

Comments
 (0)