diff --git a/VERSION_INFO b/VERSION_INFO index 9e3db2aa12..c72cf6b05f 100644 --- a/VERSION_INFO +++ b/VERSION_INFO @@ -1 +1 @@ -0.1.106 +0.1.107 diff --git a/awkward1/_numpy.py b/awkward1/_numpy.py index 6a2da6e4eb..3a59c6b6d9 100644 --- a/awkward1/_numpy.py +++ b/awkward1/_numpy.py @@ -40,230 +40,19 @@ def array_ufunc(ufunc, method, inputs, kwargs, classes, functions): if method != "__call__" or len(inputs) == 0 or "out" in kwargs: return NotImplemented - def unwrap(x): - if isinstance(x, (awkward1.highlevel.Array, awkward1.highlevel.Record)): - return x.layout - elif isinstance(x, awkward1.highlevel.FillableArray): - return x.snapshot().layout - elif isinstance(x, awkward1.layout.FillableArray): - return x.snapshot() - elif isinstance(x, (awkward1.layout.Content, awkward1.layout.Record)): - return x - elif isinstance(x, numpy.ndarray): - if issubclass(x.dtype.type, numpy.number): - return awkward1.highlevel.Array(x).layout - else: - raise ValueError("numpy.ndarray with {0} cannot be used in {1}".format(repr(x.dtype), ufunc)) - elif isinstance(x, Iterable): - return unwrap(numpy.array(x)) - else: - return x - - def checklength(inputs): - length = len(inputs[0]) - for x in inputs[1:]: - if len(x) != length: - raise ValueError("cannot broadcast {0} of length {1} with {2} of length {3}".format(type(inputs[0]).__name__, length, type(x).__name__, len(x))) - - def apply(inputs): - # handle implicit right-broadcasting (i.e. NumPy-like) - if any(isinstance(x, awkward1._util.listtypes) for x in inputs): - maxdepth = max(x.purelist_depth for x in inputs if isinstance(x, awkward1.layout.Content)) - if maxdepth > 0 and all(x.purelist_isregular for x in inputs if isinstance(x, awkward1.layout.Content)): - nextinputs = [] - for x in inputs: - if isinstance(x, awkward1.layout.Content): - while x.purelist_depth < maxdepth: - x = awkward1.layout.RegularArray(x, 1) - nextinputs.append(x) - if any(x is not y for x, y in zip(inputs, nextinputs)): - return apply(nextinputs) - - # now all lengths must agree - checklength([x for x in inputs if isinstance(x, awkward1.layout.Content)]) + inputs = [awkward1.operations.convert.tolayout(x, allowrecord=True, allowother=True) for x in inputs] + def getfunction(inputs): signature = (ufunc,) + tuple(x.parameters.get("__class__") if isinstance(x, awkward1.layout.Content) else type(x) for x in inputs) custom = awkward1._util.regular_functions(functions).get(signature) - - # the rest of this is one switch statement if custom is not None: - return custom(*inputs, **kwargs) - - elif any(isinstance(x, awkward1._util.unknowntypes) for x in inputs): - return apply([x if not isinstance(x, awkward1._util.unknowntypes) else awkward1.layout.NumpyArray(numpy.array([], dtype=numpy.int64)) for x in inputs]) - - elif any(isinstance(x, awkward1.layout.NumpyArray) and x.ndim > 1 for x in inputs): - return apply([x if not (isinstance(x, awkward1.layout.NumpyArray) and x.ndim > 1) else x.regularize_shape() for x in inputs]) - - elif any(isinstance(x, awkward1._util.indexedtypes) for x in inputs): - return apply([x if not isinstance(x, awkward1._util.indexedtypes) else x.project() for x in inputs]) - - elif any(isinstance(x, awkward1._util.uniontypes) for x in inputs): - tagslist = [] - length = None - for x in inputs: - if isinstance(x, awkward1._util.uniontypes): - tagslist.append(numpy.asarray(x.tags)) - if length is None: - length = len(tagslist[-1]) - elif length != len(tagslist[-1]): - raise ValueError("cannot broadcast UnionArray of length {0} with UnionArray of length {1}".format(length, len(tagslist[-1]))) - - combos = numpy.stack(tagslist, axis=-1) - combos = combos.view([(str(i), combos.dtype) for i in range(len(tagslist))]).reshape(length) - - tags = numpy.empty(length, dtype=numpy.int8) - index = numpy.empty(length, dtype=numpy.int64) - contents = [] - for tag, combo in enumerate(numpy.unique(combos)): - mask = (combos == combo) - tags[mask] = tag - index[mask] = numpy.arange(numpy.count_nonzero(mask)) - nextinputs = [] - for i, x in enumerate(inputs): - if isinstance(x, awkward1._util.uniontypes): - nextinputs.append(x[mask].project(combo[str(i)])) - elif isinstance(x, awkward1.layout.Content): - nextinputs.append(x[mask]) - else: - nextinputs.append(x) - contents.append(apply(nextinputs)) - - tags = awkward1.layout.Index8(tags) - index = awkward1.layout.Index64(index) - return awkward1.layout.UnionArray8_64(tags, index, contents) - - elif any(isinstance(x, awkward1._util.optiontypes) for x in inputs): - mask = None - for x in inputs: - if isinstance(x, (awkward1.layout.IndexedOptionArray32, awkward1.layout.IndexedOptionArray64)): - m = numpy.asarray(x.index) < 0 - if mask is None: - mask = m - else: - numpy.bitwise_or(mask, m, out=mask) - - nextmask = awkward1.layout.Index8(mask.view(numpy.int8)) - index = numpy.full(len(mask), -1, dtype=numpy.int64) - index[~mask] = numpy.arange(len(mask) - numpy.count_nonzero(mask), dtype=numpy.int64) - index = awkward1.layout.Index64(index) - if any(not isinstance(x, awkward1._util.optiontypes) for x in inputs): - nextindex = numpy.arange(len(mask), dtype=numpy.int64) - nextindex[mask] = -1 - nextindex = awkward1.layout.Index64(nextindex) - - nextinputs = [] - for x in inputs: - if isinstance(x, awkward1._util.optiontypes): - nextinputs.append(x.project(nextmask)) - else: - nextinputs.append(awkward1.layout.IndexedOptionArray64(nextindex, x).project(nextmask)) - - return awkward1.layout.IndexedOptionArray64(index, apply(nextinputs)) - - elif any(isinstance(x, awkward1._util.listtypes) for x in inputs): - if all(isinstance(x, awkward1.layout.RegularArray) or not isinstance(x, awkward1._util.listtypes) for x in inputs): - maxsize = max([x.size for x in inputs if isinstance(x, awkward1.layout.RegularArray)]) - for x in inputs: - if isinstance(x, awkward1.layout.RegularArray): - if maxsize > 1 and x.size == 1: - tmpindex = awkward1.layout.Index64(numpy.repeat(numpy.arange(len(x), dtype=numpy.int64), maxsize)) - nextinputs = [] - for x in inputs: - if isinstance(x, awkward1.layout.RegularArray): - if maxsize > 1 and x.size == 1: - nextinputs.append(awkward1.layout.IndexedArray64(tmpindex, x.content).project()) - elif x.size == maxsize: - nextinputs.append(x.content) - else: - raise ValueError("cannot broadcast RegularArray of size {0} with RegularArray of size {1}".format(x.size, maxsize)) - else: - nextinputs.append(x) - return awkward1.layout.RegularArray(apply(nextinputs), maxsize) - - else: - for x in inputs: - if isinstance(x, awkward1._util.listtypes) and not isinstance(x, awkward1.layout.RegularArray): - first = x - break - offsets = first.compact_offsets64() - nextinputs = [] - for x in inputs: - if isinstance(x, awkward1._util.listtypes): - nextinputs.append(x.broadcast_tooffsets64(offsets).content) - # handle implicit left-broadcasting (unlike NumPy) - elif isinstance(x, awkward1.layout.Content): - nextinputs.append(awkward1.layout.RegularArray(x, 1).broadcast_tooffsets64(offsets).content) - else: - nextinputs.append(x) - return awkward1.layout.ListOffsetArray64(offsets, apply(nextinputs)) - - elif any(isinstance(x, awkward1._util.recordtypes) for x in inputs): - keys = None - length = None - istuple = True - for x in inputs: - if isinstance(x, awkward1._util.recordtypes): - if keys is None: - keys = x.keys() - elif set(keys) != set(x.keys()): - raise ValueError("cannot broadcast records because keys don't match:\n {0}\n {1}".format(", ".join(sorted(keys)), ", ".join(sorted(x.keys())))) - if length is None: - length = len(x) - elif length != len(x): - raise ValueError("cannot broadcast RecordArray of length {0} with RecordArray of length {1}".format(length, len(x))) - if not x.istuple: - istuple = False - - if len(keys) == 0: - return awkward1.layout.RecordArray(length, istuple) - else: - contents = [] - for key in keys: - contents.append(apply([x if not isinstance(x, awkward1._util.recordtypes) else x[key] for x in inputs])) - return awkward1.layout.RecordArray(contents, keys) - - else: - assert all(isinstance(x, awkward1.layout.NumpyArray) or not isinstance(x, awkward1.layout.Content) for x in inputs) - result = getattr(ufunc, method)(*inputs, **kwargs) - return awkward1.layout.NumpyArray(result) - - isscalar = [] - - def pack(inputs): - maxlen = -1 - for x in inputs: - if isinstance(x, awkward1.layout.Content): - maxlen = max(maxlen, len(x)) - if maxlen < 0: - maxlen = 1 - nextinputs = [] - for x in inputs: - if isinstance(x, awkward1.layout.Record): - index = numpy.full(maxlen, x.at, dtype=numpy.int64) - nextinputs.append(awkward1.layout.RegularArray(x.array[index], maxlen)) - isscalar.append(True) - elif isinstance(x, awkward1.layout.Content): - nextinputs.append(awkward1.layout.RegularArray(x, len(x))) - isscalar.append(False) - else: - nextinputs.append(x) - isscalar.append(True) - return nextinputs - - def unpack(x): - if all(isscalar): - if len(x) == 0: - return x.getitem_nothing().getitem_nothing() - else: - return x[0][0] + return lambda: custom(*inputs, **kwargs) + elif all(isinstance(x, awkward1.layout.NumpyArray) or not isinstance(x, awkward1.layout.Content) for x in inputs): + return lambda: awkward1.layout.NumpyArray(getattr(ufunc, method)(*inputs, **kwargs)) else: - if len(x) == 0: - return x.getitem_nothing() - else: - return x[0] + return None - return awkward1._util.wrap(unpack(apply(pack([unwrap(x) for x in inputs]))), classes, functions) + return awkward1._util.wrap(awkward1._util.broadcast_and_apply(inputs, getfunction), classes, functions) try: NDArrayOperatorsMixin = numpy.lib.mixins.NDArrayOperatorsMixin diff --git a/awkward1/_util.py b/awkward1/_util.py index 2b722d1f77..f727f4f0d5 100644 --- a/awkward1/_util.py +++ b/awkward1/_util.py @@ -34,9 +34,9 @@ def regular_functions(functions): else: return functions -def combine_classes(arrays): +def combine_classes(*arrays): classes = None - for x in arrays: + for x in arrays[::-1]: if isinstance(x, (awkward1.highlevel.Array, awkward1.highlevel.Record, awkward1.highlevel.FillableArray)) and x._classes is not None: if classes is None: classes = dict(x._classes) @@ -44,9 +44,9 @@ def combine_classes(arrays): classes.update(x._classes) return classes -def combine_functions(arrays): +def combine_functions(*arrays): functions = None - for x in arrays: + for x in arrays[::-1]: if isinstance(x, (awkward1.highlevel.Array, awkward1.highlevel.Record, awkward1.highlevel.FillableArray)) and x._functions is not None: if functions is None: functions = dict(x._functions) @@ -114,6 +114,210 @@ def key2index(keys, key): key2index._pattern = re.compile(r"^[1-9][0-9]*$") +def broadcast_and_apply(inputs, getfunction): + def checklength(inputs): + length = len(inputs[0]) + for x in inputs[1:]: + if len(x) != length: + raise ValueError("cannot broadcast {0} of length {1} with {2} of length {3}".format(type(inputs[0]).__name__, length, type(x).__name__, len(x))) + + def apply(inputs): + # handle implicit right-broadcasting (i.e. NumPy-like) + if any(isinstance(x, listtypes) for x in inputs): + maxdepth = max(x.purelist_depth for x in inputs if isinstance(x, awkward1.layout.Content)) + if maxdepth > 0 and all(x.purelist_isregular for x in inputs if isinstance(x, awkward1.layout.Content)): + nextinputs = [] + for x in inputs: + if isinstance(x, awkward1.layout.Content): + while x.purelist_depth < maxdepth: + x = awkward1.layout.RegularArray(x, 1) + nextinputs.append(x) + if any(x is not y for x, y in zip(inputs, nextinputs)): + return apply(nextinputs) + + # now all lengths must agree + checklength([x for x in inputs if isinstance(x, awkward1.layout.Content)]) + + function = getfunction(inputs) + + # the rest of this is one switch statement + if function is not None: + return function() + + elif any(isinstance(x, unknowntypes) for x in inputs): + return apply([x if not isinstance(x, unknowntypes) else awkward1.layout.NumpyArray(numpy.array([], dtype=numpy.int64)) for x in inputs]) + + elif any(isinstance(x, awkward1.layout.NumpyArray) and x.ndim > 1 for x in inputs): + return apply([x if not (isinstance(x, awkward1.layout.NumpyArray) and x.ndim > 1) else x.toRegularArray() for x in inputs]) + + elif any(isinstance(x, indexedtypes) for x in inputs): + return apply([x if not isinstance(x, indexedtypes) else x.project() for x in inputs]) + + elif any(isinstance(x, uniontypes) for x in inputs): + tagslist = [] + length = None + for x in inputs: + if isinstance(x, uniontypes): + tagslist.append(numpy.asarray(x.tags)) + if length is None: + length = len(tagslist[-1]) + elif length != len(tagslist[-1]): + raise ValueError("cannot broadcast UnionArray of length {0} with UnionArray of length {1}".format(length, len(tagslist[-1]))) + + combos = numpy.stack(tagslist, axis=-1) + combos = combos.view([(str(i), combos.dtype) for i in range(len(tagslist))]).reshape(length) + + tags = numpy.empty(length, dtype=numpy.int8) + index = numpy.empty(length, dtype=numpy.int64) + contents = [] + for tag, combo in enumerate(numpy.unique(combos)): + mask = (combos == combo) + tags[mask] = tag + index[mask] = numpy.arange(numpy.count_nonzero(mask)) + nextinputs = [] + for i, x in enumerate(inputs): + if isinstance(x, uniontypes): + nextinputs.append(x[mask].project(combo[str(i)])) + elif isinstance(x, awkward1.layout.Content): + nextinputs.append(x[mask]) + else: + nextinputs.append(x) + contents.append(apply(nextinputs)) + + tags = awkward1.layout.Index8(tags) + index = awkward1.layout.Index64(index) + return awkward1.layout.UnionArray8_64(tags, index, contents) + + elif any(isinstance(x, optiontypes) for x in inputs): + mask = None + for x in inputs: + if isinstance(x, (awkward1.layout.IndexedOptionArray32, awkward1.layout.IndexedOptionArray64)): + m = numpy.asarray(x.index) < 0 + if mask is None: + mask = m + else: + numpy.bitwise_or(mask, m, out=mask) + + nextmask = awkward1.layout.Index8(mask.view(numpy.int8)) + index = numpy.full(len(mask), -1, dtype=numpy.int64) + index[~mask] = numpy.arange(len(mask) - numpy.count_nonzero(mask), dtype=numpy.int64) + index = awkward1.layout.Index64(index) + if any(not isinstance(x, optiontypes) for x in inputs): + nextindex = numpy.arange(len(mask), dtype=numpy.int64) + nextindex[mask] = -1 + nextindex = awkward1.layout.Index64(nextindex) + + nextinputs = [] + for x in inputs: + if isinstance(x, optiontypes): + nextinputs.append(x.project(nextmask)) + else: + nextinputs.append(awkward1.layout.IndexedOptionArray64(nextindex, x).project(nextmask)) + + return awkward1.layout.IndexedOptionArray64(index, apply(nextinputs)) + + elif any(isinstance(x, listtypes) for x in inputs): + if all(isinstance(x, awkward1.layout.RegularArray) or not isinstance(x, listtypes) for x in inputs): + maxsize = max([x.size for x in inputs if isinstance(x, awkward1.layout.RegularArray)]) + for x in inputs: + if isinstance(x, awkward1.layout.RegularArray): + if maxsize > 1 and x.size == 1: + tmpindex = awkward1.layout.Index64(numpy.repeat(numpy.arange(len(x), dtype=numpy.int64), maxsize)) + nextinputs = [] + for x in inputs: + if isinstance(x, awkward1.layout.RegularArray): + if maxsize > 1 and x.size == 1: + nextinputs.append(awkward1.layout.IndexedArray64(tmpindex, x.content[:len(x)*x.size]).project()) + elif x.size == maxsize: + nextinputs.append(x.content[:len(x)*x.size]) + else: + raise ValueError("cannot broadcast RegularArray of size {0} with RegularArray of size {1}".format(x.size, maxsize)) + else: + nextinputs.append(x) + return awkward1.layout.RegularArray(apply(nextinputs), maxsize) + + else: + for x in inputs: + if isinstance(x, listtypes) and not isinstance(x, awkward1.layout.RegularArray): + first = x + break + offsets = first.compact_offsets64() + nextinputs = [] + for x in inputs: + if isinstance(x, listtypes): + nextinputs.append(x.broadcast_tooffsets64(offsets).content) + # handle implicit left-broadcasting (unlike NumPy) + elif isinstance(x, awkward1.layout.Content): + nextinputs.append(awkward1.layout.RegularArray(x, 1).broadcast_tooffsets64(offsets).content) + else: + nextinputs.append(x) + return awkward1.layout.ListOffsetArray64(offsets, apply(nextinputs)) + + elif any(isinstance(x, recordtypes) for x in inputs): + keys = None + length = None + istuple = True + for x in inputs: + if isinstance(x, recordtypes): + if keys is None: + keys = x.keys() + elif set(keys) != set(x.keys()): + raise ValueError("cannot broadcast records because keys don't match:\n {0}\n {1}".format(", ".join(sorted(keys)), ", ".join(sorted(x.keys())))) + if length is None: + length = len(x) + elif length != len(x): + raise ValueError("cannot broadcast RecordArray of length {0} with RecordArray of length {1}".format(length, len(x))) + if not x.istuple: + istuple = False + + if len(keys) == 0: + return awkward1.layout.RecordArray(length, istuple) + else: + contents = [] + for key in keys: + contents.append(apply([x if not isinstance(x, recordtypes) else x[key] for x in inputs])) + return awkward1.layout.RecordArray(contents, keys) + + else: + raise ValueError("cannot broadcast: {0}".format(", ".join(type(x) for x in inputs))) + + isscalar = [] + + def pack(inputs): + maxlen = -1 + for x in inputs: + if isinstance(x, awkward1.layout.Content): + maxlen = max(maxlen, len(x)) + if maxlen < 0: + maxlen = 1 + nextinputs = [] + for x in inputs: + if isinstance(x, awkward1.layout.Record): + index = numpy.full(maxlen, x.at, dtype=numpy.int64) + nextinputs.append(awkward1.layout.RegularArray(x.array[index], maxlen)) + isscalar.append(True) + elif isinstance(x, awkward1.layout.Content): + nextinputs.append(awkward1.layout.RegularArray(x, len(x))) + isscalar.append(False) + else: + nextinputs.append(x) + isscalar.append(True) + return nextinputs + + def unpack(x): + if all(isscalar): + if len(x) == 0: + return x.getitem_nothing().getitem_nothing() + else: + return x[0][0] + else: + if len(x) == 0: + return x.getitem_nothing() + else: + return x[0] + + return unpack(apply(pack(inputs))) + def minimally_touching_string(limit_length, layout, classes, functions): import awkward1.layout diff --git a/awkward1/highlevel.py b/awkward1/highlevel.py index ef02c5d8ba..fb35f4a1d8 100644 --- a/awkward1/highlevel.py +++ b/awkward1/highlevel.py @@ -86,6 +86,11 @@ def __getattr__(self, where): else: raise AttributeError("no field named {0}".format(repr(where))) + def __setitem__(self, where, what): + if not isinstance(where, str): + raise ValueError("only fields may be assigned in-place (by field name)") + self._layout = awkward1.operations.structure.withfield(self._layout, what, where).layout + def __dir__(self): return sorted(set(dir(super(Array, self)) + [x for x in self._layout.keys() if _dir_pattern.match(x) and not keyword.iskeyword(x)])) @@ -190,6 +195,11 @@ def type(self): def __getitem__(self, where): return awkward1._util.wrap(self._layout[where], self._classes, self._functions) + def __setitem__(self, where, what): + if not isinstance(where, str): + raise ValueError("only fields may be assigned in-place (by field name)") + self._layout = awkward1.operations.structure.withfield(self._layout, what, where).layout + def __str__(self, limit_value=85): return awkward1._util.minimally_touching_string(limit_value + 2, self._layout, self._classes, self._functions)[1:-1] diff --git a/awkward1/operations/convert.py b/awkward1/operations/convert.py index 18994fbc17..3bfaeddbf8 100644 --- a/awkward1/operations/convert.py +++ b/awkward1/operations/convert.py @@ -230,7 +230,7 @@ def tojson(array, destination=None, pretty=False, maxdecimals=None, buffersize=6 else: return out.tojson(destination, pretty=pretty, maxdecimals=maxdecimals, buffersize=buffersize) -def tolayout(array, allowrecord=True): +def tolayout(array, allowrecord=True, allowother=False, numpytype=(numpy.number,)): import awkward1.highlevel if isinstance(array, awkward1.highlevel.Array): @@ -263,6 +263,8 @@ def tolayout(array, allowrecord=True): return out elif isinstance(array, numpy.ndarray): + if not issubclass(array.dtype.type, numpytype): + raise ValueError("NumPy {0} not allowed".format(repr(array.dtype))) out = awkward1.layout.NumpyArray(array.reshape(-1)) for size in array.shape[:0:-1]: out = awkward1.layout.RegularArray(out, size) @@ -271,7 +273,10 @@ def tolayout(array, allowrecord=True): elif isinstance(array, Iterable): return awkward1.highlevel.Array(array).layout - else: + elif not allowother: raise TypeError("{0} cannot be converted into an Awkward Array".format(array)) + else: + return array + __all__ = [x for x in list(globals()) if not x.startswith("_") and x not in ("numbers", "json", "Iterable", "numpy", "awkward1")] diff --git a/awkward1/operations/structure.py b/awkward1/operations/structure.py index bcba6d8ba6..8965488fb3 100644 --- a/awkward1/operations/structure.py +++ b/awkward1/operations/structure.py @@ -14,6 +14,22 @@ import awkward1._numpy import awkward1.operations.convert +def withfield(base, what, where=None): + base = awkward1.operations.convert.tolayout(base, allowrecord=True, allowother=False) + what = awkward1.operations.convert.tolayout(what, allowrecord=True, allowother=True) + + def getfunction(inputs): + base, what = inputs + if isinstance(base, awkward1.layout.RecordArray): + if not isinstance(what, awkward1.layout.Content): + what = awkward1.layout.NumpyArray(numpy.lib.stride_tricks.as_strided([what], shape=(len(base),), strides=(0,))) + return lambda: base.setitem_field(where, what) + else: + return None + + out = awkward1._util.broadcast_and_apply([base, what], getfunction) + return awkward1._util.wrap(out, classes=awkward1._util.combine_classes(base, what), functions=awkward1._util.combine_functions(base, what)) + def isna(array): import awkward1.highlevel @@ -39,7 +55,8 @@ def apply(layout): else: return numpy.zeros(len(layout), dtype=numpy.bool_) - return awkward1.highlevel.Array(apply(awkward1.operations.convert.tolayout(array, allowrecord=False))) + out = apply(awkward1.operations.convert.tolayout(array, allowrecord=False)) + return awkward1._util.wrap(out, classes=awkward1._util.combine_classes(array), functions=awkward1._util.combine_functions(array)) def notna(array): return ~isna(array) @@ -137,7 +154,7 @@ def concatenate(arrays, axis=0, mergebool=True): if isinstance(out, awkward1._util.uniontypes): out = out.simplify(mergebool=mergebool) - return awkward1._util.wrap(out, classes=awkward1._util.combine_classes(arrays), functions=awkward1._util.combine_functions(arrays)) + return awkward1._util.wrap(out, classes=awkward1._util.combine_classes(*arrays), functions=awkward1._util.combine_functions(*arrays)) @awkward1._numpy.implements(numpy.where) def where(condition, *args, **kwargs): @@ -171,7 +188,7 @@ def where(condition, *args, **kwargs): tmp = awkward1.layout.UnionArray8_64(tags, index, [x, y]) out = tmp.simplify(mergebool=mergebool) - return awkward1._util.wrap(out, classes=awkward1._util.combine_classes((condition,) + args), functions=awkward1._util.combine_functions((condition,) + args)) + return awkward1._util.wrap(out, classes=awkward1._util.combine_classes(*((condition,) + args)), functions=awkward1._util.combine_functions(*((condition,) + args))) else: raise TypeError("where() takes from 1 to 3 positional arguments but {0} were given".format(len(args) + 1)) diff --git a/include/awkward/array/NumpyArray.h b/include/awkward/array/NumpyArray.h index 847d666493..75a9734ef3 100644 --- a/include/awkward/array/NumpyArray.h +++ b/include/awkward/array/NumpyArray.h @@ -40,8 +40,7 @@ namespace awkward { void* byteptr(ssize_t at) const; ssize_t bytelength() const; uint8_t getbyte(ssize_t at) const; - - const std::shared_ptr regularize_shape() const; + const std::shared_ptr toRegularArray() const; bool isscalar() const override; const std::string classname() const override; diff --git a/include/awkward/array/RecordArray.h b/include/awkward/array/RecordArray.h index 3b3bfd6081..0e2470b6a4 100644 --- a/include/awkward/array/RecordArray.h +++ b/include/awkward/array/RecordArray.h @@ -17,10 +17,14 @@ namespace awkward { RecordArray(const std::shared_ptr& identities, const util::Parameters& parameters, const std::vector>& contents, const std::shared_ptr& recordlookup); RecordArray(const std::shared_ptr& identities, const util::Parameters& parameters, const std::vector>& contents); RecordArray(const std::shared_ptr& identities, const util::Parameters& parameters, int64_t length, bool istuple); + RecordArray(const std::shared_ptr& content, const std::string& key); + RecordArray(const std::shared_ptr& content); const std::vector> contents() const; const std::shared_ptr recordlookup() const; bool istuple() const; + const std::shared_ptr setitem_field(int64_t where, const std::shared_ptr& what) const; + const std::shared_ptr setitem_field(const std::string& where, const std::shared_ptr& what) const; const std::string classname() const override; void setidentities() override; @@ -65,9 +69,6 @@ namespace awkward { const std::vector>> fielditems() const; const std::shared_ptr astuple() const; - void append(const std::shared_ptr& content, const std::string& key); - void append(const std::shared_ptr& content); - protected: const std::shared_ptr getitem_next(const SliceAt& at, const Slice& tail, const Index64& advanced) const override; const std::shared_ptr getitem_next(const SliceRange& range, const Slice& tail, const Index64& advanced) const override; @@ -76,8 +77,8 @@ namespace awkward { const std::shared_ptr getitem_next(const SliceFields& fields, const Slice& tail, const Index64& advanced) const override; private: - std::vector> contents_; - std::shared_ptr recordlookup_; + const std::vector> contents_; + const std::shared_ptr recordlookup_; int64_t length_; }; } diff --git a/include/awkward/type/RecordType.h b/include/awkward/type/RecordType.h index 00eb26ac8f..e24a0f889d 100644 --- a/include/awkward/type/RecordType.h +++ b/include/awkward/type/RecordType.h @@ -37,12 +37,9 @@ namespace awkward { const std::vector>> fielditems() const; const std::shared_ptr astuple() const; - void append(const std::shared_ptr& type, const std::string& key); - void append(const std::shared_ptr& type); - private: - std::vector> types_; - std::shared_ptr recordlookup_; + const std::vector> types_; + const std::shared_ptr recordlookup_; }; } diff --git a/src/libawkward/array/EmptyArray.cpp b/src/libawkward/array/EmptyArray.cpp index 4c25b4eff5..1cf44b6862 100644 --- a/src/libawkward/array/EmptyArray.cpp +++ b/src/libawkward/array/EmptyArray.cpp @@ -7,6 +7,7 @@ #include "awkward/type/UnknownType.h" #include "awkward/type/ArrayType.h" #include "awkward/array/NumpyArray.h" +#include "awkward/array/RegularArray.h" #include "awkward/array/EmptyArray.h" diff --git a/src/libawkward/array/NumpyArray.cpp b/src/libawkward/array/NumpyArray.cpp index ce964a4a42..7b3171bc76 100644 --- a/src/libawkward/array/NumpyArray.cpp +++ b/src/libawkward/array/NumpyArray.cpp @@ -143,7 +143,7 @@ namespace awkward { return *reinterpret_cast(reinterpret_cast(ptr_.get()) + byteoffset_ + at); } - const std::shared_ptr NumpyArray::regularize_shape() const { + const std::shared_ptr NumpyArray::toRegularArray() const { if (isscalar()) { return shallow_copy(); } diff --git a/src/libawkward/array/Record.cpp b/src/libawkward/array/Record.cpp index 45c10ebc43..1389a823e6 100644 --- a/src/libawkward/array/Record.cpp +++ b/src/libawkward/array/Record.cpp @@ -158,18 +158,8 @@ namespace awkward { } const std::shared_ptr Record::getitem_fields(const std::vector& keys) const { - RecordArray out(array_.get()->identities(), parameters_, length(), istuple()); - if (istuple()) { - for (auto key : keys) { - out.append(array_.get()->field(key)); - } - } - else { - for (auto key : keys) { - out.append(array_.get()->field(key), key); - } - } - return out.getitem_at_nowrap(at_); + std::shared_ptr recordarray = array_.get()->getitem_fields(keys); + return recordarray.get()->getitem_at_nowrap(at_); } const std::shared_ptr Record::carry(const Index64& carry) const { diff --git a/src/libawkward/array/RecordArray.cpp b/src/libawkward/array/RecordArray.cpp index bb4fe77a33..be28e6615a 100644 --- a/src/libawkward/array/RecordArray.cpp +++ b/src/libawkward/array/RecordArray.cpp @@ -45,6 +45,18 @@ namespace awkward { , recordlookup_(istuple ? nullptr : new util::RecordLookup) , length_(length) { } + RecordArray::RecordArray(const std::shared_ptr& content, const std::string& key) + : Content(Identities::none(), util::Parameters()) + , contents_({ content }) + , recordlookup_(new util::RecordLookup({ key })) + , length_(0) { } + + RecordArray::RecordArray(const std::shared_ptr& content) + : Content(Identities::none(), util::Parameters()) + , contents_({ content }) + , recordlookup_(nullptr) + , length_(0) { } + const std::vector> RecordArray::contents() const { return contents_; } @@ -57,6 +69,58 @@ namespace awkward { return recordlookup_.get() == nullptr; } + const std::shared_ptr RecordArray::setitem_field(int64_t where, const std::shared_ptr& what) const { + if (where < 0) { + throw std::invalid_argument("where must be non-negative"); + } + if (what.get()->length() != length()) { + throw std::invalid_argument(std::string("array of length ") + std::to_string(what.get()->length()) + std::string(" cannot be assigned to record array of length ") + std::to_string(length())); + } + std::vector> contents; + for (size_t i = 0; i < contents_.size(); i++) { + if (where == (int64_t)i) { + contents.push_back(what); + } + contents.push_back(contents_[i]); + } + if (where >= numfields()) { + contents.push_back(what); + } + std::shared_ptr recordlookup(nullptr); + if (recordlookup_.get() != nullptr) { + recordlookup = std::make_shared(); + for (size_t i = 0; i < contents_.size(); i++) { + if (where == (int64_t)i) { + recordlookup.get()->push_back(std::to_string(where)); + } + recordlookup.get()->push_back(recordlookup_.get()->at(i)); + } + if (where >= numfields()) { + recordlookup.get()->push_back(std::to_string(where)); + } + } + return std::make_shared(identities_, parameters_, contents, recordlookup); + } + + const std::shared_ptr RecordArray::setitem_field(const std::string& where, const std::shared_ptr& what) const { + if (what.get()->length() != length()) { + throw std::invalid_argument(std::string("array of length ") + std::to_string(what.get()->length()) + std::string(" cannot be assigned to record array of length ") + std::to_string(length())); + } + std::vector> contents(contents_.begin(), contents_.end()); + contents.push_back(what); + std::shared_ptr recordlookup; + if (recordlookup_.get() != nullptr) { + recordlookup = std::make_shared(); + recordlookup.get()->insert(recordlookup.get()->end(), recordlookup_.get()->begin(), recordlookup_.get()->end()); + recordlookup.get()->push_back(where); + } + else { + recordlookup = util::init_recordlookup(numfields()); + recordlookup.get()->push_back(where); + } + return std::make_shared(identities_, parameters_, contents, recordlookup); + } + const std::string RecordArray::classname() const { return "RecordArray"; } @@ -310,18 +374,18 @@ namespace awkward { } const std::shared_ptr RecordArray::getitem_fields(const std::vector& keys) const { - RecordArray out(identities_, parameters_, length(), istuple()); - if (istuple()) { - for (auto key : keys) { - out.append(field(key).get()->getitem_range_nowrap(0, length())); - } - } - else { - for (auto key : keys) { - out.append(field(key).get()->getitem_range_nowrap(0, length()), key); + std::vector> contents; + std::shared_ptr recordlookup(nullptr); + if (recordlookup_.get() != nullptr) { + recordlookup = std::make_shared(); + } + for (auto key : keys) { + contents.push_back(field(key).get()->getitem_range_nowrap(0, length())); + if (recordlookup.get() != nullptr) { + recordlookup.get()->push_back(key); } } - return out.shallow_copy(); + return std::make_shared(identities_, parameters_, contents, recordlookup); } const std::shared_ptr RecordArray::carry(const Index64& carry) const { @@ -598,23 +662,6 @@ namespace awkward { return std::make_shared(identities_, parameters_, contents_); } - void RecordArray::append(const std::shared_ptr& content, const std::string& key) { - if (recordlookup_.get() == nullptr) { - recordlookup_ = util::init_recordlookup(numfields()); - } - contents_.push_back(content); - recordlookup_.get()->push_back(key); - } - - void RecordArray::append(const std::shared_ptr& content) { - if (recordlookup_.get() == nullptr) { - contents_.push_back(content); - } - else { - append(content, std::to_string(numfields())); - } - } - const std::shared_ptr RecordArray::getitem_next(const std::shared_ptr& head, const Slice& tail, const Index64& advanced) const { std::shared_ptr nexthead = tail.head(); Slice nexttail = tail.tail(); diff --git a/src/libawkward/type/RecordType.cpp b/src/libawkward/type/RecordType.cpp index ae844c46d4..93b009918b 100644 --- a/src/libawkward/type/RecordType.cpp +++ b/src/libawkward/type/RecordType.cpp @@ -205,21 +205,4 @@ namespace awkward { return std::make_shared(parameters_, types_, std::shared_ptr(nullptr)); } - void RecordType::append(const std::shared_ptr& type, const std::string& key) { - if (recordlookup_.get() == nullptr) { - recordlookup_ = util::init_recordlookup(numfields()); - } - types_.push_back(type); - recordlookup_.get()->push_back(key); - } - - void RecordType::append(const std::shared_ptr& type) { - if (recordlookup_.get() == nullptr) { - types_.push_back(type); - } - else { - append(type, std::to_string(numfields())); - } - } - } diff --git a/src/pyawkward.cpp b/src/pyawkward.cpp index e7e37de633..11e9a04332 100644 --- a/src/pyawkward.cpp +++ b/src/pyawkward.cpp @@ -1124,14 +1124,6 @@ py::class_, ak::Type> make_Recor } return out; }) - .def("append", [](ak::RecordType& self, py::object type, py::object key) -> void { - if (key.is(py::none())) { - self.append(unbox_type(type)); - } - else { - self.append(unbox_type(type), key.cast()); - } - }, py::arg("type"), py::arg("key") = py::none()) .def(py::pickle([](const ak::RecordType& self) { py::tuple pytypes((size_t)self.numfields()); for (int64_t i = 0; i < self.numfields(); i++) { @@ -1296,11 +1288,11 @@ py::class_, ak::Content> make_Nu .def_property_readonly("ndim", &ak::NumpyArray::ndim) .def_property_readonly("isscalar", &ak::NumpyArray::isscalar) .def_property_readonly("isempty", &ak::NumpyArray::isempty) + .def("toRegularArray", &ak::NumpyArray::toRegularArray) .def_property_readonly("iscontiguous", &ak::NumpyArray::iscontiguous) .def("contiguous", &ak::NumpyArray::contiguous) .def("become_contiguous", &ak::NumpyArray::become_contiguous) - .def("regularize_shape", &ak::NumpyArray::regularize_shape) ); } @@ -1413,6 +1405,28 @@ py::class_, ak::Content> make_ .def_property_readonly("istuple", &ak::RecordArray::istuple) .def_property_readonly("contents", &ak::RecordArray::contents) + .def("setitem_field", [](ak::RecordArray& self, py::object where, py::object what) -> py::object { + std::shared_ptr mywhat = unbox_content(what); + if (where.is(py::none())) { + return box(self.setitem_field(self.numfields(), mywhat)); + } + else { + try { + std::string mywhere = where.cast(); + return box(self.setitem_field(mywhere, mywhat)); + } + catch (py::cast_error err) { + try { + int64_t mywhere = where.cast(); + return box(self.setitem_field(mywhere, mywhat)); + } + catch (py::cast_error err) { + throw std::invalid_argument("where must be None, int, or str"); + } + } + } + }, py::arg("where"), py::arg("what")) + .def("field", [](ak::RecordArray& self, int64_t fieldindex) -> std::shared_ptr { return self.field(fieldindex); }) @@ -1442,15 +1456,6 @@ py::class_, ak::Content> make_ return box(self.astuple()); }) - .def("append", [](ak::RecordArray& self, py::object content, py::object key) -> void { - if (key.is(py::none())) { - self.append(unbox_content(content)); - } - else { - self.append(unbox_content(content), key.cast()); - } - }, py::arg("content"), py::arg("key") = py::none()) - ); } diff --git a/tests/test_PR025_record_array.py b/tests/test_PR025_record_array.py index f7642fa1cd..e912c700f7 100644 --- a/tests/test_PR025_record_array.py +++ b/tests/test_PR025_record_array.py @@ -13,11 +13,7 @@ def test_basic(): content2 = awkward1.layout.NumpyArray(numpy.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) offsets = awkward1.layout.Index64(numpy.array([0, 3, 3, 5, 6, 9])) listoffsetarray = awkward1.layout.ListOffsetArray64(offsets, content2) - recordarray = awkward1.layout.RecordArray(0) - recordarray.append(content1, "one") - recordarray.append(listoffsetarray, "two") - recordarray.append(content2) - recordarray.append(content1, "wonky") + recordarray = awkward1.layout.RecordArray([content1, listoffsetarray, content2, content1], keys=["one", "two", "2", "wonky"]) assert awkward1.tolist(recordarray.field(0)) == [1, 2, 3, 4, 5] assert awkward1.tolist(recordarray.field("two")) == [[1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]] assert awkward1.tolist(recordarray.field("wonky")) == [1, 2, 3, 4, 5] @@ -62,9 +58,7 @@ def test_scalar_record(): content2 = awkward1.layout.NumpyArray(numpy.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) offsets = awkward1.layout.Index64(numpy.array([0, 3, 3, 5, 6, 9])) listoffsetarray = awkward1.layout.ListOffsetArray64(offsets, content2) - recordarray = awkward1.layout.RecordArray(0) - recordarray.append(content1, "one") - recordarray.append(listoffsetarray, "two") + recordarray = awkward1.layout.RecordArray([content1, listoffsetarray], keys=["one", "two"]) str(recordarray) str(recordarray[2]) @@ -86,9 +80,7 @@ def test_type(): content2 = awkward1.layout.NumpyArray(numpy.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], dtype=numpy.float64)) offsets = awkward1.layout.Index64(numpy.array([0, 3, 3, 5, 6, 9])) listoffsetarray = awkward1.layout.ListOffsetArray64(offsets, content2) - recordarray = awkward1.layout.RecordArray(0, True) - recordarray.append(content1) - recordarray.append(listoffsetarray) + recordarray = awkward1.layout.RecordArray([content1, listoffsetarray]) assert str(awkward1.typeof(recordarray)) == '(int64, var * float64)' assert awkward1.typeof(recordarray) == awkward1.layout.RecordType(( @@ -98,9 +90,7 @@ def test_type(): (awkward1.layout.PrimitiveType("int64"), awkward1.layout.ListType(awkward1.layout.PrimitiveType("float64")))) - recordarray = awkward1.layout.RecordArray(0, True) - recordarray.append(content1, "one") - recordarray.append(listoffsetarray, "two") + recordarray = awkward1.layout.RecordArray([content1, listoffsetarray], keys=["one", "two"]) assert str(awkward1.typeof(recordarray)) in ('{"one": int64, "two": var * float64}', '{"two": var * float64, "one": int64}') assert str(awkward1.layout.RecordType( diff --git a/tests/test_PR086_nep13_ufunc.py b/tests/test_PR086_nep13_ufunc.py index 34c4aadbad..b8416e2412 100644 --- a/tests/test_PR086_nep13_ufunc.py +++ b/tests/test_PR086_nep13_ufunc.py @@ -42,8 +42,8 @@ def test_indexedoptionarray(): def test_regularize_shape(): array = awkward1.layout.NumpyArray(numpy.arange(2*3*5).reshape(2, 3, 5)) - assert isinstance(array.regularize_shape(), awkward1.layout.RegularArray) - assert awkward1.tolist(array.regularize_shape()) == awkward1.tolist(array) + assert isinstance(array.toRegularArray(), awkward1.layout.RegularArray) + assert awkward1.tolist(array.toRegularArray()) == awkward1.tolist(array) def test_regulararray(): array = awkward1.Array(numpy.arange(2*3*5).reshape(2, 3, 5)) diff --git a/tests/test_PR107_assign_fields_to_records.py b/tests/test_PR107_assign_fields_to_records.py new file mode 100644 index 0000000000..90208510e6 --- /dev/null +++ b/tests/test_PR107_assign_fields_to_records.py @@ -0,0 +1,83 @@ +# BSD 3-Clause License; see https://github.com/jpivarski/awkward-1.0/blob/master/LICENSE + +import sys + +import pytest +import numpy + +import awkward1 + +def test_record(): + array1 = awkward1.Array([{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, {"x": 3, "y": 3.3}]).layout + assert awkward1.tolist(array1) == [{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, {"x": 3, "y": 3.3}] + + array2 = array1.setitem_field("z", awkward1.Array([[], [1], [2, 2]]).layout) + assert awkward1.tolist(array2) == [{"x": 1, "y": 1.1, "z": []}, {"x": 2, "y": 2.2, "z": [1]}, {"x": 3, "y": 3.3, "z": [2, 2]}] + + array3 = array1.setitem_field(None, awkward1.Array([[], [1], [2, 2]]).layout) + assert awkward1.tolist(array3) == [{"x": 1, "y": 1.1, "2": []}, {"x": 2, "y": 2.2, "2": [1]}, {"x": 3, "y": 3.3, "2": [2, 2]}] + + array3 = array1.setitem_field(0, awkward1.Array([[], [1], [2, 2]]).layout) + assert awkward1.tolist(array3) == [{"x": 1, "y": 1.1, "0": []}, {"x": 2, "y": 2.2, "0": [1]}, {"x": 3, "y": 3.3, "0": [2, 2]}] + + array1 = awkward1.Array([(1, 1.1), (2, 2.2), (3, 3.3)]).layout + assert awkward1.tolist(array1) == [(1, 1.1), (2, 2.2), (3, 3.3)] + + array2 = array1.setitem_field("z", awkward1.Array([[], [1], [2, 2]]).layout) + assert awkward1.tolist(array2) == [{"0": 1, "1": 1.1, "z": []}, {"0": 2, "1": 2.2, "z": [1]}, {"0": 3, "1": 3.3, "z": [2, 2]}] + + array3 = array1.setitem_field(None, awkward1.Array([[], [1], [2, 2]]).layout) + assert awkward1.tolist(array3) == [(1, 1.1, []), (2, 2.2, [1]), (3, 3.3, [2, 2])] + + array3 = array1.setitem_field(0, awkward1.Array([[], [1], [2, 2]]).layout) + assert awkward1.tolist(array3) == [([], 1, 1.1), ([1], 2, 2.2), ([2, 2], 3, 3.3)] + + array3 = array1.setitem_field(1, awkward1.Array([[], [1], [2, 2]]).layout) + assert awkward1.tolist(array3) == [(1, [], 1.1), (2, [1], 2.2), (3, [2, 2], 3.3)] + + array3 = array1.setitem_field(100, awkward1.Array([[], [1], [2, 2]]).layout) + assert awkward1.tolist(array3) == [(1, 1.1, []), (2, 2.2, [1]), (3, 3.3, [2, 2])] + +def test_withfield(): + base = awkward1.Array([{"x": 1}, {"x": 2}, {"x": 3}]) + what = awkward1.Array([1.1, 2.2, 3.3]) + assert awkward1.tolist(awkward1.withfield(base, what)) == [{"x": 1, "1": 1.1}, {"x": 2, "1": 2.2}, {"x": 3, "1": 3.3}] + assert awkward1.tolist(awkward1.withfield(base, what, where="y")) == [{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, {"x": 3, "y": 3.3}] + + base["z"] = what + assert awkward1.tolist(base) == [{"x": 1, "z": 1.1}, {"x": 2, "z": 2.2}, {"x": 3, "z": 3.3}] + + base["q"] = 123 + assert awkward1.tolist(base) == [{"x": 1, "z": 1.1, "q": 123}, {"x": 2, "z": 2.2, "q": 123}, {"x": 3, "z": 3.3, "q": 123}] + + base = awkward1.Array([{"x": 1}, {"x": 2}, {"x": 3}])[2] + assert awkward1.tolist(awkward1.withfield(base, 100, "y")) == {"x": 3, "y": 100} + +def test_regulararray(): + content = awkward1.layout.NumpyArray(numpy.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) + recordarray = awkward1.layout.RecordArray({"x": content}) + regulararray = awkward1.Array(awkward1.layout.RegularArray(recordarray, 3)) + + content2 = awkward1.layout.NumpyArray(numpy.array([100, 200, 300])) + regulararray2 = awkward1.Array(awkward1.layout.RegularArray(content2, 1)) + assert awkward1.tolist(awkward1.withfield(regulararray, regulararray2, "y")) == [[{"x": 0.0, "y": 100}, {"x": 1.1, "y": 100}, {"x": 2.2, "y": 100}], [{"x": 3.3, "y": 200}, {"x": 4.4, "y": 200}, {"x": 5.5, "y": 200}], [{"x": 6.6, "y": 300}, {"x": 7.7, "y": 300}, {"x": 8.8, "y": 300}]] + + content2 = awkward1.layout.NumpyArray(numpy.array([100, 200, 300, 400, 500, 600, 700, 800, 900])) + regulararray2 = awkward1.Array(awkward1.layout.RegularArray(content2, 3)) + assert awkward1.tolist(awkward1.withfield(regulararray, regulararray2, "y")) == [[{"x": 0.0, "y": 100}, {"x": 1.1, "y": 200}, {"x": 2.2, "y": 300}], [{"x": 3.3, "y": 400}, {"x": 4.4, "y": 500}, {"x": 5.5, "y": 600}], [{"x": 6.6, "y": 700}, {"x": 7.7, "y": 800}, {"x": 8.8, "y": 900}]] + + content2 = awkward1.Array(awkward1.layout.NumpyArray(numpy.array([[100], [200], [300]]))) + assert awkward1.tolist(awkward1.withfield(regulararray, content2, "y")) == [[{"x": 0.0, "y": 100}, {"x": 1.1, "y": 100}, {"x": 2.2, "y": 100}], [{"x": 3.3, "y": 200}, {"x": 4.4, "y": 200}, {"x": 5.5, "y": 200}], [{"x": 6.6, "y": 300}, {"x": 7.7, "y": 300}, {"x": 8.8, "y": 300}]] + + content2 = awkward1.Array(awkward1.layout.NumpyArray(numpy.array([[100, 200, 300], [400, 500, 600], [700, 800, 900]]))) + assert awkward1.tolist(awkward1.withfield(regulararray, content2, "y")) == [[{"x": 0.0, "y": 100}, {"x": 1.1, "y": 200}, {"x": 2.2, "y": 300}], [{"x": 3.3, "y": 400}, {"x": 4.4, "y": 500}, {"x": 5.5, "y": 600}], [{"x": 6.6, "y": 700}, {"x": 7.7, "y": 800}, {"x": 8.8, "y": 900}]] + +def test_listarray(): + one = awkward1.Array([[{"x": 1}, {"x": 2}, {"x": 3}], [], [{"x": 4}, {"x": 5}]]) + two = awkward1.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]]) + assert awkward1.tolist(awkward1.withfield(one, two, "y")) == [[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, {"x": 3, "y": 3.3}], [], [{"x": 4, "y": 4.4}, {"x": 5, "y": 5.5}]] + + three = awkward1.Array([100, 200, 300]) + assert awkward1.tolist(awkward1.withfield(one, three, "y")) == [[{"x": 1, "y": 100}, {"x": 2, "y": 100}, {"x": 3, "y": 100}], [], [{"x": 4, "y": 300}, {"x": 5, "y": 300}]] + + assert awkward1.tolist(awkward1.withfield(one, [100, 200, 300], "y")) == [[{"x": 1, "y": 100}, {"x": 2, "y": 100}, {"x": 3, "y": 100}], [], [{"x": 4, "y": 300}, {"x": 5, "y": 300}]]