bolt/spark/chunk.py
def keys_to_values(self, axes, size=None):
        """
        Move indices in the keys into the values.

        Padding on these new value-dimensions is not currently supported and is set to 0.

        Parameters
        ----------
        axes : tuple
            Axes from keys to move to values.

        size : tuple, optional, default=None
            Size of chunks for the values along the new dimensions.
            If None, then no chunking for all axes (number of chunks = 1)

        Returns
        -------
        ChunkedArray
        """
        if len(axes) == 0:
            return self

        kmask = self.kmask(axes)

        if size is None:
            size = self.kshape[kmask]

        # update properties
        newplan = r_[size, self.plan]
        newsplit = self._split - len(axes)
        newshape = tuple(r_[self.kshape[~kmask], self.kshape[kmask], self.vshape].astype(int).tolist())
        newpadding = r_[zeros(len(axes), dtype=int), self.padding]

        result = self._constructor(None, shape=newshape, split=newsplit,
                                   dtype=self.dtype, plan=newplan, padding=newpadding, ordered=True)

        # convert keys into chunk + within-chunk label
        split = self.split
        def _relabel(record):
            k, data = record
            keys, chks = asarray(k[:split], 'int'), k[split:]
            movingkeys, stationarykeys = keys[kmask], keys[~kmask]
            newchks = [int(m) for m in movingkeys/size]  # element-wise integer division that works in Python 2 and 3
            labels = mod(movingkeys, size)
            return tuple(stationarykeys) + tuple(newchks) + tuple(chks) + tuple(labels), data

        rdd = self._rdd.map(_relabel)

        # group the new chunks together
        nchunks = result.getnumber(result.plan, result.vshape)
        npartitions = int(prod(result.kshape) * prod(nchunks))
        ranges = tuple(result.kshape) + tuple(nchunks)
        n = len(axes)
        if n == 0:
            s = slice(None)
        else:
            s = slice(-n)
        partitioner = lambda k: ravel_multi_index(k[s], ranges)

        rdd = rdd.partitionBy(numPartitions=npartitions, partitionFunc=partitioner)

        # reassemble the pieces in the chunks by sorting and then stacking
        uniform = result.uniform

        def _rebuild(it):
            ordered = sorted(it, key=lambda kv: kv[0][n:])
            keys, data = zip(*ordered)

            k = keys[0][s]
            labels = asarray([x[-n:] for x in keys])

            if uniform:
                labelshape = tuple(size)
            else:
                labelshape = tuple(amax(labels, axis=0) - amin(labels, axis=0) + 1)

            valshape = data[0].shape
            fullshape = labelshape + valshape
            yield k, asarray(data).reshape(fullshape)

        result._rdd = rdd.mapPartitions(_rebuild)

        if array_equal(self.vshape, [1]):
            result._rdd = result._rdd.mapValues(lambda v: squeeze(v))
            result._shape = result.shape[:-1]
            result._plan = result.plan[:-1]

        return result
def keys_to_values(self, axes, size=None):
        """
        Move indices in the keys into the values.

        Padding on these new value-dimensions is not currently supported and is set to 0.

        Parameters
        ----------
        axes : tuple
            Axes from keys to move to values.

        size : tuple, optional, default=None
            Size of chunks for the values along the new dimensions.
            If None, then no chunking for all axes (number of chunks = 1)

        Returns
        -------
        ChunkedArray
        """
        if len(axes) == 0:
            return self

        kmask = self.kmask(axes)

        if size is None:
            size = self.kshape[kmask]

        # update properties
        newplan = r_[size, self.plan]
        newsplit = self._split - len(axes)
        newshape = tuple(r_[self.kshape[~kmask], self.kshape[kmask], self.vshape].astype(int).tolist())
        newpadding = r_[zeros(len(axes), dtype=int), self.padding]

        result = self._constructor(None, shape=newshape, split=newsplit,
                                   dtype=self.dtype, plan=newplan, padding=newpadding, ordered=True)

        # convert keys into chunk + within-chunk label
        split = self.split
        def _relabel(record):
            k, data = record
            keys, chks = asarray(k[:split], 'int'), k[split:]
            movingkeys, stationarykeys = keys[kmask], keys[~kmask]
            newchks = [int(m) for m in movingkeys/size]  # element-wise integer division that works in Python 2 and 3
            labels = mod(movingkeys, size)
            return tuple(stationarykeys) + tuple(newchks) + tuple(chks) + tuple(labels), data

        rdd = self._rdd.map(_relabel)

        # group the new chunks together
        nchunks = result.getnumber(result.plan, result.vshape)
        npartitions = int(prod(result.kshape) * prod(nchunks))
        ranges = tuple(result.kshape) + tuple(nchunks)
        n = len(axes)
        if n == 0:
            s = slice(None)
        else:
            s = slice(-n)
        partitioner = lambda k: ravel_multi_index(k[s], ranges)

        rdd = rdd.partitionBy(numPartitions=npartitions, partitionFunc=partitioner)

        # reassemble the pieces in the chunks by sorting and then stacking
        uniform = result.uniform

        def _rebuild(it):
            ordered = sorted(it, key=lambda kv: kv[0][n:])
            keys, data = zip(*ordered)

            k = keys[0][s]
            labels = asarray([x[-n:] for x in keys])

            if uniform:
                labelshape = tuple(size)
            else:
                labelshape = tuple(amax(labels, axis=0) - amin(labels, axis=0) + 1)

            valshape = data[0].shape
            fullshape = labelshape + valshape
            yield k, asarray(data).reshape(fullshape)

        result._rdd = rdd.mapPartitions(_rebuild)

        if array_equal(self.vshape, [1]):
            result._rdd = result._rdd.mapValues(lambda v: squeeze(v))
            result._shape = result.shape[:-1]
            result._plan = result.plan[:-1]

        return result
Move indices in the keys into the values.

        Padding on these new value-dimensions is not currently supported and is set to 0.

        Parameters
        ----------
        axes : tuple
            Axes from keys to move to values.

        size : tuple, optional, default=None
            Size of chunks for the values along the new dimensions.
            If None, then no chunking for all axes (number of chunks = 1)

        Returns
        -------
        ChunkedArray