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