LCOV - code coverage report
Current view: top level - frmts/zarr - vsikerchunk_parquet_ref.cpp (source / functions) Hit Total Coverage
Test: gdal_filtered.info Lines: 345 369 93.5 %
Date: 2025-10-22 21:46:17 Functions: 16 16 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_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 : }

Generated by: LCOV version 1.14