diff --git a/awkward/array/base.py b/awkward/array/base.py index 8a9864a4..02ec22c0 100644 --- a/awkward/array/base.py +++ b/awkward/array/base.py @@ -473,9 +473,9 @@ def _util_flatten(cls, array, axis): return array.reshape(array.shape[:axis] + (-1,) + array.shape[axis + 2:]) @classmethod - def _util_pad(cls, array, length, maskedwhen, clip): + def _util_pad(cls, array, length, maskedwhen, clip, axis): if isinstance(array, AwkwardArray): - return array.pad(length, maskedwhen=maskedwhen, clip=clip) + return array.pad(length, maskedwhen=maskedwhen, clip=clip, axis=axis) elif len(array.shape) == 1: raise ValueError("pad cannot be applied to scalars") diff --git a/awkward/array/chunked.py b/awkward/array/chunked.py index 833bc4f0..6b542822 100644 --- a/awkward/array/chunked.py +++ b/awkward/array/chunked.py @@ -630,8 +630,8 @@ def flatten(self, axis=0): out.knowchunksizes() return out - def pad(self, length, maskedwhen=True, clip=False): - return self.copy(chunks=[self._util_pad(x, length, maskedwhen, clip) for x in self._chunks], chunksizes=self.chunksizes) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self.copy(chunks=[self._util_pad(x, length, maskedwhen, clip, axis) for x in self._chunks], chunksizes=self.chunksizes) def regular(self): self._valid() diff --git a/awkward/array/indexed.py b/awkward/array/indexed.py index ba6ab73a..85bc9056 100644 --- a/awkward/array/indexed.py +++ b/awkward/array/indexed.py @@ -247,8 +247,8 @@ def flatten(self, axis=0): self._valid() return self._util_flatten(self._content[self._index], axis) - def pad(self, length, maskedwhen=True, clip=False): - return self._util_pad(self._content[self._index], length, maskedwhen, clip) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self._util_pad(self._content[self._index], length, maskedwhen, clip, axis) def regular(self): self._valid() @@ -725,8 +725,8 @@ def flatten(self, axis=0): out[mask] = content[self._inverse[mask]] return out - def pad(self, length, maskedwhen=True, clip=False): - return self._util_pad(self._content.dense, length, maskedwhen, clip) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self._util_pad(self._content.dense, length, maskedwhen, clip, axis) def regular(self): self._valid() diff --git a/awkward/array/jagged.py b/awkward/array/jagged.py index 37a7da08..1b73255d 100644 --- a/awkward/array/jagged.py +++ b/awkward/array/jagged.py @@ -1781,7 +1781,13 @@ def ready(x): table = first.Table.named("tuple", columns1, *columns2) return first.JaggedArray(first._starts, first._stops, table) - def pad(self, length, maskedwhen=True, clip=False): + def pad(self, length, maskedwhen=True, clip=False, axis=0): + if not self._util_isinteger(axis) or axis < 0: + raise TypeError("axis must be a non-negative integer (can't count from the end)") + + if axis > 0: + return type(self).fromcounts(self.counts, self._util_pad(self._content, length, maskedwhen, clip, axis - 1)) + flatstarts = self._starts.reshape(-1) if clip: diff --git a/awkward/array/masked.py b/awkward/array/masked.py index 4520d299..1ad41d63 100644 --- a/awkward/array/masked.py +++ b/awkward/array/masked.py @@ -291,8 +291,8 @@ def flattentuple(self): def flatten(self, axis=0): return self._util_flatten(self._content[self.boolmask(maskedwhen=False)], axis) - def pad(self, length, maskedwhen=True, clip=False): - return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip)) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip, axis)) def regular(self): self._valid() diff --git a/awkward/array/objects.py b/awkward/array/objects.py index 453b1005..f811867d 100644 --- a/awkward/array/objects.py +++ b/awkward/array/objects.py @@ -262,8 +262,8 @@ def flattentuple(self): def flatten(self, axis=0): return self.copy(content=self._util_flatten(self._content, axis)) - def pad(self, length, maskedwhen=True, clip=False): - return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip)) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip, axis)) def regular(self): return self.numpy.array(self) @@ -652,7 +652,13 @@ def flatten(self, axis=0): else: return self.fromjagged(self.JaggedArray.fromcounts([len(content)], content)) - def pad(self, length, maskedwhen=None, clip=False): + def pad(self, length, maskedwhen=None, clip=False, axis=0): + if not self._util_isinteger(axis) or axis < 0: + raise TypeError("axis must be a non-negative integer (can't count from the end)") + + if axis > 0: + raise ValueError("axis too deep for StringArray") + if maskedwhen is None: maskedwhen = ord(b" ") elif not isinstance(maskedwhen, bytes) or not len(maskedwhen) == 1: diff --git a/awkward/array/table.py b/awkward/array/table.py index 2fc6b77a..5816a5be 100644 --- a/awkward/array/table.py +++ b/awkward/array/table.py @@ -767,10 +767,10 @@ def argcross(self, other, nested=False): def flatten(self, axis=0): raise ValueError("cannot flatten through a Table") - def pad(self, length, maskedwhen=True, clip=False): + def pad(self, length, maskedwhen=True, clip=False, axis=0): out = self.copy(contents={}) for n, x in self._contents.items(): - out[n] = self._util_pad(x, length, maskedwhen, clip) + out[n] = self._util_pad(x, length, maskedwhen, clip, axis) return out def regular(self): diff --git a/awkward/array/union.py b/awkward/array/union.py index 05593154..acab2dff 100644 --- a/awkward/array/union.py +++ b/awkward/array/union.py @@ -468,9 +468,9 @@ def flattentuple(self): def flatten(self, axis=0): raise NotImplementedError("flatten not yet implemented for UnionArray") - def pad(self, length, maskedwhen=True, clip=False): + def pad(self, length, maskedwhen=True, clip=False, axis=0): self._valid() - return self.copy(contents=[self._util_pad(x, length, maskedwhen, clip) for x in self._contents]) + return self.copy(contents=[self._util_pad(x, length, maskedwhen, clip, axis) for x in self._contents]) def regular(self): self._valid() diff --git a/awkward/array/virtual.py b/awkward/array/virtual.py index f5acf2a6..96c66d54 100644 --- a/awkward/array/virtual.py +++ b/awkward/array/virtual.py @@ -439,8 +439,8 @@ def flattentuple(self): def flatten(self, axis=0): return self._util_flatten(self.array, axis) - def pad(self, length, maskedwhen=True, clip=False): - return self._util_pad(self.array, length, maskedwhen, clip) + def pad(self, length, maskedwhen=True, clip=False, axis=0): + return self._util_pad(self.array, length, maskedwhen, clip, axis) def regular(self): return self._util_regular(self.array) diff --git a/awkward/version.py b/awkward/version.py index 7a464001..bbbe23da 100644 --- a/awkward/version.py +++ b/awkward/version.py @@ -4,7 +4,7 @@ import re -__version__ = "0.12.1" +__version__ = "0.12.2" version = __version__ version_info = tuple(re.split(r"[-\.]", __version__))