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