def as_rasterio()

in python/sedona/spark/raster/sedona_raster.py [0:0]

• 66 lines of code
• 13 McCabe index (conditional complexity)

    def as_rasterio(self) -> DatasetReader:
        if self.rasterio_dataset_reader is not None:
            return self.rasterio_dataset_reader

        affine = Affine.from_gdal(
            self._affine_trans.ip_x,
            self._affine_trans.scale_x,
            self._affine_trans.skew_x,
            self._affine_trans.ip_y,
            self._affine_trans.skew_y,
            self._affine_trans.scale_y,
        )
        num_bands = len(self._bands_meta)

        data_array = np.ascontiguousarray(self.as_numpy())

        dtype = data_array.dtype
        if dtype == np.uint8:
            data_type = "Byte"
        elif dtype == np.int8:
            data_type = "Int8"
        elif dtype == np.uint16:
            data_type = "Uint16"
        elif dtype == np.int16:
            data_type = "Int16"
        elif dtype == np.uint32:
            data_type = "UInt32"
        elif dtype == np.int32:
            data_type = "Int32"
        elif dtype == np.float32:
            data_type = "Float32"
        elif dtype == np.float64:
            data_type = "Float64"
        elif dtype == np.int64:
            data_type = "Int64"
        elif dtype == np.uint64:
            data_type = "Uint64"
        else:
            raise RuntimeError("unknown dtype: " + str(dtype))

        arr_if = data_array.__array_interface__
        data_pointer = arr_if["data"][0]
        geotransform = (
            f"{self._affine_trans.ip_x}/{self._affine_trans.scale_x}/{self._affine_trans.skew_x}/"
            + f"{self._affine_trans.ip_y}/{self._affine_trans.skew_y}/{self._affine_trans.scale_y}"
        )
        desc = (
            f"MEM:::DATAPOINTER={data_pointer},PIXELS={self._width},LINES={self._height},BANDS={num_bands},"
            + f"DATATYPE={data_type},GEOTRANSFORM={geotransform}"
        )

        # If we are using GDAL >= 3.7, we can use the SPATIALREFERENCE
        # parameter; otherwise we have to wrap the MEM dataset with an VRT to
        # set up the SRS.
        if GDAL_VERSION.at_least(rasterio.env.GDALVersion(3, 7)):
            escaped_srs = json.dumps(self._crs_wkt.replace("\n", ""))
            desc += f",SPATIALREFERENCE={escaped_srs}"
            dataset = _rasterio_open(desc, driver="MEM")
        else:
            # construct a VRT to wrap this MEM dataset, with SRS set up properly
            vrt_xml = _generate_vrt_xml(
                desc,
                data_type,
                self._width,
                self._height,
                geotransform.replace("/", ","),
                self._crs_wkt,
                0,
                0,
                list(range(num_bands)),
            )
            self.rasterio_memfile = MemoryFile(vrt_xml, ext=".vrt")
            dataset = _rasterio_open_memfile(self.rasterio_memfile, driver="VRT")

        # XXX: dataset does not copy the data held by data_array, so we set
        # data_array as a property of dataset to make sure that the lifetime of
        # data_array is as long as dataset, otherwise we may see band data
        # corruption.
        dataset.mem_data_array = data_array
        return dataset