LCOV - code coverage report
Current view: top level - frmts/zarr - vsikerchunk_parquet_ref.cpp (source / functions) Hit Total Coverage
Test: gdal_filtered.info Lines: 312 322 96.9 %
Date: 2025-07-09 17:50:03 Functions: 15 15 100.0 %

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

Generated by: LCOV version 1.14