Class ScalarCodec

petastorm/codecs.py:215–271

Encodes a scalar into a spark dataframe field.

class ScalarCodec(DataframeColumnCodec):
    """Encodes a scalar into a spark dataframe field."""

    def __init__(self, spark_type):
        """Constructs a codec.

        :param spark_type: an instance of a Type object from :mod:`pyspark.sql.types`
        """
        self._spark_type = spark_type

    def encode(self, unischema_field, value):
        # 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

        # We treat ndarrays with shape=() as scalars
        unsized_numpy_array = isinstance(value, np.ndarray) and value.shape == ()
        # Validate the input to be a scalar (or an unsized numpy array)
        if not unsized_numpy_array and hasattr(value, '__len__') and (not isinstance(value, str)):
            raise TypeError('Expected a scalar as a value for field \'{}\'. '
                            'Got a non-numpy type\'{}\''.format(unischema_field.name, type(value)))

        if unischema_field.shape:
            raise ValueError('The shape field of unischema_field \'%s\' must be an empty tuple (i.e. \'()\' '
                             'to indicate a scalar. However, the actual shape is %s',
                             unischema_field.name, unischema_field.shape)
        if isinstance(self._spark_type, (sql_types.ByteType, sql_types.ShortType, sql_types.IntegerType,
                                         sql_types.LongType)):
            return int(value)
        if isinstance(self._spark_type, (sql_types.FloatType, sql_types.DoubleType)):
            return float(value)
        if isinstance(self._spark_type, sql_types.BooleanType):
            return bool(value)
        if isinstance(self._spark_type, sql_types.StringType):
            if not isinstance(value, str):
                raise ValueError(
                    'Expected a string value for field {}. Got type {}'.format(unischema_field.name, type(value)))
            return str(value)

        return value

    def decode(self, unischema_field, value):
        # We are using pyarrow.serialize that does not support Decimal field types.
        # Tensorflow does not support Decimal types neither. We convert all decimals to
        # strings hence prevent Decimals from getting into anywhere in the reader. We may
        # choose to resurrect Decimals support in the future.
        return unischema_field.numpy_dtype(value)

    def spark_dtype(self):
        return self._spark_type

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

        >>> ScalarCodec(spark_type)
        """
        return f'{type(self).__name__}({type(self._spark_type).__name__}())'