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