bolt/spark/chunk.py
def map(self, func, value_shape=None, dtype=None):
        """
        Apply an array -> array function on each subarray.

        The function can change the shape of the subarray, but only along
        dimensions that are not chunked.

        Parameters
        ----------
        func : function
            Function of a single subarray to apply

        value_shape:
            Known shape of chunking plan after the map

        dtype: numpy.dtype, optional, default=None
            Known dtype of values resulting from operation

        Returns
        -------
        ChunkedArray
        """

        if value_shape is None or dtype is None:
            # try to compute the size of each mapped element by applying func to a random array
            try:
                mapped = func(random.randn(*self.plan).astype(self.dtype))
            except Exception:
                first = self._rdd.first()
                if first:
                    # eval func on the first element
                    mapped = func(first[1])
            if value_shape is None:
                value_shape = mapped.shape
            if dtype is None:
                dtype = mapped.dtype

        chunked_dims = where(self.plan != self.vshape)[0]
        unchunked_dims = where(self.plan == self.vshape)[0]

        # check that no dimensions are dropped
        if len(value_shape) != len(self.plan):
            raise NotImplementedError('map on ChunkedArray cannot drop dimensions')

        # check that chunked dimensions did not change shape
        if any([value_shape[i] != self.plan[i] for i in chunked_dims]):
            raise ValueError('map cannot change the sizes of chunked dimensions')

        def check_and_apply(v):
            new = func(v)
            if len(unchunked_dims) > 0:
                if any([new.shape[i] != value_shape[i] for i in unchunked_dims]):
                    raise Exception("Map operation did not produce values of uniform shape.")
            if len(chunked_dims) > 0:
                if any([v.shape[i] != new.shape[i] for i in chunked_dims]):
                    raise Exception("Map operation changed the size of a chunked dimension")
            return new

        rdd = self._rdd.mapValues(check_and_apply)

        vshape = [value_shape[i] if i in unchunked_dims else self.vshape[i] for i in range(len(self.vshape))]
        newshape = r_[self.kshape, vshape].astype(int).tolist()

        return self._constructor(rdd, shape=tuple(newshape), dtype=dtype,
                                 plan=asarray(value_shape)).__finalize__(self)
def map(self, func, value_shape=None, dtype=None):
        if value_shape is None or dtype is None:
            # try to compute the size of each mapped element by applying func to a random array
            try:
                mapped = func(random.randn(*self.plan).astype(self.dtype))
            except Exception:
                first = self._rdd.first()
                if first:
                    # eval func on the first element
                    mapped = func(first[1])
            if value_shape is None:
                value_shape = mapped.shape
            if dtype is None:
                dtype = mapped.dtype

        chunked_dims = where(self.plan != self.vshape)[0]
        unchunked_dims = where(self.plan == self.vshape)[0]

        # check that no dimensions are dropped
        if len(value_shape) != len(self.plan):
            raise NotImplementedError('map on ChunkedArray cannot drop dimensions')

        # check that chunked dimensions did not change shape
        if any([value_shape[i] != self.plan[i] for i in chunked_dims]):
            raise ValueError('map cannot change the sizes of chunked dimensions')

        def check_and_apply(v):
            new = func(v)
            if len(unchunked_dims) > 0:
                if any([new.shape[i] != value_shape[i] for i in unchunked_dims]):
                    raise Exception("Map operation did not produce values of uniform shape.")
            if len(chunked_dims) > 0:
                if any([v.shape[i] != new.shape[i] for i in chunked_dims]):
                    raise Exception("Map operation changed the size of a chunked dimension")
            return new

        rdd = self._rdd.mapValues(check_and_apply)

        vshape = [value_shape[i] if i in unchunked_dims else self.vshape[i] for i in range(len(self.vshape))]
        newshape = r_[self.kshape, vshape].astype(int).tolist()

        return self._constructor(rdd, shape=tuple(newshape), dtype=dtype,
                                 plan=asarray(value_shape)).__finalize__(self)
def map ( self , func , value_shape = None , dtype = None ) : if value_shape is None or dtype is None : # try to compute the size of each mapped element by applying func to a random array try : mapped = func ( random . randn ( * self . plan ) . astype ( self . dtype ) ) except Exception : first = self . _rdd . first ( ) if first : # eval func on the first element mapped = func ( first [ 1 ] ) if value_shape is None : value_shape = mapped . shape if dtype is None : dtype = mapped . dtype chunked_dims = where ( self . plan != self . vshape ) [ 0 ] unchunked_dims = where ( self . plan == self . vshape ) [ 0 ] # check that no dimensions are dropped if len ( value_shape ) != len ( self . plan ) : raise NotImplementedError ( 'map on ChunkedArray cannot drop dimensions' ) # check that chunked dimensions did not change shape if any ( [ value_shape [ i ] != self . plan [ i ] for i in chunked_dims ] ) : raise ValueError ( 'map cannot change the sizes of chunked dimensions' ) def check_and_apply ( v ) : new = func ( v ) if len ( unchunked_dims ) > 0 : if any ( [ new . shape [ i ] != value_shape [ i ] for i in unchunked_dims ] ) : raise Exception ( "Map operation did not produce values of uniform shape." ) if len ( chunked_dims ) > 0 : if any ( [ v . shape [ i ] != new . shape [ i ] for i in chunked_dims ] ) : raise Exception ( "Map operation changed the size of a chunked dimension" ) return new rdd = self . _rdd . mapValues ( check_and_apply ) vshape = [ value_shape [ i ] if i in unchunked_dims else self . vshape [ i ] for i in range ( len ( self . vshape ) ) ] newshape = r_ [ self . kshape , vshape ] . astype ( int ) . tolist ( ) return self . _constructor ( rdd , shape = tuple ( newshape ) , dtype = dtype , plan = asarray ( value_shape ) ) . __finalize__ ( self )
Apply an array -> array function on each subarray.

        The function can change the shape of the subarray, but only along
        dimensions that are not chunked.

        Parameters
        ----------
        func : function
            Function of a single subarray to apply

        value_shape:
            Known shape of chunking plan after the map

        dtype: numpy.dtype, optional, default=None
            Known dtype of values resulting from operation

        Returns
        -------
        ChunkedArray
Apply an array - > array function on each subarray .