diff --git a/awkward/array/chunked.py b/awkward/array/chunked.py index 2ce6ffa3..fa8cc2d0 100644 --- a/awkward/array/chunked.py +++ b/awkward/array/chunked.py @@ -323,7 +323,7 @@ def __getitem__(self, where): where = (where,) head, tail = where[0], where[1:] - if isinstance(head, self.ChunkedArray): + if isinstance(head, ChunkedArray): if not self._aligned(head): raise ValueError("A ChunkedArray can only be used as a slice of a ChunkedArray if they have the same chunk sizes") chunks = [] @@ -630,14 +630,32 @@ def argpairs(self, nested=False): return out def cross(self, other, nested=False): - out = self.copy(chunks=[x.cross(other, nested=nested) for x in self._chunks]) - out.knowchunksizes() - return out + if not isinstance(other, ChunkedArray) or not self._aligned(other): + raise ValueError("A ChunkedArray can only be crossed with a ChunkedArray if they have the same chunk sizes") + chunks = [] + chunksizes = [] + for c, h in zip(self.chunks, other.chunks): + if isinstance(c, awkward.array.virtual.VirtualArray): + c = c.array + if isinstance(h, awkward.array.virtual.VirtualArray): + h = h.array + chunks.append(c.cross(h, nested=nested)) + chunksizes.append(len(chunks[-1])) + return self.copy(chunks=chunks, chunksizes=chunksizes) def argcross(self, other, nested=False): - out = self.copy(chunks=[x.argcross(other, nested=nested) for x in self._chunks]) - out.knowchunksizes() - return out + if not isinstance(other, ChunkedArray) or not self._aligned(other): + raise ValueError("A ChunkedArray can only be crossed with a ChunkedArray if they have the same chunk sizes") + chunks = [] + chunksizes = [] + for c, h in zip(self.chunks, other.chunks): + if isinstance(c, awkward.array.virtual.VirtualArray): + c = c.array + if isinstance(h, awkward.array.virtual.VirtualArray): + h = h.array + chunks.append(c.argcross(h, nested=nested)) + chunksizes.append(len(chunks[-1])) + return self.copy(chunks=chunks, chunksizes=chunksizes) def flattentuple(self): return self.copy(chunks=[self._util_flattentuple(x) for x in self._chunks], chunksizes=self._chunksizes)