/src/gdal/frmts/zarr/zarr_v3_codec_sharding.cpp
Line | Count | Source |
1 | | /****************************************************************************** |
2 | | * |
3 | | * Project: GDAL |
4 | | * Purpose: Zarr driver, "sharding_indexed" codec |
5 | | * Author: Even Rouault <even dot rouault at spatialys.com> |
6 | | * |
7 | | ****************************************************************************** |
8 | | * Copyright (c) 2026, Development Seed |
9 | | * |
10 | | * SPDX-License-Identifier: MIT |
11 | | ****************************************************************************/ |
12 | | |
13 | | #include "zarr_v3_codec.h" |
14 | | |
15 | | #include "cpl_mem_cache.h" |
16 | | #include "cpl_vsi_virtual.h" |
17 | | #include "cpl_worker_thread_pool.h" |
18 | | #include "gdal_thread_pool.h" |
19 | | |
20 | | #include <algorithm> |
21 | | #include <cinttypes> |
22 | | #include <limits> |
23 | | #include <mutex> |
24 | | |
25 | | // Process-wide LRU cache for shard indices. |
26 | | // Key: shard filepath/URL. Value: flat byte buffer of Location entries |
27 | | // (each 16 bytes: uint64_t nOffset + uint64_t nSize, host byte order). |
28 | | // Populated lazily by BatchDecodePartial. |
29 | | // Shards whose index exceeds GDAL_ZARR_SHARD_INDEX_CACHE_MAX_BYTES |
30 | | // (default 1 MiB) fall back to per-entry ReadMultiRange. |
31 | | // Thread safety: lru11::Cache<..., std::mutex> serialises all access. |
32 | | static constexpr size_t SHARD_INDEX_CACHE_ENTRIES = 64; |
33 | | static lru11::Cache<std::string, std::vector<GByte>, std::mutex> |
34 | | g_oShardIndexCache{SHARD_INDEX_CACHE_ENTRIES, 0}; |
35 | | |
36 | | // Implements https://zarr-specs.readthedocs.io/en/latest/v3/codecs/sharding-indexed/index.html |
37 | | |
38 | | void ZarrClearShardIndexCache() |
39 | 0 | { |
40 | 0 | g_oShardIndexCache.clear(); |
41 | 0 | } |
42 | | |
43 | | void ZarrEraseShardIndexFromCache(const std::string &osFilename) |
44 | 0 | { |
45 | 0 | g_oShardIndexCache.remove(osFilename); |
46 | 0 | } |
47 | | |
48 | | /************************************************************************/ |
49 | | /* ZarrV3CodecShardingIndexed() */ |
50 | | /************************************************************************/ |
51 | | |
52 | 0 | ZarrV3CodecShardingIndexed::ZarrV3CodecShardingIndexed() : ZarrV3Codec(NAME) |
53 | 0 | { |
54 | 0 | } |
55 | | |
56 | | /************************************************************************/ |
57 | | /* ZarrV3CodecShardingIndexed::Clone() */ |
58 | | /************************************************************************/ |
59 | | |
60 | | std::unique_ptr<ZarrV3Codec> ZarrV3CodecShardingIndexed::Clone() const |
61 | 0 | { |
62 | 0 | auto psClone = std::make_unique<ZarrV3CodecShardingIndexed>(); |
63 | 0 | ZarrArrayMetadata oOutputArrayMetadata; |
64 | 0 | psClone->InitFromConfiguration(std::string(), m_oConfiguration, |
65 | 0 | m_oInputArrayMetadata, oOutputArrayMetadata, |
66 | 0 | /* bEmitWarnings = */ false); |
67 | 0 | return psClone; |
68 | 0 | } |
69 | | |
70 | | /************************************************************************/ |
71 | | /* ZarrV3CodecShardingIndexed::InitFromConfiguration() */ |
72 | | /************************************************************************/ |
73 | | |
74 | | bool ZarrV3CodecShardingIndexed::InitFromConfiguration( |
75 | | const std::string &osArrayName, const CPLJSONObject &configuration, |
76 | | const ZarrArrayMetadata &oInputArrayMetadata, |
77 | | ZarrArrayMetadata &oOutputArrayMetadata, bool bEmitWarnings) |
78 | 0 | { |
79 | 0 | if (oInputArrayMetadata.anBlockSizes.empty()) |
80 | 0 | { |
81 | 0 | CPLError( |
82 | 0 | CE_Failure, CPLE_AppDefined, |
83 | 0 | "Codec sharding_indexed: sharding not supported for scalar array"); |
84 | 0 | return false; |
85 | 0 | } |
86 | | |
87 | 0 | m_oConfiguration = configuration.Clone(); |
88 | 0 | m_oInputArrayMetadata = oInputArrayMetadata; |
89 | |
|
90 | 0 | if (!configuration.IsValid() || |
91 | 0 | configuration.GetType() != CPLJSONObject::Type::Object) |
92 | 0 | { |
93 | 0 | CPLError( |
94 | 0 | CE_Failure, CPLE_AppDefined, |
95 | 0 | "Codec sharding_indexed: configuration missing or not an object"); |
96 | 0 | return false; |
97 | 0 | } |
98 | | |
99 | 0 | const auto oChunkShape = configuration["chunk_shape"].ToArray(); |
100 | 0 | if (!oChunkShape.IsValid() || |
101 | 0 | oChunkShape.GetType() != CPLJSONObject::Type::Array) |
102 | 0 | { |
103 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
104 | 0 | "Codec sharding_indexed: configuration.chunk_shape missing or " |
105 | 0 | "not an array"); |
106 | 0 | return false; |
107 | 0 | } |
108 | 0 | if (static_cast<size_t>(oChunkShape.Size()) != |
109 | 0 | m_oInputArrayMetadata.anBlockSizes.size()) |
110 | 0 | { |
111 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
112 | 0 | "Codec sharding_indexed: configuration.chunk_shape should " |
113 | 0 | "have the same shape as the array"); |
114 | 0 | return false; |
115 | 0 | } |
116 | 0 | std::vector<size_t> anCountInnerChunks; |
117 | 0 | for (int i = 0; i < oChunkShape.Size(); ++i) |
118 | 0 | { |
119 | 0 | if (oChunkShape[i].GetType() != CPLJSONObject::Type::Integer && |
120 | 0 | oChunkShape[i].GetType() != CPLJSONObject::Type::Long) |
121 | 0 | { |
122 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
123 | 0 | "Codec sharding_indexed: configuration.chunk_shape[%d] " |
124 | 0 | "should be an integer", |
125 | 0 | i); |
126 | 0 | return false; |
127 | 0 | } |
128 | 0 | const int64_t nVal = oChunkShape[i].ToLong(); |
129 | 0 | if (nVal <= 0 || |
130 | 0 | static_cast<uint64_t>(nVal) > |
131 | 0 | m_oInputArrayMetadata.anBlockSizes[i] || |
132 | 0 | (m_oInputArrayMetadata.anBlockSizes[i] % nVal) != 0) |
133 | 0 | { |
134 | 0 | CPLError( |
135 | 0 | CE_Failure, CPLE_AppDefined, |
136 | 0 | "Codec sharding_indexed: configuration.chunk_shape[%d]=%" PRId64 |
137 | 0 | " should be a strictly positive value that is a divisor of " |
138 | 0 | "%" PRIu64, |
139 | 0 | i, nVal, |
140 | 0 | static_cast<uint64_t>(m_oInputArrayMetadata.anBlockSizes[i])); |
141 | 0 | return false; |
142 | 0 | } |
143 | 0 | #ifndef __COVERITY__ |
144 | | // The following cast is safe since ZarrArray::ParseChunkSize() has |
145 | | // previously validated that m_oInputArrayMetadata.anBlockSizes[i] fits |
146 | | // on size_t |
147 | | if constexpr (sizeof(size_t) < sizeof(uint64_t)) |
148 | | { |
149 | | // coverity[result_independent_of_operands] |
150 | | CPLAssert(nVal <= std::numeric_limits<size_t>::max()); |
151 | | } |
152 | 0 | #endif |
153 | 0 | m_anInnerBlockSize.push_back(static_cast<size_t>(nVal)); |
154 | 0 | anCountInnerChunks.push_back( |
155 | 0 | static_cast<size_t>(m_oInputArrayMetadata.anBlockSizes[i] / nVal)); |
156 | 0 | } |
157 | | |
158 | 0 | const auto oCodecs = configuration["codecs"]; |
159 | 0 | if (!oCodecs.IsValid() || oCodecs.GetType() != CPLJSONObject::Type::Array) |
160 | 0 | { |
161 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
162 | 0 | "Codec sharding_indexed: configuration.codecs missing or " |
163 | 0 | "not an array"); |
164 | 0 | return false; |
165 | 0 | } |
166 | 0 | if (oCodecs.ToArray().Size() == 0) |
167 | 0 | { |
168 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
169 | 0 | "Codec sharding_indexed: configuration.codecs[] is empty"); |
170 | 0 | return false; |
171 | 0 | } |
172 | 0 | ZarrArrayMetadata inputArrayMetadataCodecs = m_oInputArrayMetadata; |
173 | 0 | inputArrayMetadataCodecs.anBlockSizes = m_anInnerBlockSize; |
174 | 0 | m_poCodecSequence = |
175 | 0 | std::make_unique<ZarrV3CodecSequence>(inputArrayMetadataCodecs); |
176 | 0 | if (!m_poCodecSequence->InitFromJson(osArrayName, oCodecs, |
177 | 0 | oOutputArrayMetadata)) |
178 | 0 | { |
179 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
180 | 0 | "Codec sharding_indexed: initialization of codecs failed"); |
181 | 0 | return false; |
182 | 0 | } |
183 | | |
184 | 0 | if (bEmitWarnings && m_poCodecSequence->SupportsPartialDecoding()) |
185 | 0 | { |
186 | | // Implementation limitation |
187 | 0 | CPLError(CE_Warning, CPLE_AppDefined, |
188 | 0 | "Nested sharding detected. For now, partial decoding is only " |
189 | 0 | "implemented on the outer-most shard level"); |
190 | 0 | } |
191 | |
|
192 | 0 | const auto oIndexCodecs = configuration["index_codecs"]; |
193 | 0 | if (!oIndexCodecs.IsValid() || |
194 | 0 | oIndexCodecs.GetType() != CPLJSONObject::Type::Array) |
195 | 0 | { |
196 | 0 | CPLError( |
197 | 0 | CE_Failure, CPLE_AppDefined, |
198 | 0 | "Codec sharding_indexed: configuration.index_codecs missing or " |
199 | 0 | "not an array"); |
200 | 0 | return false; |
201 | 0 | } |
202 | 0 | if (oIndexCodecs.ToArray().Size() == 0) |
203 | 0 | { |
204 | 0 | CPLError( |
205 | 0 | CE_Failure, CPLE_AppDefined, |
206 | 0 | "Codec sharding_indexed: configuration.index_codecs[] is empty"); |
207 | 0 | return false; |
208 | 0 | } |
209 | 0 | ZarrArrayMetadata inputArrayMetadataIndex; |
210 | 0 | inputArrayMetadataIndex.oElt.nativeType = |
211 | 0 | DtypeElt::NativeType::UNSIGNED_INT; |
212 | 0 | inputArrayMetadataIndex.oElt.nativeSize = sizeof(uint64_t); |
213 | 0 | inputArrayMetadataIndex.oElt.gdalType = |
214 | 0 | GDALExtendedDataType::Create(GDT_UInt64); |
215 | 0 | inputArrayMetadataIndex.oElt.gdalSize = sizeof(uint64_t); |
216 | 0 | inputArrayMetadataIndex.anBlockSizes = std::move(anCountInnerChunks); |
217 | | // 2 for offset and size |
218 | 0 | inputArrayMetadataIndex.anBlockSizes.push_back(2); |
219 | 0 | m_poIndexCodecSequence = |
220 | 0 | std::make_unique<ZarrV3CodecSequence>(inputArrayMetadataIndex); |
221 | 0 | ZarrArrayMetadata oOutputArrayMetadataIndex; |
222 | 0 | if (!m_poIndexCodecSequence->InitFromJson(osArrayName, oIndexCodecs, |
223 | 0 | oOutputArrayMetadataIndex)) |
224 | 0 | { |
225 | 0 | CPLError( |
226 | 0 | CE_Failure, CPLE_AppDefined, |
227 | 0 | "Codec sharding_indexed: initialization of index_codecs failed"); |
228 | 0 | return false; |
229 | 0 | } |
230 | 0 | const auto &indexCodecs = m_poIndexCodecSequence->GetCodecs(); |
231 | 0 | if (indexCodecs.empty()) |
232 | 0 | { |
233 | | // ok, there is only a "bytes" codec, optimized away if the order |
234 | | // is the one of the native architecture |
235 | 0 | } |
236 | 0 | else if (indexCodecs[0]->GetName() == ZarrV3CodecBytes::NAME || |
237 | 0 | indexCodecs[0]->GetName() == ZarrV3CodecCRC32C::NAME) |
238 | 0 | { |
239 | | // ok |
240 | 0 | } |
241 | 0 | else if (indexCodecs.size() == 2 && |
242 | 0 | indexCodecs[1]->GetName() == ZarrV3CodecCRC32C::NAME) |
243 | 0 | { |
244 | | // ok |
245 | 0 | } |
246 | 0 | else |
247 | 0 | { |
248 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
249 | 0 | "Codec sharding_indexed: this implementation only supports " |
250 | 0 | "Bytes, possibly followed by CRC32C, as index_codecs"); |
251 | 0 | return false; |
252 | 0 | } |
253 | 0 | m_bIndexHasCRC32 = (!indexCodecs.empty() && indexCodecs.back()->GetName() == |
254 | 0 | ZarrV3CodecCRC32C::NAME); |
255 | |
|
256 | 0 | const std::string osIndexLocation = |
257 | 0 | configuration.GetString("index_location", "end"); |
258 | 0 | if (osIndexLocation != "start" && osIndexLocation != "end") |
259 | 0 | { |
260 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
261 | 0 | "Codec sharding_indexed: invalid value for index_location"); |
262 | 0 | return false; |
263 | 0 | } |
264 | 0 | m_bIndexLocationAtEnd = (osIndexLocation == "end"); |
265 | |
|
266 | 0 | return true; |
267 | 0 | } |
268 | | |
269 | | /************************************************************************/ |
270 | | /* ExtractSubArrayFromLargerOne() */ |
271 | | /************************************************************************/ |
272 | | |
273 | | // Inverse of CopySubArrayIntoLargerOne(): extract a contiguous inner chunk |
274 | | // from its position in the larger shard buffer. |
275 | | static void |
276 | | ExtractSubArrayFromLargerOne(const ZarrByteVectorQuickResize &abySrc, |
277 | | const std::vector<size_t> &anSrcBlockSize, |
278 | | const std::vector<size_t> &anInnerBlockSize, |
279 | | const std::vector<size_t> &anInnerBlockIndices, |
280 | | ZarrByteVectorQuickResize &abyChunk, |
281 | | const size_t nDTSize) |
282 | 0 | { |
283 | 0 | const auto nDims = anInnerBlockSize.size(); |
284 | 0 | CPLAssert(nDims > 0); |
285 | 0 | CPLAssert(nDims == anInnerBlockIndices.size()); |
286 | 0 | CPLAssert(nDims == anSrcBlockSize.size()); |
287 | | // +1 to avoid gcc -Wnull-dereference false positives |
288 | 0 | std::vector<const GByte *> srcPtrStack(nDims + 1); |
289 | 0 | std::vector<size_t> count(nDims + 1); |
290 | 0 | std::vector<size_t> srcStride(nDims + 1); |
291 | |
|
292 | 0 | size_t nSrcStride = nDTSize; |
293 | 0 | for (size_t iDim = nDims; iDim > 0;) |
294 | 0 | { |
295 | 0 | --iDim; |
296 | 0 | srcStride[iDim] = nSrcStride; |
297 | 0 | nSrcStride *= anSrcBlockSize[iDim]; |
298 | 0 | } |
299 | |
|
300 | 0 | srcPtrStack[0] = abySrc.data(); |
301 | 0 | for (size_t iDim = 0; iDim < nDims; ++iDim) |
302 | 0 | { |
303 | 0 | CPLAssert((anInnerBlockIndices[iDim] + 1) * anInnerBlockSize[iDim] <= |
304 | 0 | anSrcBlockSize[iDim]); |
305 | 0 | srcPtrStack[0] += anInnerBlockIndices[iDim] * anInnerBlockSize[iDim] * |
306 | 0 | srcStride[iDim]; |
307 | 0 | } |
308 | 0 | GByte *pabyDst = abyChunk.data(); |
309 | |
|
310 | 0 | const size_t nLastDimSize = anInnerBlockSize.back() * nDTSize; |
311 | 0 | size_t dimIdx = 0; |
312 | 0 | lbl_next_depth: |
313 | 0 | if (dimIdx + 1 == nDims) |
314 | 0 | { |
315 | 0 | memcpy(pabyDst, srcPtrStack[dimIdx], nLastDimSize); |
316 | 0 | pabyDst += nLastDimSize; |
317 | 0 | } |
318 | 0 | else |
319 | 0 | { |
320 | 0 | count[dimIdx] = anInnerBlockSize[dimIdx]; |
321 | 0 | while (true) |
322 | 0 | { |
323 | 0 | dimIdx++; |
324 | 0 | srcPtrStack[dimIdx] = srcPtrStack[dimIdx - 1]; |
325 | 0 | goto lbl_next_depth; |
326 | 0 | lbl_return_to_caller: |
327 | 0 | dimIdx--; |
328 | 0 | if (--count[dimIdx] == 0) |
329 | 0 | break; |
330 | 0 | srcPtrStack[dimIdx] += srcStride[dimIdx]; |
331 | 0 | } |
332 | 0 | } |
333 | 0 | if (dimIdx > 0) |
334 | 0 | goto lbl_return_to_caller; |
335 | 0 | } |
336 | | |
337 | | /************************************************************************/ |
338 | | /* IsAllNoData() */ |
339 | | /************************************************************************/ |
340 | | |
341 | | // Check if a buffer consists entirely of the fill value. |
342 | | static bool IsAllNoData(const ZarrByteVectorQuickResize &abyChunk, |
343 | | const ZarrArrayMetadata &metadata) |
344 | 0 | { |
345 | 0 | const size_t nDTSize = metadata.oElt.nativeSize; |
346 | 0 | const size_t nBytes = abyChunk.size(); |
347 | |
|
348 | 0 | if (metadata.abyNoData.empty() || |
349 | 0 | metadata.abyNoData == std::vector<GByte>(nDTSize, 0)) |
350 | 0 | { |
351 | | // Zero fill value: byte-by-byte check (compiler auto-vectorizes) |
352 | 0 | const GByte *p = abyChunk.data(); |
353 | 0 | for (size_t i = 0; i < nBytes; ++i) |
354 | 0 | { |
355 | 0 | if (p[i] != 0) |
356 | 0 | return false; |
357 | 0 | } |
358 | 0 | return true; |
359 | 0 | } |
360 | 0 | else |
361 | 0 | { |
362 | | // Non-zero fill value: element-wise compare |
363 | 0 | CPLAssert(metadata.abyNoData.size() == nDTSize); |
364 | 0 | const GByte *p = abyChunk.data(); |
365 | 0 | const size_t nCount = nBytes / nDTSize; |
366 | 0 | for (size_t i = 0; i < nCount; ++i) |
367 | 0 | { |
368 | 0 | if (memcmp(p + i * nDTSize, metadata.abyNoData.data(), nDTSize) != |
369 | 0 | 0) |
370 | 0 | return false; |
371 | 0 | } |
372 | 0 | return true; |
373 | 0 | } |
374 | 0 | } |
375 | | |
376 | | /************************************************************************/ |
377 | | /* ZarrV3CodecShardingIndexed::Encode() */ |
378 | | /************************************************************************/ |
379 | | |
380 | | bool ZarrV3CodecShardingIndexed::Encode(const ZarrByteVectorQuickResize &abySrc, |
381 | | ZarrByteVectorQuickResize &abyDst) const |
382 | 0 | { |
383 | | // Compute total number of inner chunks |
384 | 0 | size_t nInnerChunks = 1; |
385 | 0 | for (size_t i = 0; i < m_anInnerBlockSize.size(); ++i) |
386 | 0 | { |
387 | 0 | const size_t nCountInnerChunksThisDim = |
388 | 0 | m_oInputArrayMetadata.anBlockSizes[i] / m_anInnerBlockSize[i]; |
389 | 0 | nInnerChunks *= nCountInnerChunksThisDim; |
390 | 0 | } |
391 | |
|
392 | 0 | const auto nDTSize = m_oInputArrayMetadata.oElt.nativeSize; |
393 | 0 | const size_t nExpectedSrcSize = |
394 | 0 | nDTSize * MultiplyElements(m_oInputArrayMetadata.anBlockSizes); |
395 | 0 | if (abySrc.size() != nExpectedSrcSize) |
396 | 0 | { |
397 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
398 | 0 | "ZarrV3CodecShardingIndexed::Encode(): input buffer size " |
399 | 0 | "(%" PRIu64 ") != expected (%" PRIu64 ")", |
400 | 0 | static_cast<uint64_t>(abySrc.size()), |
401 | 0 | static_cast<uint64_t>(nExpectedSrcSize)); |
402 | 0 | return false; |
403 | 0 | } |
404 | | |
405 | | // Index: one Location per inner chunk, initially all empty |
406 | 0 | std::vector<Location> anLocations(nInnerChunks, |
407 | 0 | {std::numeric_limits<uint64_t>::max(), |
408 | 0 | std::numeric_limits<uint64_t>::max()}); |
409 | | |
410 | | // Accumulate encoded inner chunk data |
411 | 0 | ZarrByteVectorQuickResize abyData; |
412 | | // Reserve approximate capacity: decoded chunk size is close to the upper |
413 | | // bound per chunk (compression may slightly increase size in pathological |
414 | | // cases), but avoids most reallocations. |
415 | 0 | const size_t nDecodedChunkSize = |
416 | 0 | nDTSize * MultiplyElements(m_anInnerBlockSize); |
417 | | // resize+clear = reserve (ZarrByteVectorQuickResize has no reserve()) |
418 | 0 | try |
419 | 0 | { |
420 | 0 | abyData.resize(nInnerChunks * nDecodedChunkSize); |
421 | 0 | } |
422 | 0 | catch (const std::exception &) |
423 | 0 | { |
424 | 0 | CPLError(CE_Failure, CPLE_OutOfMemory, |
425 | 0 | "Cannot allocate memory for accumulated shard data"); |
426 | 0 | return false; |
427 | 0 | } |
428 | 0 | abyData.clear(); |
429 | 0 | ZarrByteVectorQuickResize abyChunk; |
430 | |
|
431 | 0 | uint64_t nCurrentOffset = 0; |
432 | 0 | std::vector<size_t> anChunkIndices(m_anInnerBlockSize.size(), 0); |
433 | |
|
434 | 0 | for (size_t iChunk = 0; iChunk < nInnerChunks; ++iChunk) |
435 | 0 | { |
436 | | // Update chunk coordinates (same iteration order as Decode) |
437 | 0 | if (iChunk > 0) |
438 | 0 | { |
439 | 0 | size_t iDim = m_anInnerBlockSize.size() - 1; |
440 | 0 | while (++anChunkIndices[iDim] == |
441 | 0 | m_oInputArrayMetadata.anBlockSizes[iDim] / |
442 | 0 | m_anInnerBlockSize[iDim]) |
443 | 0 | { |
444 | 0 | anChunkIndices[iDim] = 0; |
445 | 0 | --iDim; |
446 | 0 | } |
447 | 0 | } |
448 | | |
449 | | // Extract this inner chunk from the shard buffer |
450 | 0 | try |
451 | 0 | { |
452 | 0 | abyChunk.resize(nDecodedChunkSize); |
453 | 0 | } |
454 | 0 | catch (const std::exception &) |
455 | 0 | { |
456 | 0 | CPLError(CE_Failure, CPLE_OutOfMemory, |
457 | 0 | "Cannot allocate memory for inner chunk"); |
458 | 0 | return false; |
459 | 0 | } |
460 | 0 | ExtractSubArrayFromLargerOne(abySrc, m_oInputArrayMetadata.anBlockSizes, |
461 | 0 | m_anInnerBlockSize, anChunkIndices, |
462 | 0 | abyChunk, nDTSize); |
463 | | |
464 | | // Skip empty (all-nodata) chunks |
465 | 0 | if (IsAllNoData(abyChunk, m_oInputArrayMetadata)) |
466 | 0 | continue; |
467 | | |
468 | | // Encode inner chunk through codec chain (bytes + compressor) |
469 | 0 | if (!m_poCodecSequence->Encode(abyChunk)) |
470 | 0 | { |
471 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
472 | 0 | "ZarrV3CodecShardingIndexed::Encode(): cannot encode " |
473 | 0 | "inner chunk %" PRIu64, |
474 | 0 | static_cast<uint64_t>(iChunk)); |
475 | 0 | return false; |
476 | 0 | } |
477 | | |
478 | 0 | anLocations[iChunk] = {nCurrentOffset, abyChunk.size()}; |
479 | 0 | nCurrentOffset += abyChunk.size(); |
480 | 0 | abyData.insert(abyData.end(), abyChunk.begin(), abyChunk.end()); |
481 | 0 | } |
482 | | |
483 | | // All inner chunks are nodata: signal empty shard via zero-length output |
484 | 0 | if (abyData.empty()) |
485 | 0 | { |
486 | 0 | abyDst.clear(); |
487 | 0 | return true; |
488 | 0 | } |
489 | | |
490 | | // Build index buffer |
491 | 0 | ZarrByteVectorQuickResize abyIndex; |
492 | 0 | const size_t nIndexRawSize = nInnerChunks * sizeof(Location); |
493 | | |
494 | | // If index is at start, data offsets must account for the index prefix. |
495 | | // Encode the index once to determine its encoded size (includes CRC32C |
496 | | // overhead), then re-encode below with the adjusted offsets. |
497 | | // Note: currently unreachable in write mode (creation hardcodes "end"), |
498 | | // but kept for completeness with the Decode() start-index path. |
499 | 0 | if (!m_bIndexLocationAtEnd) |
500 | 0 | { |
501 | | // Encode a temporary copy of the index to determine its encoded size |
502 | 0 | abyIndex.resize(nIndexRawSize); |
503 | 0 | memcpy(abyIndex.data(), anLocations.data(), nIndexRawSize); |
504 | 0 | if (!m_poIndexCodecSequence->Encode(abyIndex)) |
505 | 0 | { |
506 | 0 | CPLError( |
507 | 0 | CE_Failure, CPLE_AppDefined, |
508 | 0 | "ZarrV3CodecShardingIndexed::Encode(): cannot encode index"); |
509 | 0 | return false; |
510 | 0 | } |
511 | 0 | const uint64_t nIndexEncodedSize = abyIndex.size(); |
512 | | // Adjust offsets |
513 | 0 | for (auto &loc : anLocations) |
514 | 0 | { |
515 | 0 | if (loc.nOffset != std::numeric_limits<uint64_t>::max()) |
516 | 0 | loc.nOffset += nIndexEncodedSize; |
517 | 0 | } |
518 | 0 | } |
519 | | |
520 | | // (Re-)encode index with final offsets |
521 | 0 | abyIndex.resize(nIndexRawSize); |
522 | 0 | memcpy(abyIndex.data(), anLocations.data(), nIndexRawSize); |
523 | 0 | if (!m_poIndexCodecSequence->Encode(abyIndex)) |
524 | 0 | { |
525 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
526 | 0 | "ZarrV3CodecShardingIndexed::Encode(): cannot encode index"); |
527 | 0 | return false; |
528 | 0 | } |
529 | | |
530 | | // Assemble output: data + index (or index + data) |
531 | 0 | if (m_bIndexLocationAtEnd) |
532 | 0 | { |
533 | 0 | abyDst = std::move(abyData); |
534 | 0 | abyDst.insert(abyDst.end(), abyIndex.begin(), abyIndex.end()); |
535 | 0 | } |
536 | 0 | else |
537 | 0 | { |
538 | 0 | abyDst = std::move(abyIndex); |
539 | 0 | abyDst.insert(abyDst.end(), abyData.begin(), abyData.end()); |
540 | 0 | } |
541 | |
|
542 | 0 | return true; |
543 | 0 | } |
544 | | |
545 | | /************************************************************************/ |
546 | | /* CopySubArrayIntoLargerOne() */ |
547 | | /************************************************************************/ |
548 | | |
549 | | static void |
550 | | CopySubArrayIntoLargerOne(const ZarrByteVectorQuickResize &abyChunk, |
551 | | const std::vector<size_t> &anInnerBlockSize, |
552 | | const std::vector<size_t> &anInnerBlockIndices, |
553 | | ZarrByteVectorQuickResize &abyDst, |
554 | | const std::vector<size_t> &anDstBlockSize, |
555 | | const size_t nDTSize) |
556 | 0 | { |
557 | 0 | const auto nDims = anInnerBlockSize.size(); |
558 | 0 | CPLAssert(nDims > 0); |
559 | 0 | CPLAssert(nDims == anInnerBlockIndices.size()); |
560 | 0 | CPLAssert(nDims == anDstBlockSize.size()); |
561 | | // +1 just to make some gcc versions not emit -Wnull-dereference false positives |
562 | 0 | std::vector<GByte *> dstPtrStack(nDims + 1); |
563 | 0 | std::vector<size_t> count(nDims + 1); |
564 | 0 | std::vector<size_t> dstStride(nDims + 1); |
565 | |
|
566 | 0 | size_t nDstStride = nDTSize; |
567 | 0 | for (size_t iDim = nDims; iDim > 0;) |
568 | 0 | { |
569 | 0 | --iDim; |
570 | 0 | dstStride[iDim] = nDstStride; |
571 | 0 | nDstStride *= anDstBlockSize[iDim]; |
572 | 0 | } |
573 | |
|
574 | 0 | dstPtrStack[0] = abyDst.data(); |
575 | 0 | for (size_t iDim = 0; iDim < nDims; ++iDim) |
576 | 0 | { |
577 | 0 | CPLAssert((anInnerBlockIndices[iDim] + 1) * anInnerBlockSize[iDim] <= |
578 | 0 | anDstBlockSize[iDim]); |
579 | 0 | dstPtrStack[0] += anInnerBlockIndices[iDim] * anInnerBlockSize[iDim] * |
580 | 0 | dstStride[iDim]; |
581 | 0 | } |
582 | 0 | const GByte *pabySrc = abyChunk.data(); |
583 | |
|
584 | 0 | const size_t nLastDimSize = anInnerBlockSize.back() * nDTSize; |
585 | 0 | size_t dimIdx = 0; |
586 | 0 | lbl_next_depth: |
587 | 0 | if (dimIdx + 1 == nDims) |
588 | 0 | { |
589 | 0 | memcpy(dstPtrStack[dimIdx], pabySrc, nLastDimSize); |
590 | 0 | pabySrc += nLastDimSize; |
591 | 0 | } |
592 | 0 | else |
593 | 0 | { |
594 | 0 | count[dimIdx] = anInnerBlockSize[dimIdx]; |
595 | 0 | while (true) |
596 | 0 | { |
597 | 0 | dimIdx++; |
598 | 0 | dstPtrStack[dimIdx] = dstPtrStack[dimIdx - 1]; |
599 | 0 | goto lbl_next_depth; |
600 | 0 | lbl_return_to_caller: |
601 | 0 | dimIdx--; |
602 | 0 | if (--count[dimIdx] == 0) |
603 | 0 | break; |
604 | 0 | dstPtrStack[dimIdx] += dstStride[dimIdx]; |
605 | 0 | } |
606 | 0 | } |
607 | 0 | if (dimIdx > 0) |
608 | 0 | goto lbl_return_to_caller; |
609 | 0 | } |
610 | | |
611 | | /************************************************************************/ |
612 | | /* FillWithNoData() */ |
613 | | /************************************************************************/ |
614 | | |
615 | | static void FillWithNoData(ZarrByteVectorQuickResize &abyDst, |
616 | | const size_t nCount, |
617 | | const ZarrArrayMetadata &metadata) |
618 | 0 | { |
619 | 0 | const size_t nDTSize = metadata.oElt.nativeSize; |
620 | 0 | if (metadata.abyNoData.empty() || |
621 | 0 | metadata.abyNoData == std::vector<GByte>(nDTSize, 0)) |
622 | 0 | { |
623 | 0 | memset(abyDst.data(), 0, nDTSize * nCount); |
624 | 0 | } |
625 | 0 | else |
626 | 0 | { |
627 | 0 | CPLAssert(metadata.abyNoData.size() == nDTSize); |
628 | 0 | for (size_t i = 0; i < nCount; ++i) |
629 | 0 | { |
630 | 0 | memcpy(abyDst.data() + i * nDTSize, metadata.abyNoData.data(), |
631 | 0 | nDTSize); |
632 | 0 | } |
633 | 0 | } |
634 | 0 | } |
635 | | |
636 | | /************************************************************************/ |
637 | | /* ZarrV3CodecShardingIndexed::Decode() */ |
638 | | /************************************************************************/ |
639 | | |
640 | | bool ZarrV3CodecShardingIndexed::Decode(const ZarrByteVectorQuickResize &abySrc, |
641 | | ZarrByteVectorQuickResize &abyDst) const |
642 | 0 | { |
643 | 0 | size_t nInnerChunks = 1; |
644 | 0 | for (size_t i = 0; i < m_anInnerBlockSize.size(); ++i) |
645 | 0 | { |
646 | 0 | const size_t nCountInnerChunksThisdim = |
647 | 0 | m_oInputArrayMetadata.anBlockSizes[i] / m_anInnerBlockSize[i]; |
648 | 0 | nInnerChunks *= nCountInnerChunksThisdim; |
649 | 0 | } |
650 | |
|
651 | 0 | const size_t nIndexEncodedSize = nInnerChunks * sizeof(Location) + |
652 | 0 | (m_bIndexHasCRC32 ? sizeof(uint32_t) : 0); |
653 | 0 | if (abySrc.size() < nIndexEncodedSize) |
654 | 0 | { |
655 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
656 | 0 | "ZarrV3CodecShardingIndexed::Decode(): input buffer is too " |
657 | 0 | "small to hold the shard index"); |
658 | 0 | return false; |
659 | 0 | } |
660 | 0 | ZarrByteVectorQuickResize abyIndex; |
661 | 0 | if (m_bIndexLocationAtEnd) |
662 | 0 | { |
663 | 0 | abyIndex.insert(abyIndex.end(), |
664 | 0 | abySrc.begin() + (abySrc.size() - nIndexEncodedSize), |
665 | 0 | abySrc.end()); |
666 | 0 | } |
667 | 0 | else |
668 | 0 | { |
669 | 0 | abyIndex.insert(abyIndex.end(), abySrc.begin(), |
670 | 0 | abySrc.end() + nIndexEncodedSize); |
671 | 0 | } |
672 | |
|
673 | 0 | if (!m_poIndexCodecSequence->Decode(abyIndex)) |
674 | 0 | { |
675 | 0 | CPLError( |
676 | 0 | CE_Failure, CPLE_NotSupported, |
677 | 0 | "ZarrV3CodecShardingIndexed::Decode(): cannot decode shard index"); |
678 | 0 | return false; |
679 | 0 | } |
680 | | |
681 | 0 | if (abyIndex.size() != nInnerChunks * sizeof(Location)) |
682 | 0 | { |
683 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
684 | 0 | "ZarrV3CodecShardingIndexed::Decode(): shard index has not " |
685 | 0 | "expected size"); |
686 | 0 | return false; |
687 | 0 | } |
688 | | |
689 | 0 | const Location *panLocations = |
690 | 0 | reinterpret_cast<const Location *>(abyIndex.data()); |
691 | |
|
692 | 0 | ZarrByteVectorQuickResize abyChunk; |
693 | 0 | const auto nDTSize = m_oInputArrayMetadata.oElt.nativeSize; |
694 | 0 | const size_t nExpectedDecodedChunkSize = |
695 | 0 | nDTSize * MultiplyElements(m_anInnerBlockSize); |
696 | 0 | const size_t nDstCount = |
697 | 0 | MultiplyElements(m_oInputArrayMetadata.anBlockSizes); |
698 | |
|
699 | 0 | try |
700 | 0 | { |
701 | 0 | abyDst.resize(nDstCount * nDTSize); |
702 | 0 | } |
703 | 0 | catch (const std::exception &) |
704 | 0 | { |
705 | 0 | CPLError(CE_Failure, CPLE_OutOfMemory, |
706 | 0 | "Cannot allocate memory for decoded shard"); |
707 | 0 | return false; |
708 | 0 | } |
709 | | |
710 | 0 | FillWithNoData(abyDst, nDstCount, m_oInputArrayMetadata); |
711 | |
|
712 | 0 | std::vector<size_t> anChunkIndices(m_anInnerBlockSize.size(), 0); |
713 | 0 | for (size_t iChunk = 0; iChunk < nInnerChunks; ++iChunk) |
714 | 0 | { |
715 | 0 | if (iChunk > 0) |
716 | 0 | { |
717 | | // Update chunk coordinates |
718 | 0 | size_t iDim = m_anInnerBlockSize.size() - 1; |
719 | 0 | while (++anChunkIndices[iDim] == |
720 | 0 | m_oInputArrayMetadata.anBlockSizes[iDim] / |
721 | 0 | m_anInnerBlockSize[iDim]) |
722 | 0 | { |
723 | 0 | anChunkIndices[iDim] = 0; |
724 | 0 | --iDim; |
725 | 0 | } |
726 | 0 | } |
727 | |
|
728 | | #ifdef DEBUG_VERBOSE |
729 | | CPLDebug("ZARR", "Chunk %" PRIu64 ": offset %" PRIu64 ", size %" PRIu64, |
730 | | static_cast<uint64_t>(iChunk), panLocations[iChunk].nOffset, |
731 | | panLocations[iChunk].nSize); |
732 | | #endif |
733 | |
|
734 | 0 | if (panLocations[iChunk].nOffset == |
735 | 0 | std::numeric_limits<uint64_t>::max() && |
736 | 0 | panLocations[iChunk].nSize == std::numeric_limits<uint64_t>::max()) |
737 | 0 | { |
738 | | // Empty chunk |
739 | 0 | continue; |
740 | 0 | } |
741 | | |
742 | 0 | if (panLocations[iChunk].nOffset >= abySrc.size() || |
743 | 0 | panLocations[iChunk].nSize > |
744 | 0 | abySrc.size() - panLocations[iChunk].nOffset) |
745 | 0 | { |
746 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
747 | 0 | "ZarrV3CodecShardingIndexed::Decode(): invalid chunk " |
748 | 0 | "location for chunk %" PRIu64 ": offset=%" PRIu64 |
749 | 0 | ", size=%" PRIu64, |
750 | 0 | static_cast<uint64_t>(iChunk), |
751 | 0 | panLocations[iChunk].nOffset, panLocations[iChunk].nSize); |
752 | 0 | return false; |
753 | 0 | } |
754 | | |
755 | 0 | abyChunk.clear(); |
756 | 0 | abyChunk.insert( |
757 | 0 | abyChunk.end(), |
758 | 0 | abySrc.begin() + static_cast<size_t>(panLocations[iChunk].nOffset), |
759 | 0 | abySrc.begin() + static_cast<size_t>(panLocations[iChunk].nOffset + |
760 | 0 | panLocations[iChunk].nSize)); |
761 | 0 | if (!m_poCodecSequence->Decode(abyChunk)) |
762 | 0 | { |
763 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
764 | 0 | "ZarrV3CodecShardingIndexed::Decode(): cannot decode " |
765 | 0 | "chunk %" PRIu64, |
766 | 0 | static_cast<uint64_t>(iChunk)); |
767 | 0 | return false; |
768 | 0 | } |
769 | | |
770 | 0 | if (abyChunk.size() != nExpectedDecodedChunkSize) |
771 | 0 | { |
772 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
773 | 0 | "ZarrV3CodecShardingIndexed::Decode(): decoded size for " |
774 | 0 | "chunk %" PRIu64 " is %" PRIu64 " whereas %" PRIu64 |
775 | 0 | " is expected", |
776 | 0 | static_cast<uint64_t>(iChunk), |
777 | 0 | static_cast<uint64_t>(abyChunk.size()), |
778 | 0 | static_cast<uint64_t>(nExpectedDecodedChunkSize)); |
779 | 0 | return false; |
780 | 0 | } |
781 | | |
782 | 0 | CopySubArrayIntoLargerOne(abyChunk, m_anInnerBlockSize, anChunkIndices, |
783 | 0 | abyDst, m_oInputArrayMetadata.anBlockSizes, |
784 | 0 | nDTSize); |
785 | 0 | } |
786 | | |
787 | 0 | return true; |
788 | 0 | } |
789 | | |
790 | | /************************************************************************/ |
791 | | /* ZarrV3CodecShardingIndexed::DecodePartial() */ |
792 | | /************************************************************************/ |
793 | | |
794 | | bool ZarrV3CodecShardingIndexed::DecodePartial( |
795 | | VSIVirtualHandle *poFile, const ZarrByteVectorQuickResize & /* abySrc */, |
796 | | ZarrByteVectorQuickResize &abyDst, std::vector<size_t> &anStartIdx, |
797 | | std::vector<size_t> &anCount) |
798 | 0 | { |
799 | 0 | CPLAssert(anStartIdx.size() == m_oInputArrayMetadata.anBlockSizes.size()); |
800 | 0 | CPLAssert(anStartIdx.size() == anCount.size()); |
801 | |
|
802 | 0 | size_t nInnerChunkCount = 1; |
803 | 0 | size_t nInnerChunkIdx = 0; |
804 | 0 | for (size_t i = 0; i < anStartIdx.size(); ++i) |
805 | 0 | { |
806 | 0 | CPLAssert(anStartIdx[i] + anCount[i] <= |
807 | 0 | m_oInputArrayMetadata.anBlockSizes[i]); |
808 | 0 | if ((anStartIdx[i] % m_anInnerBlockSize[i]) != 0 || |
809 | 0 | anCount[i] != m_anInnerBlockSize[i]) |
810 | 0 | { |
811 | | // Should not happen with the current call sites. |
812 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
813 | 0 | "ZarrV3CodecShardingIndexed::DecodePartial() only " |
814 | 0 | "supported on an exact inner chunk"); |
815 | 0 | return false; |
816 | 0 | } |
817 | | |
818 | 0 | const size_t nCountInnerChunksThisDim = |
819 | 0 | m_oInputArrayMetadata.anBlockSizes[i] / m_anInnerBlockSize[i]; |
820 | 0 | nInnerChunkIdx = nInnerChunkIdx * nCountInnerChunksThisDim + |
821 | 0 | anStartIdx[i] / m_anInnerBlockSize[i]; |
822 | 0 | nInnerChunkCount *= nCountInnerChunksThisDim; |
823 | 0 | } |
824 | | |
825 | 0 | abyDst.clear(); |
826 | |
|
827 | 0 | const auto nDTSize = m_oInputArrayMetadata.oElt.nativeSize; |
828 | 0 | const auto nExpectedDecodedChunkSize = nDTSize * MultiplyElements(anCount); |
829 | |
|
830 | 0 | vsi_l_offset nLocationOffset = |
831 | 0 | static_cast<vsi_l_offset>(nInnerChunkIdx) * sizeof(Location); |
832 | 0 | if (m_bIndexLocationAtEnd) |
833 | 0 | { |
834 | 0 | poFile->Seek(0, SEEK_END); |
835 | 0 | const auto nFileSize = poFile->Tell(); |
836 | 0 | vsi_l_offset nIndexSize = |
837 | 0 | static_cast<vsi_l_offset>(nInnerChunkCount) * sizeof(Location); |
838 | 0 | if (m_bIndexHasCRC32) |
839 | 0 | nIndexSize += sizeof(uint32_t); |
840 | 0 | if (nFileSize < nIndexSize) |
841 | 0 | { |
842 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
843 | 0 | "ZarrV3CodecShardingIndexed::DecodePartial(): shard file " |
844 | 0 | "too small"); |
845 | 0 | return false; |
846 | 0 | } |
847 | 0 | nLocationOffset += nFileSize - nIndexSize; |
848 | 0 | } |
849 | | |
850 | 0 | Location loc; |
851 | 0 | if (poFile->Seek(nLocationOffset, SEEK_SET) != 0 || |
852 | 0 | poFile->Read(&loc, 1, sizeof(loc)) != sizeof(loc)) |
853 | 0 | { |
854 | |
|
855 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
856 | 0 | "ZarrV3CodecShardingIndexed::DecodePartial(): " |
857 | 0 | "cannot read index for chunk %" PRIu64, |
858 | 0 | static_cast<uint64_t>(nInnerChunkIdx)); |
859 | 0 | return false; |
860 | 0 | } |
861 | | |
862 | 0 | if (!m_poIndexCodecSequence->GetCodecs().empty() && |
863 | 0 | m_poIndexCodecSequence->GetCodecs().front()->GetName() == |
864 | 0 | ZarrV3CodecBytes::NAME && |
865 | 0 | !m_poIndexCodecSequence->GetCodecs().front()->IsNoOp()) |
866 | 0 | { |
867 | 0 | CPL_SWAP64PTR(&(loc.nOffset)); |
868 | 0 | CPL_SWAP64PTR(&(loc.nSize)); |
869 | 0 | } |
870 | |
|
871 | 0 | if (loc.nOffset == std::numeric_limits<uint64_t>::max() && |
872 | 0 | loc.nSize == std::numeric_limits<uint64_t>::max()) |
873 | 0 | { |
874 | | // Empty chunk |
875 | 0 | try |
876 | 0 | { |
877 | 0 | abyDst.resize(nExpectedDecodedChunkSize); |
878 | 0 | } |
879 | 0 | catch (const std::exception &) |
880 | 0 | { |
881 | 0 | CPLError(CE_Failure, CPLE_OutOfMemory, |
882 | 0 | "Cannot allocate memory for decoded shard"); |
883 | 0 | return false; |
884 | 0 | } |
885 | 0 | FillWithNoData(abyDst, MultiplyElements(anCount), |
886 | 0 | m_oInputArrayMetadata); |
887 | 0 | return true; |
888 | 0 | } |
889 | | |
890 | 0 | constexpr size_t THRESHOLD = 10 * 1024 * 1024; |
891 | 0 | if (loc.nSize > THRESHOLD) |
892 | 0 | { |
893 | | // When the chunk size is above a certain threshold, check it against |
894 | | // the actual file size to avoid excessive memory allocation attempts. |
895 | |
|
896 | 0 | poFile->Seek(0, SEEK_END); |
897 | 0 | const auto nFileSize = poFile->Tell(); |
898 | |
|
899 | 0 | if (loc.nOffset >= nFileSize || loc.nSize > nFileSize - loc.nOffset) |
900 | 0 | { |
901 | 0 | CPLError( |
902 | 0 | CE_Failure, CPLE_NotSupported, |
903 | 0 | "ZarrV3CodecShardingIndexed::DecodePartial(): invalid chunk " |
904 | 0 | "location for chunk %" PRIu64 ": offset=%" PRIu64 |
905 | 0 | ", size=%" PRIu64, |
906 | 0 | static_cast<uint64_t>(nInnerChunkIdx), loc.nOffset, loc.nSize); |
907 | 0 | return false; |
908 | 0 | } |
909 | 0 | } |
910 | | |
911 | | if constexpr (sizeof(size_t) < sizeof(uint64_t)) |
912 | | { |
913 | | #ifndef __COVERITY__ |
914 | | // coverity[result_independent_of_operands] |
915 | | if (loc.nSize > std::numeric_limits<size_t>::max()) |
916 | | { |
917 | | CPLError( |
918 | | CE_Failure, CPLE_NotSupported, |
919 | | "ZarrV3CodecShardingIndexed::DecodePartial(): too large chunk " |
920 | | "size for chunk %" PRIu64 " for this platform: size=%" PRIu64, |
921 | | static_cast<uint64_t>(nInnerChunkIdx), loc.nSize); |
922 | | return false; |
923 | | } |
924 | | #endif |
925 | | } |
926 | |
|
927 | 0 | try |
928 | 0 | { |
929 | 0 | abyDst.resize(static_cast<size_t>(loc.nSize)); |
930 | 0 | } |
931 | 0 | catch (const std::exception &) |
932 | 0 | { |
933 | 0 | CPLError(CE_Failure, CPLE_OutOfMemory, |
934 | 0 | "Cannot allocate memory for decoded shard"); |
935 | 0 | return false; |
936 | 0 | } |
937 | | |
938 | 0 | if (poFile->Seek(loc.nOffset, SEEK_SET) != 0 || |
939 | 0 | poFile->Read(abyDst.data(), 1, abyDst.size()) != abyDst.size()) |
940 | 0 | { |
941 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
942 | 0 | "ZarrV3CodecShardingIndexed::DecodePartial(): cannot read " |
943 | 0 | "data for chunk %" PRIu64 ": offset=%" PRIu64 |
944 | 0 | ", size=%" PRIu64, |
945 | 0 | static_cast<uint64_t>(nInnerChunkIdx), loc.nOffset, loc.nSize); |
946 | 0 | return false; |
947 | 0 | } |
948 | | |
949 | 0 | if (!m_poCodecSequence->Decode(abyDst)) |
950 | 0 | { |
951 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
952 | 0 | "ZarrV3CodecShardingIndexed::DecodePartial(): cannot decode " |
953 | 0 | "chunk %" PRIu64, |
954 | 0 | static_cast<uint64_t>(nInnerChunkIdx)); |
955 | 0 | return false; |
956 | 0 | } |
957 | | |
958 | 0 | if (abyDst.size() != nExpectedDecodedChunkSize) |
959 | 0 | { |
960 | 0 | CPLError( |
961 | 0 | CE_Failure, CPLE_NotSupported, |
962 | 0 | "ZarrV3CodecShardingIndexed::DecodePartial(): decoded size for " |
963 | 0 | "chunk %" PRIu64 " is %" PRIu64 " whereas %" PRIu64 " is expected", |
964 | 0 | static_cast<uint64_t>(nInnerChunkIdx), |
965 | 0 | static_cast<uint64_t>(abyDst.size()), |
966 | 0 | static_cast<uint64_t>(nExpectedDecodedChunkSize)); |
967 | 0 | return false; |
968 | 0 | } |
969 | | |
970 | 0 | return true; |
971 | 0 | } |
972 | | |
973 | | /************************************************************************/ |
974 | | /* ZarrV3CodecShardingIndexed::BatchDecodePartial() */ |
975 | | /************************************************************************/ |
976 | | |
977 | | bool ZarrV3CodecShardingIndexed::BatchDecodePartial( |
978 | | VSIVirtualHandle *poFile, const char *pszFilename, |
979 | | const std::vector<std::pair<std::vector<size_t>, std::vector<size_t>>> |
980 | | &anRequests, |
981 | | std::vector<ZarrByteVectorQuickResize> &aResults) |
982 | 0 | { |
983 | 0 | if (anRequests.empty()) |
984 | 0 | return true; |
985 | | |
986 | 0 | const auto nDTSize = m_oInputArrayMetadata.oElt.nativeSize; |
987 | | |
988 | | // --- Compute inner chunk count and per-request inner chunk indices --- |
989 | 0 | size_t nInnerChunkCount = 1; |
990 | 0 | for (size_t i = 0; i < m_oInputArrayMetadata.anBlockSizes.size(); ++i) |
991 | 0 | { |
992 | 0 | nInnerChunkCount *= |
993 | 0 | m_oInputArrayMetadata.anBlockSizes[i] / m_anInnerBlockSize[i]; |
994 | 0 | } |
995 | | |
996 | | // Determine whether index codec requires byte-swapping |
997 | 0 | const bool bSwapIndex = |
998 | 0 | !m_poIndexCodecSequence->GetCodecs().empty() && |
999 | 0 | m_poIndexCodecSequence->GetCodecs().front()->GetName() == |
1000 | 0 | ZarrV3CodecBytes::NAME && |
1001 | 0 | !m_poIndexCodecSequence->GetCodecs().front()->IsNoOp(); |
1002 | | |
1003 | | // Compute index base offset. For index-at-end, we need the file size. |
1004 | 0 | vsi_l_offset nIndexBaseOffset = 0; |
1005 | 0 | if (m_bIndexLocationAtEnd) |
1006 | 0 | { |
1007 | 0 | poFile->Seek(0, SEEK_END); |
1008 | 0 | const auto nFileSize = poFile->Tell(); |
1009 | 0 | vsi_l_offset nIndexSize = |
1010 | 0 | static_cast<vsi_l_offset>(nInnerChunkCount) * sizeof(Location); |
1011 | 0 | if (m_bIndexHasCRC32) |
1012 | 0 | nIndexSize += sizeof(uint32_t); |
1013 | 0 | if (nFileSize < nIndexSize) |
1014 | 0 | { |
1015 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
1016 | 0 | "BatchDecodePartial: shard file too small"); |
1017 | 0 | return false; |
1018 | 0 | } |
1019 | 0 | nIndexBaseOffset = nFileSize - nIndexSize; |
1020 | 0 | } |
1021 | | |
1022 | | // Build per-request inner chunk indices |
1023 | 0 | std::vector<size_t> anInnerChunkIndices(anRequests.size()); |
1024 | 0 | for (size_t iReq = 0; iReq < anRequests.size(); ++iReq) |
1025 | 0 | { |
1026 | 0 | const auto &anStartIdx = anRequests[iReq].first; |
1027 | 0 | CPLAssert(anStartIdx.size() == |
1028 | 0 | m_oInputArrayMetadata.anBlockSizes.size()); |
1029 | |
|
1030 | 0 | size_t nInnerChunkIdx = 0; |
1031 | 0 | for (size_t i = 0; i < anStartIdx.size(); ++i) |
1032 | 0 | { |
1033 | 0 | const size_t nCountInnerChunksThisDim = |
1034 | 0 | m_oInputArrayMetadata.anBlockSizes[i] / m_anInnerBlockSize[i]; |
1035 | 0 | nInnerChunkIdx = nInnerChunkIdx * nCountInnerChunksThisDim + |
1036 | 0 | anStartIdx[i] / m_anInnerBlockSize[i]; |
1037 | 0 | } |
1038 | 0 | anInnerChunkIndices[iReq] = nInnerChunkIdx; |
1039 | 0 | } |
1040 | | |
1041 | | // --- Pass 1: read index entries, using per-shard cache when possible --- |
1042 | 0 | std::vector<Location> aLocations(anRequests.size()); |
1043 | |
|
1044 | 0 | const size_t nIndexBytes = nInnerChunkCount * sizeof(Location); |
1045 | 0 | const size_t nMaxIndexBytes = |
1046 | 0 | static_cast<size_t>(CPLAtoGIntBig(CPLGetConfigOption( |
1047 | 0 | "GDAL_ZARR_SHARD_INDEX_CACHE_MAX_BYTES", "1048576"))); |
1048 | 0 | if (pszFilename != nullptr && nIndexBytes <= nMaxIndexBytes) |
1049 | 0 | { |
1050 | | // Try to serve from the in-process index cache. |
1051 | | // g_oShardIndexCache serialises its own access via std::mutex. |
1052 | 0 | std::vector<GByte> abyIndexBuf; |
1053 | 0 | bool bCacheHit = g_oShardIndexCache.tryGet(pszFilename, abyIndexBuf); |
1054 | 0 | if (bCacheHit && abyIndexBuf.size() != nIndexBytes) |
1055 | 0 | bCacheHit = false; // stale entry from different layout |
1056 | |
|
1057 | 0 | if (!bCacheHit) |
1058 | 0 | { |
1059 | | // Cache miss: read the full index in one contiguous I/O call. |
1060 | 0 | abyIndexBuf.resize(nIndexBytes); |
1061 | 0 | poFile->Seek(nIndexBaseOffset, SEEK_SET); |
1062 | 0 | if (poFile->Read(abyIndexBuf.data(), nIndexBytes, 1) != 1) |
1063 | 0 | { |
1064 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
1065 | 0 | "BatchDecodePartial: failed to read shard index"); |
1066 | 0 | return false; |
1067 | 0 | } |
1068 | 0 | if (bSwapIndex) |
1069 | 0 | { |
1070 | 0 | for (size_t j = 0; j < nInnerChunkCount; ++j) |
1071 | 0 | { |
1072 | 0 | GByte *p = abyIndexBuf.data() + j * sizeof(Location); |
1073 | 0 | CPL_SWAP64PTR(p); |
1074 | 0 | CPL_SWAP64PTR(p + sizeof(uint64_t)); |
1075 | 0 | } |
1076 | 0 | } |
1077 | 0 | g_oShardIndexCache.insert(pszFilename, abyIndexBuf); |
1078 | 0 | } |
1079 | | |
1080 | 0 | for (size_t i = 0; i < anRequests.size(); ++i) |
1081 | 0 | memcpy(&aLocations[i], |
1082 | 0 | abyIndexBuf.data() + |
1083 | 0 | anInnerChunkIndices[i] * sizeof(Location), |
1084 | 0 | sizeof(Location)); |
1085 | 0 | } |
1086 | 0 | else |
1087 | 0 | { |
1088 | | // Shard index too large for cache or no filename: fall back to |
1089 | | // per-entry ReadMultiRange. |
1090 | 0 | std::vector<vsi_l_offset> anIdxOffsets(anRequests.size()); |
1091 | 0 | std::vector<size_t> anIdxSizes(anRequests.size(), sizeof(Location)); |
1092 | 0 | std::vector<void *> ppIdxData(anRequests.size()); |
1093 | |
|
1094 | 0 | for (size_t i = 0; i < anRequests.size(); ++i) |
1095 | 0 | { |
1096 | 0 | anIdxOffsets[i] = nIndexBaseOffset + static_cast<vsi_l_offset>( |
1097 | 0 | anInnerChunkIndices[i]) * |
1098 | 0 | sizeof(Location); |
1099 | 0 | #ifndef __COVERITY__ |
1100 | 0 | ppIdxData[i] = &aLocations[i]; |
1101 | 0 | #endif |
1102 | 0 | } |
1103 | |
|
1104 | 0 | if (poFile->ReadMultiRange(static_cast<int>(anRequests.size()), |
1105 | 0 | ppIdxData.data(), anIdxOffsets.data(), |
1106 | 0 | anIdxSizes.data()) != 0) |
1107 | 0 | { |
1108 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
1109 | 0 | "BatchDecodePartial: ReadMultiRange() failed for index"); |
1110 | 0 | return false; |
1111 | 0 | } |
1112 | | |
1113 | 0 | if (bSwapIndex) |
1114 | 0 | { |
1115 | 0 | for (auto &loc : aLocations) |
1116 | 0 | { |
1117 | 0 | CPL_SWAP64PTR(&(loc.nOffset)); |
1118 | 0 | CPL_SWAP64PTR(&(loc.nSize)); |
1119 | 0 | } |
1120 | 0 | } |
1121 | 0 | } |
1122 | | |
1123 | | // --- Classify requests: empty chunks vs data chunks --- |
1124 | 0 | aResults.resize(anRequests.size()); |
1125 | |
|
1126 | 0 | struct DataRange |
1127 | 0 | { |
1128 | 0 | size_t nReqIdx; |
1129 | 0 | }; |
1130 | |
|
1131 | 0 | std::vector<DataRange> aDataRanges; |
1132 | 0 | std::vector<vsi_l_offset> anDataOffsets; |
1133 | 0 | std::vector<size_t> anDataSizes; |
1134 | |
|
1135 | 0 | for (size_t iReq = 0; iReq < anRequests.size(); ++iReq) |
1136 | 0 | { |
1137 | 0 | const auto &anCount = anRequests[iReq].second; |
1138 | 0 | const auto nExpectedDecodedChunkSize = |
1139 | 0 | nDTSize * MultiplyElements(anCount); |
1140 | 0 | const Location &loc = aLocations[iReq]; |
1141 | |
|
1142 | 0 | if (loc.nOffset == std::numeric_limits<uint64_t>::max() && |
1143 | 0 | loc.nSize == std::numeric_limits<uint64_t>::max()) |
1144 | 0 | { |
1145 | | // Empty chunk - fill with nodata |
1146 | 0 | try |
1147 | 0 | { |
1148 | 0 | aResults[iReq].resize(nExpectedDecodedChunkSize); |
1149 | 0 | } |
1150 | 0 | catch (const std::exception &) |
1151 | 0 | { |
1152 | 0 | CPLError(CE_Failure, CPLE_OutOfMemory, |
1153 | 0 | "Cannot allocate memory for decoded chunk"); |
1154 | 0 | return false; |
1155 | 0 | } |
1156 | 0 | FillWithNoData(aResults[iReq], MultiplyElements(anCount), |
1157 | 0 | m_oInputArrayMetadata); |
1158 | 0 | continue; |
1159 | 0 | } |
1160 | | |
1161 | | if constexpr (sizeof(size_t) < sizeof(uint64_t)) |
1162 | | { |
1163 | | #ifndef __COVERITY__ |
1164 | | // coverity[result_independent_of_operands] |
1165 | | if (loc.nSize > std::numeric_limits<size_t>::max()) |
1166 | | { |
1167 | | CPLError(CE_Failure, CPLE_NotSupported, |
1168 | | "BatchDecodePartial: too large chunk size"); |
1169 | | return false; |
1170 | | } |
1171 | | #endif |
1172 | | } |
1173 | |
|
1174 | 0 | aDataRanges.push_back({iReq}); |
1175 | 0 | anDataOffsets.push_back(loc.nOffset); |
1176 | 0 | anDataSizes.push_back(static_cast<size_t>(loc.nSize)); |
1177 | 0 | } |
1178 | | |
1179 | 0 | if (aDataRanges.empty()) |
1180 | 0 | return true; |
1181 | | |
1182 | | // Validate against file size (same threshold as DecodePartial) |
1183 | 0 | constexpr size_t THRESHOLD = 10 * 1024 * 1024; |
1184 | 0 | { |
1185 | 0 | uint64_t nAccSize = 0; |
1186 | 0 | for (const auto &nSize : anDataSizes) |
1187 | 0 | { |
1188 | 0 | if (nAccSize > std::numeric_limits<uint64_t>::max() - nSize) |
1189 | 0 | { |
1190 | 0 | nAccSize = std::numeric_limits<uint64_t>::max(); |
1191 | 0 | break; |
1192 | 0 | } |
1193 | 0 | nAccSize += nSize; |
1194 | 0 | } |
1195 | 0 | if (nAccSize > THRESHOLD) |
1196 | 0 | { |
1197 | 0 | poFile->Seek(0, SEEK_END); |
1198 | 0 | const auto nFileSize = poFile->Tell(); |
1199 | 0 | for (size_t i = 0; i < aDataRanges.size(); ++i) |
1200 | 0 | { |
1201 | 0 | if (anDataOffsets[i] >= nFileSize || |
1202 | 0 | anDataSizes[i] > nFileSize - anDataOffsets[i]) |
1203 | 0 | { |
1204 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
1205 | 0 | "BatchDecodePartial: invalid chunk location: " |
1206 | 0 | "offset=%" PRIu64 ", size=%" PRIu64, |
1207 | 0 | static_cast<uint64_t>(anDataOffsets[i]), |
1208 | 0 | static_cast<uint64_t>(anDataSizes[i])); |
1209 | 0 | return false; |
1210 | 0 | } |
1211 | 0 | } |
1212 | 0 | } |
1213 | 0 | } |
1214 | | |
1215 | | // --- Pass 2: ReadMultiRange for data chunks --- |
1216 | 0 | std::vector<ZarrByteVectorQuickResize> aCompressed(aDataRanges.size()); |
1217 | 0 | std::vector<void *> ppData(aDataRanges.size()); |
1218 | |
|
1219 | 0 | for (size_t i = 0; i < aDataRanges.size(); ++i) |
1220 | 0 | { |
1221 | 0 | try |
1222 | 0 | { |
1223 | 0 | aCompressed[i].resize(anDataSizes[i]); |
1224 | 0 | } |
1225 | 0 | catch (const std::exception &) |
1226 | 0 | { |
1227 | 0 | CPLError(CE_Failure, CPLE_OutOfMemory, |
1228 | 0 | "Cannot allocate memory for compressed chunk"); |
1229 | 0 | return false; |
1230 | 0 | } |
1231 | 0 | ppData[i] = aCompressed[i].data(); |
1232 | 0 | } |
1233 | | |
1234 | 0 | CPLDebugOnly("ZARR", |
1235 | 0 | "BatchDecodePartial: ReadMultiRange() with %d data ranges", |
1236 | 0 | static_cast<int>(aDataRanges.size())); |
1237 | |
|
1238 | 0 | if (poFile->ReadMultiRange(static_cast<int>(aDataRanges.size()), |
1239 | 0 | ppData.data(), anDataOffsets.data(), |
1240 | 0 | anDataSizes.data()) != 0) |
1241 | 0 | { |
1242 | 0 | CPLError(CE_Failure, CPLE_AppDefined, |
1243 | 0 | "BatchDecodePartial: ReadMultiRange() failed for data"); |
1244 | 0 | return false; |
1245 | 0 | } |
1246 | | |
1247 | | // --- Decode compressed chunks (parallel when GDAL_NUM_THREADS > 1) --- |
1248 | 0 | const int nMaxThreads = GDALGetNumThreads(); |
1249 | 0 | const int nChunks = static_cast<int>(aDataRanges.size()); |
1250 | 0 | const int nThreads = std::min(std::max(1, nMaxThreads), nChunks); |
1251 | | |
1252 | | // Try parallel decode when multiple threads are available |
1253 | 0 | CPLWorkerThreadPool *wtp = (nThreads > 1 && nChunks > 1) |
1254 | 0 | ? GDALGetGlobalThreadPool(nMaxThreads) |
1255 | 0 | : nullptr; |
1256 | |
|
1257 | 0 | if (!wtp) |
1258 | 0 | { |
1259 | | // Sequential fallback |
1260 | 0 | for (size_t i = 0; i < aDataRanges.size(); ++i) |
1261 | 0 | { |
1262 | 0 | const size_t iReq = aDataRanges[i].nReqIdx; |
1263 | 0 | const auto &anCount = anRequests[iReq].second; |
1264 | 0 | const auto nExpectedDecodedChunkSize = |
1265 | 0 | nDTSize * MultiplyElements(anCount); |
1266 | |
|
1267 | 0 | if (!m_poCodecSequence->Decode(aCompressed[i])) |
1268 | 0 | { |
1269 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
1270 | 0 | "BatchDecodePartial: cannot decode chunk %" PRIu64, |
1271 | 0 | static_cast<uint64_t>(anInnerChunkIndices[iReq])); |
1272 | 0 | return false; |
1273 | 0 | } |
1274 | | |
1275 | 0 | if (aCompressed[i].size() != nExpectedDecodedChunkSize) |
1276 | 0 | { |
1277 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
1278 | 0 | "BatchDecodePartial: decoded size %" PRIu64 |
1279 | 0 | " != expected %" PRIu64, |
1280 | 0 | static_cast<uint64_t>(aCompressed[i].size()), |
1281 | 0 | static_cast<uint64_t>(nExpectedDecodedChunkSize)); |
1282 | 0 | return false; |
1283 | 0 | } |
1284 | | |
1285 | 0 | aResults[iReq] = std::move(aCompressed[i]); |
1286 | 0 | } |
1287 | 0 | return true; |
1288 | 0 | } |
1289 | | |
1290 | 0 | CPLDebugOnly("ZARR", |
1291 | 0 | "BatchDecodePartial: parallel decode with %d threads " |
1292 | 0 | "for %d chunks", |
1293 | 0 | nThreads, nChunks); |
1294 | |
|
1295 | 0 | { |
1296 | 0 | bool bGlobalOK = true; |
1297 | 0 | std::mutex oMutex; |
1298 | | |
1299 | | // Clone codecs per thread on the main thread (Clone() is not |
1300 | | // thread-safe due to JSON object cloning) |
1301 | 0 | std::vector<std::unique_ptr<ZarrV3CodecSequence>> apoCodecs(nThreads); |
1302 | 0 | for (int t = 0; t < nThreads; ++t) |
1303 | 0 | apoCodecs[t] = m_poCodecSequence->Clone(); |
1304 | |
|
1305 | 0 | auto poJobQueue = wtp->CreateJobQueue(); |
1306 | 0 | for (int t = 0; t < nThreads; ++t) |
1307 | 0 | { |
1308 | 0 | const int iFirst = |
1309 | 0 | static_cast<int>(static_cast<int64_t>(t) * nChunks / nThreads); |
1310 | 0 | const int iEnd = static_cast<int>(static_cast<int64_t>(t + 1) * |
1311 | 0 | nChunks / nThreads); |
1312 | |
|
1313 | 0 | poJobQueue->SubmitJob( |
1314 | 0 | [iFirst, iEnd, t, &aDataRanges, &anRequests, &aCompressed, |
1315 | 0 | &aResults, &apoCodecs, &bGlobalOK, &oMutex, nDTSize]() |
1316 | 0 | { |
1317 | 0 | for (int i = iFirst; i < iEnd; ++i) |
1318 | 0 | { |
1319 | 0 | { |
1320 | 0 | std::lock_guard<std::mutex> oLock(oMutex); |
1321 | 0 | if (!bGlobalOK) |
1322 | 0 | return; |
1323 | 0 | } |
1324 | | |
1325 | 0 | const size_t iReq = aDataRanges[i].nReqIdx; |
1326 | 0 | const auto &anCount = anRequests[iReq].second; |
1327 | 0 | const auto nExpected = |
1328 | 0 | nDTSize * MultiplyElements(anCount); |
1329 | |
|
1330 | 0 | if (!apoCodecs[t]->Decode(aCompressed[i]) || |
1331 | 0 | aCompressed[i].size() != nExpected) |
1332 | 0 | { |
1333 | 0 | std::lock_guard<std::mutex> oLock(oMutex); |
1334 | 0 | bGlobalOK = false; |
1335 | 0 | return; |
1336 | 0 | } |
1337 | | |
1338 | | // Each job writes to a unique iReq slot - no lock |
1339 | 0 | aResults[iReq] = std::move(aCompressed[i]); |
1340 | 0 | } |
1341 | 0 | }); |
1342 | 0 | } |
1343 | 0 | poJobQueue->WaitCompletion(); |
1344 | |
|
1345 | 0 | if (!bGlobalOK) |
1346 | 0 | { |
1347 | 0 | CPLError(CE_Failure, CPLE_NotSupported, |
1348 | 0 | "BatchDecodePartial: parallel decode failed"); |
1349 | 0 | return false; |
1350 | 0 | } |
1351 | 0 | } |
1352 | | |
1353 | 0 | return true; |
1354 | 0 | } |
1355 | | |
1356 | | /************************************************************************/ |
1357 | | /* ZarrV3CodecShardingIndexed::GetInnerMostBlockSize() */ |
1358 | | /************************************************************************/ |
1359 | | |
1360 | | std::vector<size_t> ZarrV3CodecShardingIndexed::GetInnerMostBlockSize( |
1361 | | const std::vector<size_t> &) const |
1362 | 0 | { |
1363 | 0 | return m_anInnerBlockSize; |
1364 | | // TODO if we one day properly support nested sharding |
1365 | | // return m_poCodecSequence->GetInnerMostBlockSize(m_anInnerBlockSize); |
1366 | 0 | } |