Line data Source code
1 : /******************************************************************************
2 : *
3 : * Project: GDAL
4 : * Purpose: Zarr driver. Virtual file system for
5 : * https://fsspec.github.io/kerchunk/spec.html#parquet-references
6 : * Author: Even Rouault <even dot rouault at spatialys.com>
7 : *
8 : ******************************************************************************
9 : * Copyright (c) 2025, Even Rouault <even dot rouault at spatialys.com>
10 : *
11 : * SPDX-License-Identifier: MIT
12 : ****************************************************************************/
13 :
14 : #include "vsikerchunk.h"
15 :
16 : #include "cpl_json.h"
17 : #include "cpl_mem_cache.h"
18 : #include "cpl_vsi_virtual.h"
19 :
20 : #include "gdal_priv.h"
21 : #include "ogrsf_frmts.h"
22 :
23 : #include <algorithm>
24 : #include <cinttypes>
25 : #include <functional>
26 : #include <limits>
27 : #include <mutex>
28 : #include <set>
29 : #include <utility>
30 :
31 : extern "C" int CPL_DLL GDALIsInGlobalDestructor();
32 :
33 : /************************************************************************/
34 : /* VSIZarrArrayInfo */
35 : /************************************************************************/
36 :
37 : struct VSIZarrArrayInfo
38 : {
39 : std::vector<uint64_t> anChunkCount{};
40 : };
41 :
42 : /************************************************************************/
43 : /* VSIKerchunkParquetRefFile */
44 : /************************************************************************/
45 :
46 : struct VSIKerchunkParquetRefFile
47 : {
48 : int m_nRecordSize = 0;
49 : std::map<std::string, std::vector<GByte>> m_oMapKeys{};
50 : std::map<std::string, VSIZarrArrayInfo> m_oMapArrayInfo{};
51 : };
52 :
53 : /************************************************************************/
54 : /* VSIKerchunkParquetRefFileSystem */
55 : /************************************************************************/
56 :
57 : class VSIKerchunkParquetRefFileSystem final : public VSIFilesystemHandler
58 : {
59 : public:
60 1653 : VSIKerchunkParquetRefFileSystem()
61 1653 : {
62 1653 : IsFileSystemInstantiated() = true;
63 1653 : }
64 :
65 : ~VSIKerchunkParquetRefFileSystem();
66 :
67 4427 : static bool &IsFileSystemInstantiated()
68 : {
69 : static bool bIsFileSystemInstantiated = false;
70 4427 : return bIsFileSystemInstantiated;
71 : }
72 :
73 : VSIVirtualHandle *Open(const char *pszFilename, const char *pszAccess,
74 : bool bSetError, CSLConstList papszOptions) override;
75 :
76 : int Stat(const char *pszFilename, VSIStatBufL *pStatBuf,
77 : int nFlags) override;
78 :
79 : char **ReadDirEx(const char *pszDirname, int nMaxFiles) override;
80 :
81 : void CleanCache();
82 :
83 : private:
84 : lru11::Cache<std::string, std::shared_ptr<VSIKerchunkParquetRefFile>,
85 : std::mutex>
86 : m_oCache{};
87 :
88 : std::mutex m_oParquetCacheMutex{};
89 : lru11::Cache<std::string, std::shared_ptr<GDALDataset>> *m_poParquetCache{};
90 :
91 : static std::pair<std::string, std::string>
92 : SplitFilename(const char *pszFilename);
93 :
94 : std::shared_ptr<VSIKerchunkParquetRefFile>
95 : Load(const std::string &osRootFilename);
96 :
97 : struct ChunkInfo
98 : {
99 : std::string osParquetFileDirectory{};
100 : std::unique_ptr<OGRFeature> poFeature{};
101 : int iPathField = -1;
102 : int iOffsetField = -1;
103 : int iSizeField = -1;
104 : int iRawField = -1;
105 : };
106 :
107 : ChunkInfo
108 : GetChunkInfo(const std::string &osRootFilename,
109 : const std::shared_ptr<VSIKerchunkParquetRefFile> &refFile,
110 : const std::string &osKey);
111 :
112 : CPL_DISALLOW_COPY_ASSIGN(VSIKerchunkParquetRefFileSystem)
113 : };
114 :
115 : /************************************************************************/
116 : /* ~VSIKerchunkParquetRefFileSystem() */
117 : /************************************************************************/
118 :
119 2242 : VSIKerchunkParquetRefFileSystem::~VSIKerchunkParquetRefFileSystem()
120 : {
121 1121 : CleanCache();
122 1121 : IsFileSystemInstantiated() = false;
123 2242 : }
124 :
125 : /************************************************************************/
126 : /* VSIKerchunkParquetRefFileSystem::CleanCache() */
127 : /************************************************************************/
128 :
129 2176 : void VSIKerchunkParquetRefFileSystem::CleanCache()
130 : {
131 : // If we are in the unloading of the library do not try to close
132 : // datasets to avoid crashes and prefer leaking memory...
133 2176 : if (!GDALIsInGlobalDestructor())
134 : {
135 3690 : std::lock_guard<std::mutex> oLock(m_oParquetCacheMutex);
136 1845 : if (m_poParquetCache)
137 : {
138 8 : m_poParquetCache->clear();
139 8 : delete m_poParquetCache;
140 8 : m_poParquetCache = nullptr;
141 : }
142 : }
143 2176 : }
144 :
145 : /************************************************************************/
146 : /* VSIKerchunkParquetRefFileSystem::SplitFilename() */
147 : /************************************************************************/
148 :
149 : /*static*/
150 : std::pair<std::string, std::string>
151 213 : VSIKerchunkParquetRefFileSystem::SplitFilename(const char *pszFilename)
152 : {
153 213 : if (!STARTS_WITH(pszFilename, PARQUET_REF_FS_PREFIX))
154 0 : return {std::string(), std::string()};
155 :
156 426 : std::string osRootFilename;
157 :
158 213 : pszFilename += strlen(PARQUET_REF_FS_PREFIX);
159 :
160 213 : if (*pszFilename == '{')
161 : {
162 : // Parse /vsikerchunk_parquet_ref/{/path/to/some/parquet_root}[key]
163 199 : int nLevel = 1;
164 199 : ++pszFilename;
165 12632 : for (; *pszFilename; ++pszFilename)
166 : {
167 12622 : if (*pszFilename == '{')
168 : {
169 0 : ++nLevel;
170 : }
171 12622 : else if (*pszFilename == '}')
172 : {
173 189 : --nLevel;
174 189 : if (nLevel == 0)
175 : {
176 189 : ++pszFilename;
177 189 : break;
178 : }
179 : }
180 12433 : osRootFilename += *pszFilename;
181 : }
182 199 : if (nLevel != 0)
183 : {
184 10 : CPLError(CE_Failure, CPLE_AppDefined,
185 : "Invalid %s syntax: should be "
186 : "%s{/path/to/some/file}[/optional_key]",
187 : PARQUET_REF_FS_PREFIX, PARQUET_REF_FS_PREFIX);
188 10 : return {std::string(), std::string()};
189 : }
190 :
191 : return {osRootFilename,
192 378 : *pszFilename == '/' ? pszFilename + 1 : pszFilename};
193 : }
194 : else
195 : {
196 14 : CPLError(CE_Failure, CPLE_AppDefined,
197 : "Invalid %s syntax: should be "
198 : "%s{/path/to/root/dir}[/optional_key]",
199 : PARQUET_REF_FS_PREFIX, PARQUET_REF_FS_PREFIX);
200 14 : return {std::string(), std::string()};
201 : }
202 : }
203 :
204 : /************************************************************************/
205 : /* VSIKerchunkParquetRefFileSystem::Load() */
206 : /************************************************************************/
207 :
208 : std::shared_ptr<VSIKerchunkParquetRefFile>
209 189 : VSIKerchunkParquetRefFileSystem::Load(const std::string &osRootFilename)
210 : {
211 189 : std::shared_ptr<VSIKerchunkParquetRefFile> refFile;
212 189 : if (m_oCache.tryGet(osRootFilename, refFile))
213 156 : return refFile;
214 :
215 66 : CPLJSONDocument oDoc;
216 :
217 : const std::string osZMetataFilename =
218 66 : CPLFormFilenameSafe(osRootFilename.c_str(), ".zmetadata", nullptr);
219 33 : if (!oDoc.Load(osZMetataFilename))
220 : {
221 11 : CPLError(CE_Failure, CPLE_AppDefined,
222 : "VSIKerchunkParquetRefFileSystem: cannot open %s",
223 : osZMetataFilename.c_str());
224 11 : return nullptr;
225 : }
226 :
227 44 : const auto oRoot = oDoc.GetRoot();
228 66 : const auto oRecordSize = oRoot.GetObj("record_size");
229 43 : if (!oRecordSize.IsValid() ||
230 21 : oRecordSize.GetType() != CPLJSONObject::Type::Integer)
231 : {
232 2 : CPLError(CE_Failure, CPLE_AppDefined,
233 : "VSIKerchunkParquetRefFileSystem: key 'record_size' missing "
234 : "or not of type integer");
235 2 : return nullptr;
236 : }
237 :
238 60 : const auto oMetadata = oRoot.GetObj("metadata");
239 39 : if (!oMetadata.IsValid() ||
240 19 : oMetadata.GetType() != CPLJSONObject::Type::Object)
241 : {
242 2 : CPLError(CE_Failure, CPLE_AppDefined,
243 : "VSIKerchunkParquetRefFileSystem: key 'metadata' missing "
244 : "or not of type dict");
245 2 : return nullptr;
246 : }
247 :
248 18 : refFile = std::make_shared<VSIKerchunkParquetRefFile>();
249 18 : refFile->m_nRecordSize = oRecordSize.ToInteger();
250 18 : if (refFile->m_nRecordSize < 1)
251 : {
252 1 : CPLError(CE_Failure, CPLE_AppDefined,
253 : "VSIKerchunkParquetRefFileSystem: Invalid 'record_size'");
254 1 : return nullptr;
255 : }
256 :
257 61 : for (const auto &oEntry : oMetadata.GetChildren())
258 : {
259 52 : const std::string osKeyName = oEntry.GetName();
260 52 : if (oEntry.GetType() == CPLJSONObject::Type::Object)
261 : {
262 : const std::string osSerialized =
263 51 : oEntry.Format(CPLJSONObject::PrettyFormat::Plain);
264 51 : std::vector<GByte> abyValue;
265 : abyValue.insert(
266 51 : abyValue.end(),
267 51 : reinterpret_cast<const GByte *>(osSerialized.data()),
268 51 : reinterpret_cast<const GByte *>(osSerialized.data()) +
269 102 : osSerialized.size());
270 :
271 51 : refFile->m_oMapKeys[osKeyName] = std::move(abyValue);
272 :
273 51 : if (cpl::ends_with(osKeyName, "/.zarray"))
274 : {
275 32 : const auto oShape = oEntry.GetArray("shape");
276 32 : const auto oChunks = oEntry.GetArray("chunks");
277 16 : if (!oShape.IsValid())
278 : {
279 1 : CPLError(CE_Failure, CPLE_AppDefined,
280 : "VSIKerchunkParquetRefFileSystem: "
281 : "missing 'shape' entry for key '%s'",
282 : osKeyName.c_str());
283 1 : return nullptr;
284 : }
285 15 : else if (!oChunks.IsValid())
286 : {
287 1 : CPLError(CE_Failure, CPLE_AppDefined,
288 : "VSIKerchunkParquetRefFileSystem: "
289 : "missing 'chunks' entry for key '%s'",
290 : osKeyName.c_str());
291 1 : return nullptr;
292 : }
293 14 : else if (oShape.Size() != oChunks.Size())
294 : {
295 1 : CPLError(CE_Failure, CPLE_AppDefined,
296 : "VSIKerchunkParquetRefFileSystem: "
297 : "'shape' and 'chunks' entries have not the same "
298 : "number of values for key '%s'",
299 : osKeyName.c_str());
300 1 : return nullptr;
301 : }
302 13 : else if (oShape.Size() > 32)
303 : {
304 1 : CPLError(CE_Failure, CPLE_AppDefined,
305 : "VSIKerchunkParquetRefFileSystem: "
306 : "'shape' has too many dimensions for key '%s'",
307 : osKeyName.c_str());
308 1 : return nullptr;
309 : }
310 : else
311 : {
312 12 : VSIZarrArrayInfo arrayInfo;
313 12 : uint64_t nTotalChunks = 1;
314 19 : for (int i = 0; i < oShape.Size(); ++i)
315 : {
316 10 : const uint64_t nSize = oShape[i].ToLong();
317 10 : const uint64_t nChunkSize = oChunks[i].ToLong();
318 10 : if (nSize == 0)
319 : {
320 1 : CPLError(CE_Failure, CPLE_AppDefined,
321 : "VSIKerchunkParquetRefFileSystem: "
322 : "shape[%d]=0 in "
323 : "array definition for key '%s'",
324 : i, osKeyName.c_str());
325 3 : return nullptr;
326 : }
327 9 : else if (nChunkSize == 0)
328 : {
329 1 : CPLError(CE_Failure, CPLE_AppDefined,
330 : "VSIKerchunkParquetRefFileSystem: "
331 : "chunks[%d]=0 in "
332 : "array definition for key '%s'",
333 : i, osKeyName.c_str());
334 1 : return nullptr;
335 : }
336 8 : const auto nChunkCount =
337 8 : DIV_ROUND_UP(nSize, nChunkSize);
338 8 : if (nChunkCount >
339 8 : std::numeric_limits<uint64_t>::max() / nTotalChunks)
340 : {
341 1 : CPLError(
342 : CE_Failure, CPLE_AppDefined,
343 : "VSIKerchunkParquetRefFileSystem: "
344 : "product(shape[]) > UINT64_MAX for key '%s'",
345 : osKeyName.c_str());
346 1 : return nullptr;
347 : }
348 7 : nTotalChunks *= nChunkCount;
349 7 : arrayInfo.anChunkCount.push_back(nChunkCount);
350 : }
351 : const std::string osArrayDir = osKeyName.substr(
352 18 : 0, osKeyName.size() - strlen("/.zarray"));
353 9 : refFile->m_oMapArrayInfo[osArrayDir] = std::move(arrayInfo);
354 : }
355 : }
356 : }
357 : else
358 : {
359 1 : CPLError(CE_Failure, CPLE_AppDefined,
360 : "VSIKerchunkParquetRefFileSystem: invalid value type for "
361 : "key '%s'",
362 : osKeyName.c_str());
363 1 : return nullptr;
364 : }
365 : }
366 :
367 9 : m_oCache.insert(osRootFilename, refFile);
368 9 : return refFile;
369 : }
370 :
371 : /************************************************************************/
372 : /* VSIKerchunkParquetRefFileSystem::GetChunkInfo() */
373 : /************************************************************************/
374 :
375 : VSIKerchunkParquetRefFileSystem::ChunkInfo
376 78 : VSIKerchunkParquetRefFileSystem::GetChunkInfo(
377 : const std::string &osRootFilename,
378 : const std::shared_ptr<VSIKerchunkParquetRefFile> &refFile,
379 : const std::string &osKey)
380 : {
381 78 : ChunkInfo info;
382 :
383 156 : const std::string osArrayPath = CPLGetPathSafe(osKey.c_str());
384 78 : const auto oIterArray = refFile->m_oMapArrayInfo.find(osArrayPath);
385 156 : const std::string osIndices = CPLGetFilename(osKey.c_str());
386 121 : if (oIterArray != refFile->m_oMapArrayInfo.end() && !osIndices.empty() &&
387 121 : osIndices[0] >= '0' && osIndices[0] <= '9')
388 : {
389 29 : const auto &oArrayInfo = oIterArray->second;
390 : const CPLStringList aosIndices(
391 29 : CSLTokenizeString2(osIndices.c_str(), ".", 0));
392 29 : if ((static_cast<size_t>(aosIndices.size()) ==
393 39 : oArrayInfo.anChunkCount.size()) ||
394 20 : (aosIndices.size() == 1 && strcmp(aosIndices[0], "0") == 0 &&
395 10 : oArrayInfo.anChunkCount.empty()))
396 : {
397 29 : std::vector<uint64_t> anIndices;
398 55 : for (size_t i = 0; i < oArrayInfo.anChunkCount.size(); ++i)
399 : {
400 30 : char *endptr = nullptr;
401 30 : anIndices.push_back(std::strtoull(aosIndices[i], &endptr, 10));
402 30 : if (aosIndices[i][0] == '-' ||
403 58 : endptr != aosIndices[i] + strlen(aosIndices[i]) ||
404 28 : anIndices[i] >= oArrayInfo.anChunkCount[i])
405 : {
406 4 : return info;
407 : }
408 : }
409 :
410 25 : uint64_t nLinearIndex = 0;
411 25 : uint64_t nMulFactor = 1;
412 48 : for (size_t i = anIndices.size(); i > 0;)
413 : {
414 23 : --i;
415 23 : nLinearIndex += anIndices[i] * nMulFactor;
416 23 : nMulFactor *= oArrayInfo.anChunkCount[i];
417 : }
418 :
419 25 : CPLDebugOnly("VSIKerchunkParquetRefFileSystem",
420 : "Linear chunk index %" PRIu64, nLinearIndex);
421 :
422 25 : const uint64_t nParquetIdx = nLinearIndex / refFile->m_nRecordSize;
423 : const int nIdxInParquet =
424 25 : static_cast<int>(nLinearIndex % refFile->m_nRecordSize);
425 :
426 : const std::string osParquetFilename = CPLFormFilenameSafe(
427 25 : CPLFormFilenameSafe(osRootFilename.c_str(), osArrayPath.c_str(),
428 : nullptr)
429 : .c_str(),
430 75 : CPLSPrintf("refs.%" PRIu64 ".parq", nParquetIdx), nullptr);
431 25 : CPLDebugOnly("VSIKerchunkParquetRefFileSystem",
432 : "Looking for entry %d in Parquet file %s",
433 : nIdxInParquet, osParquetFilename.c_str());
434 :
435 50 : std::lock_guard<std::mutex> oLock(m_oParquetCacheMutex);
436 25 : std::shared_ptr<GDALDataset> poDS;
437 25 : if (!m_poParquetCache)
438 : {
439 8 : m_poParquetCache = std::make_unique<lru11::Cache<
440 16 : std::string, std::shared_ptr<GDALDataset>>>()
441 8 : .release();
442 : }
443 25 : if (!m_poParquetCache->tryGet(osParquetFilename, poDS))
444 : {
445 14 : const char *const apszAllowedDrivers[] = {"PARQUET", "ADBC",
446 : nullptr};
447 : CPLConfigOptionSetter oSetter(
448 28 : "OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL", "NO", false);
449 14 : poDS.reset(
450 : GDALDataset::Open(osParquetFilename.c_str(),
451 : GDAL_OF_VECTOR | GDAL_OF_VERBOSE_ERROR,
452 : apszAllowedDrivers, nullptr, nullptr));
453 14 : if (poDS)
454 14 : m_poParquetCache->insert(osParquetFilename, poDS);
455 : }
456 :
457 25 : if (poDS && poDS->GetLayerCount() == 1)
458 : {
459 48 : const auto IsIntOrInt64 = [](OGRFieldType eType)
460 48 : { return eType == OFTInteger || eType == OFTInteger64; };
461 25 : auto poLayer = poDS->GetLayer(0);
462 25 : const auto poDefn = poLayer->GetLayerDefn();
463 25 : info.iPathField = poDefn->GetFieldIndex("path");
464 25 : info.iOffsetField = poDefn->GetFieldIndex("offset");
465 25 : info.iSizeField = poDefn->GetFieldIndex("size");
466 25 : info.iRawField = poDefn->GetFieldIndex("raw");
467 25 : if (info.iPathField >= 0 && info.iOffsetField >= 0 &&
468 49 : info.iSizeField >= 0 && info.iRawField >= 0 &&
469 24 : poDefn->GetFieldDefn(info.iPathField)->GetType() ==
470 24 : OFTString &&
471 24 : IsIntOrInt64(
472 48 : poDefn->GetFieldDefn(info.iOffsetField)->GetType()) &&
473 24 : IsIntOrInt64(
474 74 : poDefn->GetFieldDefn(info.iSizeField)->GetType()) &&
475 24 : poDefn->GetFieldDefn(info.iRawField)->GetType() ==
476 : OFTBinary)
477 : {
478 : info.osParquetFileDirectory =
479 24 : CPLGetPathSafe(osParquetFilename.c_str());
480 24 : info.poFeature.reset(poLayer->GetFeature(nIdxInParquet));
481 : }
482 : else
483 : {
484 1 : CPLError(CE_Failure, CPLE_AppDefined,
485 : "%s has an unexpected field structure",
486 : osParquetFilename.c_str());
487 : }
488 : }
489 : }
490 : }
491 74 : return info;
492 : }
493 :
494 : /************************************************************************/
495 : /* VSIKerchunkParquetRefFileSystem::Open() */
496 : /************************************************************************/
497 :
498 63 : VSIVirtualHandle *VSIKerchunkParquetRefFileSystem::Open(
499 : const char *pszFilename, const char *pszAccess, bool /* bSetError */,
500 : CSLConstList /* papszOptions */)
501 : {
502 63 : CPLDebugOnly("VSIKerchunkParquetRefFileSystem", "Open(%s)", pszFilename);
503 63 : if (strcmp(pszAccess, "r") != 0 && strcmp(pszAccess, "rb") != 0)
504 0 : return nullptr;
505 :
506 126 : const auto [osRootFilename, osKey] = SplitFilename(pszFilename);
507 63 : if (osRootFilename.empty())
508 6 : return nullptr;
509 :
510 114 : const auto refFile = Load(osRootFilename);
511 57 : if (!refFile)
512 3 : return nullptr;
513 :
514 54 : const auto oIter = refFile->m_oMapKeys.find(osKey);
515 54 : if (oIter == refFile->m_oMapKeys.end())
516 : {
517 48 : const auto info = GetChunkInfo(osRootFilename, refFile, osKey);
518 24 : if (info.poFeature)
519 : {
520 13 : if (info.poFeature->IsFieldSetAndNotNull(info.iRawField))
521 : {
522 1 : auto psField = info.poFeature->GetRawFieldRef(info.iRawField);
523 : // Borrow binary data to feature
524 1 : GByte *abyData = psField->Binary.paData;
525 1 : int nSize = psField->Binary.nCount;
526 1 : psField->Binary.paData = nullptr;
527 1 : psField->Binary.nCount = 0;
528 : // and transmit its ownership to the VSIMem file
529 1 : return VSIFileFromMemBuffer(nullptr, abyData, nSize,
530 1 : /* bTakeOwnership = */ true);
531 : }
532 : else
533 : {
534 : const uint64_t nOffset =
535 12 : info.poFeature->GetFieldAsInteger64(info.iOffsetField);
536 : const int nSize =
537 12 : info.poFeature->GetFieldAsInteger(info.iSizeField);
538 :
539 : std::string osVSIPath = VSIKerchunkMorphURIToVSIPath(
540 12 : info.poFeature->GetFieldAsString(info.iPathField),
541 36 : info.osParquetFileDirectory);
542 12 : if (osVSIPath.empty())
543 0 : return nullptr;
544 :
545 : const std::string osPath =
546 : nSize ? CPLSPrintf("/vsisubfile/%" PRIu64 "_%u,%s", nOffset,
547 : nSize, osVSIPath.c_str())
548 24 : : std::move(osVSIPath);
549 12 : CPLDebugOnly("VSIKerchunkParquetRefFileSystem", "Opening %s",
550 : osPath.c_str());
551 : CPLConfigOptionSetter oSetter("GDAL_DISABLE_READDIR_ON_OPEN",
552 24 : "EMPTY_DIR", false);
553 12 : return VSIFOpenL(osPath.c_str(), "rb");
554 : }
555 : }
556 :
557 11 : return nullptr;
558 : }
559 :
560 30 : const auto &abyValue = oIter->second;
561 30 : return VSIFileFromMemBuffer(nullptr, const_cast<GByte *>(abyValue.data()),
562 60 : abyValue.size(), /* bTakeOwnership = */ false);
563 : }
564 :
565 : /************************************************************************/
566 : /* VSIKerchunkParquetRefFileSystem::Stat() */
567 : /************************************************************************/
568 :
569 123 : int VSIKerchunkParquetRefFileSystem::Stat(const char *pszFilename,
570 : VSIStatBufL *pStatBuf, int nFlags)
571 : {
572 123 : CPLDebugOnly("VSIKerchunkParquetRefFileSystem", "Stat(%s)", pszFilename);
573 123 : memset(pStatBuf, 0, sizeof(VSIStatBufL));
574 :
575 246 : const auto [osRootFilename, osKey] = SplitFilename(pszFilename);
576 123 : if (osRootFilename.empty())
577 12 : return -1;
578 :
579 222 : const auto refFile = Load(osRootFilename);
580 111 : if (!refFile)
581 18 : return -1;
582 :
583 93 : if (osKey.empty())
584 : {
585 4 : pStatBuf->st_mode = S_IFDIR;
586 4 : return 0;
587 : }
588 :
589 89 : const auto oIter = refFile->m_oMapKeys.find(osKey);
590 89 : if (oIter == refFile->m_oMapKeys.end())
591 : {
592 108 : const auto info = GetChunkInfo(osRootFilename, refFile, osKey);
593 54 : if (info.poFeature)
594 : {
595 11 : if (info.poFeature->IsFieldSetAndNotNull(info.iRawField))
596 : {
597 0 : int nSize = 0;
598 0 : info.poFeature->GetFieldAsBinary(info.iRawField, &nSize);
599 0 : pStatBuf->st_size = nSize;
600 : }
601 : else
602 : {
603 11 : pStatBuf->st_size =
604 11 : info.poFeature->GetFieldAsInteger64(info.iSizeField);
605 11 : if (pStatBuf->st_size == 0)
606 : {
607 : const std::string osVSIPath = VSIKerchunkMorphURIToVSIPath(
608 5 : info.poFeature->GetFieldAsString(info.iPathField),
609 15 : info.osParquetFileDirectory);
610 5 : if (osVSIPath.empty())
611 0 : return -1;
612 5 : return VSIStatExL(osVSIPath.c_str(), pStatBuf, nFlags);
613 : }
614 : }
615 6 : pStatBuf->st_mode = S_IFREG;
616 6 : return 0;
617 : }
618 :
619 129 : if (cpl::contains(refFile->m_oMapKeys, osKey + "/.zgroup") ||
620 86 : cpl::contains(refFile->m_oMapKeys, osKey + "/.zarray"))
621 : {
622 7 : pStatBuf->st_mode = S_IFDIR;
623 7 : return 0;
624 : }
625 :
626 36 : return -1;
627 : }
628 :
629 35 : const auto &abyValue = oIter->second;
630 35 : pStatBuf->st_size = abyValue.size();
631 35 : pStatBuf->st_mode = S_IFREG;
632 :
633 35 : return 0;
634 : }
635 :
636 : /************************************************************************/
637 : /* VSIKerchunkParquetRefFileSystem::ReadDirEx() */
638 : /************************************************************************/
639 :
640 27 : char **VSIKerchunkParquetRefFileSystem::ReadDirEx(const char *pszDirname,
641 : int nMaxFiles)
642 : {
643 27 : CPLDebugOnly("VSIKerchunkParquetRefFileSystem", "ReadDir(%s)", pszDirname);
644 :
645 54 : const auto [osRootFilename, osAskedKey] = SplitFilename(pszDirname);
646 27 : if (osRootFilename.empty())
647 6 : return nullptr;
648 :
649 42 : const auto refFile = Load(osRootFilename);
650 21 : if (!refFile)
651 3 : return nullptr;
652 :
653 36 : std::set<std::string> set;
654 90 : for (const auto &[key, value] : refFile->m_oMapKeys)
655 : {
656 72 : if (osAskedKey.empty())
657 : {
658 28 : const auto nPos = key.find('/');
659 28 : if (nPos == std::string::npos)
660 14 : set.insert(key);
661 : else
662 14 : set.insert(key.substr(0, nPos));
663 : }
664 44 : else if (key.size() > osAskedKey.size() &&
665 60 : cpl::starts_with(key, osAskedKey) &&
666 16 : key[osAskedKey.size()] == '/')
667 : {
668 32 : std::string subKey = key.substr(osAskedKey.size() + 1);
669 16 : const auto nPos = subKey.find('/');
670 16 : if (nPos == std::string::npos)
671 16 : set.insert(std::move(subKey));
672 : else
673 0 : set.insert(subKey.substr(0, nPos));
674 : }
675 : }
676 :
677 36 : CPLStringList aosRet;
678 55 : for (const std::string &v : set)
679 : {
680 : // CPLDebugOnly("VSIKerchunkParquetRefFileSystem", ".. %s", v.c_str());
681 37 : aosRet.AddString(v.c_str());
682 : }
683 :
684 : // Synthetize file names for x.y.z chunks
685 18 : const auto oIterArray = refFile->m_oMapArrayInfo.find(osAskedKey);
686 18 : if (oIterArray != refFile->m_oMapArrayInfo.end())
687 : {
688 8 : const auto &oArrayInfo = oIterArray->second;
689 8 : if (oArrayInfo.anChunkCount.empty())
690 : {
691 4 : aosRet.AddString("0");
692 : }
693 : else
694 : {
695 8 : std::string osCurElt;
696 8 : std::function<bool(size_t)> Enumerate;
697 4 : if (nMaxFiles <= 0)
698 3 : nMaxFiles = 100 * 1024 * 1024;
699 :
700 8 : Enumerate = [nMaxFiles, &aosRet, &oArrayInfo, &osCurElt,
701 98 : &Enumerate](size_t iDim)
702 : {
703 8 : const size_t sizeBefore = osCurElt.size();
704 19 : for (uint64_t i = 0; i < oArrayInfo.anChunkCount[iDim]; ++i)
705 : {
706 13 : osCurElt += CPLSPrintf("%" PRIu64, i);
707 13 : if (iDim + 1 < oArrayInfo.anChunkCount.size())
708 : {
709 4 : osCurElt += '.';
710 4 : if (!Enumerate(iDim + 1))
711 1 : return false;
712 : }
713 : else
714 : {
715 9 : if (aosRet.size() >= nMaxFiles)
716 1 : return false;
717 8 : aosRet.AddString(osCurElt);
718 : }
719 11 : osCurElt.resize(sizeBefore);
720 : }
721 6 : return true;
722 4 : };
723 :
724 4 : Enumerate(0);
725 : }
726 : }
727 :
728 18 : return aosRet.StealList();
729 : }
730 :
731 : /************************************************************************/
732 : /* VSIInstallKerchunkParquetRefFileSystem() */
733 : /************************************************************************/
734 :
735 1653 : void VSIInstallKerchunkParquetRefFileSystem()
736 : {
737 : static std::mutex oMutex;
738 3306 : std::lock_guard<std::mutex> oLock(oMutex);
739 : // cppcheck-suppress knownConditionTrueFalse
740 1653 : if (!VSIKerchunkParquetRefFileSystem::IsFileSystemInstantiated())
741 : {
742 1653 : VSIFileManager::InstallHandler(
743 : PARQUET_REF_FS_PREFIX,
744 3306 : std::make_unique<VSIKerchunkParquetRefFileSystem>().release());
745 : }
746 1653 : }
747 :
748 : /************************************************************************/
749 : /* VSIKerchunkParquetRefFileSystemCleanCache() */
750 : /************************************************************************/
751 :
752 1055 : void VSIKerchunkParquetRefFileSystemCleanCache()
753 : {
754 0 : auto poFS = dynamic_cast<VSIKerchunkParquetRefFileSystem *>(
755 1055 : VSIFileManager::GetHandler(PARQUET_REF_FS_PREFIX));
756 1055 : if (poFS)
757 1055 : poFS->CleanCache();
758 1055 : }
|
