Skip to content

Commit

Permalink
Implement VRTProcessedDataset::IRasterIO()
Browse files Browse the repository at this point in the history
  • Loading branch information
@rouault
rouault committed Jan 5, 2025
1 parent 2c7d458 commit f8b9b26
Show file tree
Hide file tree
Showing 3 changed files with 285 additions and 8 deletions.
184 changes: 184 additions & 0 deletions autotest/gdrivers/vrtprocesseddataset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"""<VRTDataset subclass='VRTProcessedDataset'>
<Input>
<SourceFilename>{src_filename}</SourceFilename>
</Input>
<ProcessingSteps>
<Step name="Affine combination of band values">
<Algorithm>BandAffineCombination</Algorithm>
<Argument name="coefficients_1">0,0,1,0,0</Argument>
<Argument name="coefficients_2">0,0,0,1,0</Argument>
<Argument name="coefficients_3">0,0,0,0,1</Argument>
<Argument name="coefficients_4">0,1,0,0,0</Argument>
</Step>
</ProcessingSteps>
</VRTDataset>
"""
)
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"""<VRTDataset subclass='VRTProcessedDataset'>
<Input>
<SourceFilename>{src_filename}</SourceFilename>
</Input>
<ProcessingSteps>
<Step name="Affine combination of band values">
<Algorithm>BandAffineCombination</Algorithm>
<Argument name="coefficients_1">0,0,1,0,0</Argument>
<Argument name="coefficients_2">0,0,0,1,0</Argument>
<Argument name="coefficients_3">0,0,0,0,1</Argument>
<Argument name="coefficients_4">0,1,0,0,0</Argument>
</Step>
</ProcessingSteps>
</VRTDataset>
"""
)

# 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")
9 changes: 9 additions & 0 deletions frmts/vrt/vrtdataset.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
100 changes: 92 additions & 8 deletions frmts/vrt/vrtprocesseddataset.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<size_t>(iY) * nBufXSize *
nOutBands) *
nLastDTSize,
eLastDT, nLastDTSize * nOutBands,
pDst +
static_cast<size_t>(iY) * nBlockXSize * nDTSize,
eDataType, nDTSize, nBufXSize);
GDALCopyWords64(poVRTDS->m_abyInput.data() +
(iDstBand + static_cast<size_t>(iY) *
nBufXSize * nOutBands) *
nLastDTSize,
eLastDT, nLastDTSize * nOutBands,
pDst +
static_cast<size_t>(iY) * nBlockXSize * nDTSize,
eDataType, nDTSize, nBufXSize);
}
if (poBlock)
poBlock->DropLock();
Expand All @@ -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<size_t>(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<GByte *>(pData) + iBand * nBandSpace,
eBufType, nBufTypeSize,
static_cast<size_t>(nBufXSize) * nBufYSize);
}
return CE_None;
}

return VRTDataset::IRasterIO(eRWFlag, nXOff, nYOff, nXSize, nYSize, pData,
nBufXSize, nBufYSize, eBufType, nBandCount,
panBandMap, nPixelSpace, nLineSpace,
nBandSpace, psExtraArg);
}

/*! @endcond */

/************************************************************************/
Expand Down

0 comments on commit f8b9b26

Please sign in to comment.