From f8b9b265052ddff28538bce4039b702b427c627d Mon Sep 17 00:00:00 2001 From: Even Rouault Date: Sun, 5 Jan 2025 18:27:32 +0100 Subject: [PATCH] Implement VRTProcessedDataset::IRasterIO() --- autotest/gdrivers/vrtprocesseddataset.py | 184 +++++++++++++++++++++++ frmts/vrt/vrtdataset.h | 9 ++ frmts/vrt/vrtprocesseddataset.cpp | 100 +++++++++++- 3 files changed, 285 insertions(+), 8 deletions(-) diff --git a/autotest/gdrivers/vrtprocesseddataset.py b/autotest/gdrivers/vrtprocesseddataset.py index e3ff9a1bf787..dfbaa3e7025c 100755 --- a/autotest/gdrivers/vrtprocesseddataset.py +++ b/autotest/gdrivers/vrtprocesseddataset.py @@ -1326,3 +1326,187 @@ def test_vrtprocesseddataset_OutputBands(): match="Invalid value for OutputBands.dataType", ): gdal.Open("data/vrt/processed_OutputBands_USER_PROVIDED_invalid_type.vrt") + + +############################################################################### +# Test VRTProcessedDataset::RasterIO() + + +def test_vrtprocesseddataset_RasterIO(tmp_vsimem): + + src_filename = str(tmp_vsimem / "src.tif") + src_ds = gdal.GetDriverByName("GTiff").Create(src_filename, 2, 3, 4) + src_ds.GetRasterBand(1).WriteArray(np.array([[1, 2], [3, 4], [5, 6]])) + src_ds.GetRasterBand(2).WriteArray(np.array([[7, 8], [9, 10], [11, 12]])) + src_ds.GetRasterBand(3).WriteArray(np.array([[13, 14], [15, 16], [17, 18]])) + src_ds.GetRasterBand(4).WriteArray(np.array([[19, 20], [21, 22], [23, 24]])) + src_ds.BuildOverviews("NEAR", [2]) + src_ds = None + + ds = gdal.Open( + f""" + + {src_filename} + + + + BandAffineCombination + 0,0,1,0,0 + 0,0,0,1,0 + 0,0,0,0,1 + 0,1,0,0,0 + + + + """ + ) + assert ds.RasterXSize == 2 + assert ds.RasterYSize == 3 + assert ds.RasterCount == 4 + + # Optimized code path with INTERLEAVE=BAND + np.testing.assert_equal( + ds.ReadAsArray(), + np.array( + [ + [[7, 8], [9, 10], [11, 12]], + [[13, 14], [15, 16], [17, 18]], + [[19, 20], [21, 22], [23, 24]], + [[1, 2], [3, 4], [5, 6]], + ] + ), + ) + + # Optimized code path with INTERLEAVE=BAND but buf_type != native type + np.testing.assert_equal( + ds.ReadAsArray(buf_type=gdal.GDT_Int16), + np.array( + [ + [[7, 8], [9, 10], [11, 12]], + [[13, 14], [15, 16], [17, 18]], + [[19, 20], [21, 22], [23, 24]], + [[1, 2], [3, 4], [5, 6]], + ] + ), + ) + + # Optimized code path with INTERLEAVE=BAND + np.testing.assert_equal( + ds.ReadAsArray(1, 2, 1, 1), + np.array([[[12]], [[18]], [[24]], [[6]]]), + ) + + # Optimized code path with INTERLEAVE=PIXEL + np.testing.assert_equal( + ds.ReadAsArray(interleave="PIXEL"), + np.array( + [ + [[7, 13, 19, 1], [8, 14, 20, 2]], + [[9, 15, 21, 3], [10, 16, 22, 4]], + [[11, 17, 23, 5], [12, 18, 24, 6]], + ] + ), + ) + + # Optimized code path with INTERLEAVE=PIXEL but buf_type != native type + np.testing.assert_equal( + ds.ReadAsArray(interleave="PIXEL", buf_type=gdal.GDT_Int16), + np.array( + [ + [[7, 13, 19, 1], [8, 14, 20, 2]], + [[9, 15, 21, 3], [10, 16, 22, 4]], + [[11, 17, 23, 5], [12, 18, 24, 6]], + ] + ), + ) + + # Optimized code path with INTERLEAVE=PIXEL + np.testing.assert_equal( + ds.ReadAsArray(1, 2, 1, 1, interleave="PIXEL"), + np.array([[[12, 18, 24, 6]]]), + ) + + # Not optimized INTERLEAVE=BAND because not enough bands + np.testing.assert_equal( + ds.ReadAsArray(band_list=[1, 2, 3]), + np.array( + [ + [[7, 8], [9, 10], [11, 12]], + [[13, 14], [15, 16], [17, 18]], + [[19, 20], [21, 22], [23, 24]], + ] + ), + ) + + # Not optimized INTERLEAVE=BAND because of out-of-order band list + np.testing.assert_equal( + ds.ReadAsArray(band_list=[4, 1, 2, 3]), + np.array( + [ + [[1, 2], [3, 4], [5, 6]], + [[7, 8], [9, 10], [11, 12]], + [[13, 14], [15, 16], [17, 18]], + [[19, 20], [21, 22], [23, 24]], + ] + ), + ) + + # Not optimized INTERLEAVE=PIXEL because of out-of-order band list + np.testing.assert_equal( + ds.ReadAsArray(interleave="PIXEL", band_list=[4, 1, 2, 3]), + np.array( + [ + [[1, 7, 13, 19], [2, 8, 14, 20]], + [[3, 9, 15, 21], [4, 10, 16, 22]], + [[5, 11, 17, 23], [6, 12, 18, 24]], + ] + ), + ) + + # Optimized code path with overviews + assert ds.GetRasterBand(1).GetOverview(0).XSize == 1 + assert ds.GetRasterBand(1).GetOverview(0).YSize == 2 + np.testing.assert_equal( + ds.ReadAsArray(buf_xsize=1, buf_ysize=2), + np.array([[[7], [11]], [[13], [17]], [[19], [23]], [[1], [5]]]), + ) + + # Non-optimized code path with overviews + np.testing.assert_equal( + ds.ReadAsArray(buf_xsize=1, buf_ysize=1), + np.array([[[11]], [[17]], [[23]], [[5]]]), + ) + + # I/O error + gdal.GetDriverByName("GTiff").Create( + src_filename, 1024, 1024, 4, options=["TILED=YES"] + ) + f = gdal.VSIFOpenL(src_filename, "rb+") + gdal.VSIFTruncateL(f, 4096) + gdal.VSIFCloseL(f) + + ds = gdal.Open( + f""" + + {src_filename} + + + + BandAffineCombination + 0,0,1,0,0 + 0,0,0,1,0 + 0,0,0,0,1 + 0,1,0,0,0 + + + + """ + ) + + # Error in INTERLEAVE=BAND optimized code path + with pytest.raises(Exception): + ds.ReadAsArray() + + # Error in INTERLEAVE=PIXEL optimized code path + with pytest.raises(Exception): + ds.ReadAsArray(interleave="PIXEL") diff --git a/frmts/vrt/vrtdataset.h b/frmts/vrt/vrtdataset.h index bdb148343f1d..e596fcad1e5d 100644 --- a/frmts/vrt/vrtdataset.h +++ b/frmts/vrt/vrtdataset.h @@ -665,6 +665,15 @@ class VRTProcessedDataset final : public VRTDataset #pragma GCC diagnostic pop #endif + protected: + virtual CPLErr IRasterIO(GDALRWFlag eRWFlag, int nXOff, int nYOff, + int nXSize, int nYSize, void *pData, int nBufXSize, + int nBufYSize, GDALDataType eBufType, + int nBandCount, BANDMAP_TYPE panBandMap, + GSpacing nPixelSpace, GSpacing nLineSpace, + GSpacing nBandSpace, + GDALRasterIOExtraArg *psExtraArg) override; + private: friend class VRTProcessedRasterBand; diff --git a/frmts/vrt/vrtprocesseddataset.cpp b/frmts/vrt/vrtprocesseddataset.cpp index d8b0875618a4..e21445ec1fc2 100644 --- a/frmts/vrt/vrtprocesseddataset.cpp +++ b/frmts/vrt/vrtprocesseddataset.cpp @@ -1301,14 +1301,14 @@ CPLErr VRTProcessedRasterBand::IReadBlock(int nBlockXOff, int nBlockYOff, } for (int iY = 0; iY < nBufYSize; ++iY) { - GDALCopyWords(poVRTDS->m_abyInput.data() + - (iDstBand + static_cast(iY) * nBufXSize * - nOutBands) * - nLastDTSize, - eLastDT, nLastDTSize * nOutBands, - pDst + - static_cast(iY) * nBlockXSize * nDTSize, - eDataType, nDTSize, nBufXSize); + GDALCopyWords64(poVRTDS->m_abyInput.data() + + (iDstBand + static_cast(iY) * + nBufXSize * nOutBands) * + nLastDTSize, + eLastDT, nLastDTSize * nOutBands, + pDst + + static_cast(iY) * nBlockXSize * nDTSize, + eDataType, nDTSize, nBufXSize); } if (poBlock) poBlock->DropLock(); @@ -1317,6 +1317,90 @@ CPLErr VRTProcessedRasterBand::IReadBlock(int nBlockXOff, int nBlockYOff, return CE_None; } +/************************************************************************/ +/* VRTProcessedDataset::IRasterIO() */ +/************************************************************************/ + +CPLErr VRTProcessedDataset::IRasterIO( + GDALRWFlag eRWFlag, int nXOff, int nYOff, int nXSize, int nYSize, + void *pData, int nBufXSize, int nBufYSize, GDALDataType eBufType, + int nBandCount, BANDMAP_TYPE panBandMap, GSpacing nPixelSpace, + GSpacing nLineSpace, GSpacing nBandSpace, GDALRasterIOExtraArg *psExtraArg) +{ + // Try to pass the request to the most appropriate overview dataset. + if (nBufXSize < nXSize && nBufYSize < nYSize) + { + int bTried = FALSE; + const CPLErr eErr = TryOverviewRasterIO( + eRWFlag, nXOff, nYOff, nXSize, nYSize, pData, nBufXSize, nBufYSize, + eBufType, nBandCount, panBandMap, nPixelSpace, nLineSpace, + nBandSpace, psExtraArg, &bTried); + if (bTried) + return eErr; + } + + const auto IsSequentialBandMap = [panBandMap, nBandCount]() + { + for (int i = 0; i < nBandCount; ++i) + { + if (panBandMap[i] != i + 1) + { + return false; + } + } + return true; + }; + + // Optimize INTERLEAVE=PIXEL reading of all bands at nominal resolution. + if (eRWFlag == GF_Read && nXSize == nBufXSize && nYSize == nBufYSize && + nBandCount == nBands && + nBandSpace == GDALGetDataTypeSizeBytes(eBufType) && + nPixelSpace == nBandSpace * nBands && + nLineSpace == nPixelSpace * nBufXSize && IsSequentialBandMap()) + { + if (!ProcessRegion(nXOff, nYOff, nBufXSize, nBufYSize)) + { + return CE_Failure; + } + const auto eLastDT = m_aoSteps.back().eOutDT; + const int nLastDTSize = GDALGetDataTypeSizeBytes(eLastDT); + GDALCopyWords64(m_abyInput.data(), eLastDT, nLastDTSize, pData, + eBufType, GDALGetDataTypeSizeBytes(eBufType), + static_cast(nBufXSize) * nBufYSize * nBands); + return CE_None; + } + + // Optimize INTERLEAVE=BAND reading of all bands at nominal resolution. + else if (eRWFlag == GF_Read && nXSize == nBufXSize && nYSize == nBufYSize && + nBandCount == nBands && + nPixelSpace == GDALGetDataTypeSizeBytes(eBufType) && + nLineSpace == nPixelSpace * nBufXSize && + nBandSpace == nLineSpace * nBands && IsSequentialBandMap()) + { + if (!ProcessRegion(nXOff, nYOff, nBufXSize, nBufYSize)) + { + return CE_Failure; + } + const auto eLastDT = m_aoSteps.back().eOutDT; + const int nLastDTSize = GDALGetDataTypeSizeBytes(eLastDT); + const int nBufTypeSize = GDALGetDataTypeSizeBytes(eBufType); + for (int iBand = 0; iBand < nBands; ++iBand) + { + GDALCopyWords64(m_abyInput.data() + iBand * nLastDTSize, eLastDT, + nLastDTSize * nBands, + static_cast(pData) + iBand * nBandSpace, + eBufType, nBufTypeSize, + static_cast(nBufXSize) * nBufYSize); + } + return CE_None; + } + + return VRTDataset::IRasterIO(eRWFlag, nXOff, nYOff, nXSize, nYSize, pData, + nBufXSize, nBufYSize, eBufType, nBandCount, + panBandMap, nPixelSpace, nLineSpace, + nBandSpace, psExtraArg); +} + /*! @endcond */ /************************************************************************/