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