Line data Source code
1 : /******************************************************************************
2 : *
3 : * Project: GDAL
4 : * Purpose: Zarr driver, ZarrV3CodecSequence class
5 : * Author: Even Rouault <even dot rouault at spatialys.com>
6 : *
7 : ******************************************************************************
8 : * Copyright (c) 2023, Even Rouault <even dot rouault at spatialys.com>
9 : *
10 : * SPDX-License-Identifier: MIT
11 : ****************************************************************************/
12 :
13 : #include "zarr_v3_codec.h"
14 :
15 : /************************************************************************/
16 : /* ZarrV3CodecSequence::Clone() */
17 : /************************************************************************/
18 :
19 96 : std::unique_ptr<ZarrV3CodecSequence> ZarrV3CodecSequence::Clone() const
20 : {
21 96 : auto poClone = std::make_unique<ZarrV3CodecSequence>(m_oInputArrayMetadata);
22 164 : for (const auto &poCodec : m_apoCodecs)
23 68 : poClone->m_apoCodecs.emplace_back(poCodec->Clone());
24 96 : poClone->m_oCodecArray = m_oCodecArray.Clone();
25 96 : poClone->m_bPartialDecodingPossible = m_bPartialDecodingPossible;
26 96 : return poClone;
27 : }
28 :
29 : /************************************************************************/
30 : /* ZarrV3CodecSequence::InitFromJson() */
31 : /************************************************************************/
32 :
33 2678 : bool ZarrV3CodecSequence::InitFromJson(const CPLJSONObject &oCodecs,
34 : ZarrArrayMetadata &oOutputArrayMetadata)
35 : {
36 2678 : if (oCodecs.GetType() != CPLJSONObject::Type::Array)
37 : {
38 0 : CPLError(CE_Failure, CPLE_AppDefined, "codecs is not an array");
39 0 : return false;
40 : }
41 5356 : auto oCodecsArray = oCodecs.ToArray();
42 :
43 5356 : ZarrArrayMetadata oInputArrayMetadata = m_oInputArrayMetadata;
44 2678 : ZarrV3Codec::IOType eLastType = ZarrV3Codec::IOType::ARRAY;
45 5356 : std::string osLastCodec;
46 :
47 : const auto InsertImplicitEndianCodecIfNeeded =
48 8121 : [this, &oInputArrayMetadata, &eLastType, &osLastCodec]()
49 : {
50 4057 : CPL_IGNORE_RET_VAL(this);
51 4057 : if (eLastType == ZarrV3Codec::IOType::ARRAY &&
52 1 : oInputArrayMetadata.oElt.nativeSize > 1 &&
53 1 : oInputArrayMetadata.oElt.nativeType !=
54 1 : DtypeElt::NativeType::STRING_ASCII &&
55 1 : oInputArrayMetadata.oElt.nativeType !=
56 : DtypeElt::NativeType::STRING_UNICODE)
57 : {
58 1 : CPLError(CE_Warning, CPLE_AppDefined,
59 : "'bytes' codec missing. Assuming little-endian storage, "
60 : "but such tolerance may be removed in future versions");
61 2 : auto poEndianCodec = std::make_unique<ZarrV3CodecBytes>();
62 2 : ZarrArrayMetadata oTmpOutputArrayMetadata;
63 2 : poEndianCodec->InitFromConfiguration(
64 2 : ZarrV3CodecBytes::GetConfiguration(true), oInputArrayMetadata,
65 : oTmpOutputArrayMetadata, /* bEmitWarnings = */ true);
66 1 : oInputArrayMetadata = std::move(oTmpOutputArrayMetadata);
67 1 : eLastType = poEndianCodec->GetOutputType();
68 1 : osLastCodec = poEndianCodec->GetName();
69 : if constexpr (!CPL_IS_LSB)
70 : {
71 : // Insert a little endian codec if we are on a big endian target
72 : m_apoCodecs.emplace_back(std::move(poEndianCodec));
73 : }
74 : }
75 4057 : };
76 :
77 2678 : bool bShardingFound = false;
78 5356 : std::vector<size_t> anBlockSizesBeforeSharding;
79 6817 : for (const auto &oCodec : oCodecsArray)
80 : {
81 4159 : if (oCodec.GetType() != CPLJSONObject::Type::Object)
82 : {
83 0 : CPLError(CE_Failure, CPLE_AppDefined, "codecs[] is not an object");
84 20 : return false;
85 : }
86 8318 : const auto osName = oCodec["name"].ToString();
87 0 : std::unique_ptr<ZarrV3Codec> poCodec;
88 4159 : if (osName == ZarrV3CodecGZip::NAME)
89 59 : poCodec = std::make_unique<ZarrV3CodecGZip>();
90 4100 : else if (osName == ZarrV3CodecBlosc::NAME)
91 4 : poCodec = std::make_unique<ZarrV3CodecBlosc>();
92 4096 : else if (osName == ZarrV3CodecZstd::NAME)
93 527 : poCodec = std::make_unique<ZarrV3CodecZstd>();
94 5165 : else if (osName == ZarrV3CodecBytes::NAME ||
95 1596 : osName == "endian" /* endian is the old name */)
96 1973 : poCodec = std::make_unique<ZarrV3CodecBytes>();
97 1596 : else if (osName == ZarrV3CodecTranspose::NAME)
98 83 : poCodec = std::make_unique<ZarrV3CodecTranspose>();
99 1513 : else if (osName == ZarrV3CodecCRC32C::NAME)
100 809 : poCodec = std::make_unique<ZarrV3CodecCRC32C>();
101 704 : else if (osName == ZarrV3CodecVLenUTF8::NAME)
102 14 : poCodec = std::make_unique<ZarrV3CodecVLenUTF8>();
103 690 : else if (osName == ZarrV3CodecShardingIndexed::NAME)
104 : {
105 688 : bShardingFound = true;
106 688 : poCodec = std::make_unique<ZarrV3CodecShardingIndexed>();
107 : }
108 : else
109 : {
110 2 : CPLError(CE_Failure, CPLE_NotSupported, "Unsupported codec: %s",
111 : osName.c_str());
112 2 : return false;
113 : }
114 :
115 4157 : if (poCodec->GetInputType() == ZarrV3Codec::IOType::ARRAY)
116 : {
117 2758 : if (eLastType == ZarrV3Codec::IOType::BYTES)
118 : {
119 0 : CPLError(CE_Failure, CPLE_AppDefined,
120 : "Cannot chain codec %s with %s",
121 0 : poCodec->GetName().c_str(), osLastCodec.c_str());
122 0 : return false;
123 : }
124 : }
125 : else
126 : {
127 1399 : InsertImplicitEndianCodecIfNeeded();
128 : }
129 :
130 4157 : ZarrArrayMetadata oStepOutputArrayMetadata;
131 4157 : if (osName == ZarrV3CodecShardingIndexed::NAME)
132 : {
133 688 : anBlockSizesBeforeSharding = oInputArrayMetadata.anBlockSizes;
134 : }
135 8314 : if (!poCodec->InitFromConfiguration(oCodec["configuration"],
136 : oInputArrayMetadata,
137 : oStepOutputArrayMetadata,
138 4157 : /* bEmitWarnings = */ true))
139 : {
140 18 : return false;
141 : }
142 4139 : oInputArrayMetadata = std::move(oStepOutputArrayMetadata);
143 4139 : eLastType = poCodec->GetOutputType();
144 4139 : osLastCodec = poCodec->GetName();
145 :
146 4139 : if (!poCodec->IsNoOp())
147 2219 : m_apoCodecs.emplace_back(std::move(poCodec));
148 : }
149 :
150 2658 : if (bShardingFound)
151 : {
152 670 : m_bPartialDecodingPossible =
153 670 : (m_apoCodecs.back()->GetName() == ZarrV3CodecShardingIndexed::NAME);
154 670 : if (!m_bPartialDecodingPossible)
155 : {
156 138 : m_bPartialDecodingPossible = false;
157 : // This is not an implementation limitation, but the result of a
158 : // badly thought dataset. Zarr-Python also emits a similar warning.
159 138 : CPLError(
160 : CE_Warning, CPLE_AppDefined,
161 : "Sharding codec found, but not in last position. Consequently "
162 : "partial shard decoding will not be possible");
163 : oInputArrayMetadata.anBlockSizes =
164 138 : std::move(anBlockSizesBeforeSharding);
165 : }
166 : }
167 :
168 2658 : InsertImplicitEndianCodecIfNeeded();
169 :
170 2658 : m_oCodecArray = oCodecs.Clone();
171 2658 : oOutputArrayMetadata = std::move(oInputArrayMetadata);
172 2658 : return true;
173 : }
174 :
175 : /************************************************************************/
176 : /* ZarrV3CodecBytes::AllocateBuffer() */
177 : /************************************************************************/
178 :
179 44406 : bool ZarrV3CodecSequence::AllocateBuffer(ZarrByteVectorQuickResize &abyBuffer,
180 : size_t nEltCount)
181 : {
182 44406 : if (!m_apoCodecs.empty())
183 : {
184 33029 : const size_t nRawSize =
185 33029 : nEltCount * m_oInputArrayMetadata.oElt.nativeSize;
186 : // Grow the temporary buffer a bit beyond the uncompressed size
187 33029 : const size_t nMaxSize = nRawSize + nRawSize / 3 + 64;
188 : try
189 : {
190 33029 : m_abyTmp.resize(nMaxSize);
191 : }
192 0 : catch (const std::exception &e)
193 : {
194 0 : CPLError(CE_Failure, CPLE_OutOfMemory, "%s", e.what());
195 0 : return false;
196 : }
197 33029 : m_abyTmp.resize(nRawSize);
198 :
199 : // Grow the input/output buffer too if we have several steps
200 33029 : if (m_apoCodecs.size() >= 2 && abyBuffer.capacity() < nMaxSize)
201 : {
202 222 : const size_t nSize = abyBuffer.size();
203 : try
204 : {
205 222 : abyBuffer.resize(nMaxSize);
206 : }
207 0 : catch (const std::exception &e)
208 : {
209 0 : CPLError(CE_Failure, CPLE_OutOfMemory, "%s", e.what());
210 0 : return false;
211 : }
212 222 : abyBuffer.resize(nSize);
213 : }
214 : }
215 44406 : return true;
216 : }
217 :
218 : /************************************************************************/
219 : /* ZarrV3CodecSequence::Encode() */
220 : /************************************************************************/
221 :
222 15236 : bool ZarrV3CodecSequence::Encode(ZarrByteVectorQuickResize &abyBuffer)
223 : {
224 15236 : if (!AllocateBuffer(abyBuffer,
225 15236 : MultiplyElements(m_oInputArrayMetadata.anBlockSizes)))
226 0 : return false;
227 24982 : for (const auto &poCodec : m_apoCodecs)
228 : {
229 9746 : if (!poCodec->Encode(abyBuffer, m_abyTmp))
230 0 : return false;
231 9746 : std::swap(abyBuffer, m_abyTmp);
232 : }
233 15236 : return true;
234 : }
235 :
236 : /************************************************************************/
237 : /* ZarrV3CodecSequence::Decode() */
238 : /************************************************************************/
239 :
240 28143 : bool ZarrV3CodecSequence::Decode(ZarrByteVectorQuickResize &abyBuffer)
241 : {
242 28143 : if (!AllocateBuffer(abyBuffer,
243 28143 : MultiplyElements(m_oInputArrayMetadata.anBlockSizes)))
244 0 : return false;
245 49738 : for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
246 : {
247 22488 : const auto &poCodec = *iter;
248 22488 : if (!poCodec->Decode(abyBuffer, m_abyTmp))
249 893 : return false;
250 21595 : std::swap(abyBuffer, m_abyTmp);
251 : }
252 27250 : return true;
253 : }
254 :
255 : /************************************************************************/
256 : /* ZarrV3CodecSequence::DecodePartial() */
257 : /************************************************************************/
258 :
259 1027 : bool ZarrV3CodecSequence::DecodePartial(VSIVirtualHandle *poFile,
260 : ZarrByteVectorQuickResize &abyBuffer,
261 : const std::vector<size_t> &anStartIdxIn,
262 : const std::vector<size_t> &anCountIn)
263 : {
264 1027 : CPLAssert(anStartIdxIn.size() == m_oInputArrayMetadata.anBlockSizes.size());
265 1027 : CPLAssert(anStartIdxIn.size() == anCountIn.size());
266 :
267 1027 : if (!AllocateBuffer(abyBuffer, MultiplyElements(anCountIn)))
268 0 : return false;
269 :
270 : // anStartIdxIn and anCountIn are expressed in the shape *before* encoding
271 : // We need to apply the potential transpositions before submitting them
272 : // to the decoder of the Array->Bytes decoder
273 2054 : std::vector<size_t> anStartIdx(anStartIdxIn);
274 2054 : std::vector<size_t> anCount(anCountIn);
275 2079 : for (auto &poCodec : m_apoCodecs)
276 : {
277 1052 : poCodec->ChangeArrayShapeForward(anStartIdx, anCount);
278 : }
279 :
280 1723 : for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
281 : {
282 1052 : const auto &poCodec = *iter;
283 :
284 2104 : if (!poCodec->DecodePartial(poFile, abyBuffer, m_abyTmp, anStartIdx,
285 1052 : anCount))
286 356 : return false;
287 696 : std::swap(abyBuffer, m_abyTmp);
288 : }
289 671 : return true;
290 : }
291 :
292 : /************************************************************************/
293 : /* ZarrV3CodecSequence::BatchDecodePartial() */
294 : /************************************************************************/
295 :
296 4031 : bool ZarrV3CodecSequence::BatchDecodePartial(
297 : VSIVirtualHandle *poFile, const char *pszFilename,
298 : const std::vector<std::pair<std::vector<size_t>, std::vector<size_t>>>
299 : &anRequests,
300 : std::vector<ZarrByteVectorQuickResize> &aResults)
301 : {
302 : // Only batch-decode when sharding is the sole codec. If other codecs
303 : // (e.g. transpose) precede it, indices and output need codec-specific
304 : // transformations that BatchDecodePartial does not handle.
305 4031 : if (m_apoCodecs.size() == 1)
306 : {
307 4023 : auto *poSharding = dynamic_cast<ZarrV3CodecShardingIndexed *>(
308 8046 : m_apoCodecs.back().get());
309 4023 : if (poSharding)
310 : {
311 4023 : return poSharding->BatchDecodePartial(poFile, pszFilename,
312 4023 : anRequests, aResults);
313 : }
314 : }
315 :
316 : // Fallback: sequential DecodePartial for non-sharding codec chains
317 8 : aResults.resize(anRequests.size());
318 32 : for (size_t i = 0; i < anRequests.size(); ++i)
319 : {
320 24 : if (!DecodePartial(poFile, aResults[i], anRequests[i].first,
321 24 : anRequests[i].second))
322 0 : return false;
323 : }
324 8 : return true;
325 : }
326 :
327 : /************************************************************************/
328 : /* ZarrV3CodecSequence::GetInnerMostBlockSize() */
329 : /************************************************************************/
330 :
331 1294 : std::vector<size_t> ZarrV3CodecSequence::GetInnerMostBlockSize(
332 : const std::vector<size_t> &anOuterBlockSize) const
333 : {
334 1294 : auto chunkSize = anOuterBlockSize;
335 2316 : for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
336 : {
337 1022 : const auto &poCodec = *iter;
338 1512 : if (m_bPartialDecodingPossible ||
339 490 : poCodec->GetName() != ZarrV3CodecShardingIndexed::NAME)
340 : {
341 884 : chunkSize = poCodec->GetInnerMostBlockSize(chunkSize);
342 : }
343 : }
344 1294 : return chunkSize;
345 : }
|
