Class NdarrayCodec

petastorm/codecs.py:133–171

Encodes numpy ndarray into, or decodes an ndarray from, a spark dataframe field.

class NdarrayCodec(DataframeColumnCodec):
    """Encodes numpy ndarray into, or decodes an ndarray from, a spark dataframe field."""

    def encode(self, unischema_field, value):
        expected_dtype = unischema_field.numpy_dtype
        if isinstance(value, np.ndarray):
            if expected_dtype != value.dtype.type:
                raise ValueError('Unexpected type of {} feature. '
                                 'Expected {}. Got {}'.format(unischema_field.name, expected_dtype, value.dtype))

            expected_shape = unischema_field.shape
            if not _is_compliant_shape(value.shape, expected_shape):
                raise ValueError('Unexpected dimensions of {} feature. '
                                 'Expected {}. Got {}'.format(unischema_field.name, expected_shape, value.shape))
        else:
            raise ValueError('Unexpected type of {} feature. '
                             'Expected ndarray of {}. Got {}'.format(unischema_field.name, expected_dtype, type(value)))

        memfile = BytesIO()
        np.save(memfile, value)
        return bytearray(memfile.getvalue())

    def decode(self, unischema_field, value):
        memfile = BytesIO(value)
        return np.load(memfile)

    def spark_dtype(self):
        # Lazy loading pyspark to avoid creating pyspark dependency on data reading code path
        # (currently works only with make_batch_reader). We should move all pyspark related code into a separate module
        import pyspark.sql.types as sql_types

        return sql_types.BinaryType()

    def __str__(self):
        """Represent this as the following form:

        >>> NdarrayCodec()
        """
        return f'{type(self).__name__}()'


class CompressedNdarrayCodec(DataframeColumnCodec):