Line data    Source code

       1             : /******************************************************************************
       2             :  *
       3             :  * Project:  GDAL
       4             :  * Purpose:  Zarr driver, ZarrV3CodecSequence class
       5             :  * Author:   Even Rouault <even dot rouault at spatialys.com>
       6             :  *
       7             :  ******************************************************************************
       8             :  * Copyright (c) 2023, Even Rouault <even dot rouault at spatialys.com>
       9             :  *
      10             :  * SPDX-License-Identifier: MIT
      11             :  ****************************************************************************/
      12             : 
      13             : #include "zarr_v3_codec.h"
      14             : 
      15             : /************************************************************************/
      16             : /*                     ZarrV3CodecSequence::Clone()                     */
      17             : /************************************************************************/
      18             : 
      19          88 : std::unique_ptr<ZarrV3CodecSequence> ZarrV3CodecSequence::Clone() const
      20             : {
      21          88 :     auto poClone = std::make_unique<ZarrV3CodecSequence>(m_oInputArrayMetadata);
      22         156 :     for (const auto &poCodec : m_apoCodecs)
      23          68 :         poClone->m_apoCodecs.emplace_back(poCodec->Clone());
      24          88 :     poClone->m_oCodecArray = m_oCodecArray.Clone();
      25          88 :     poClone->m_bPartialDecodingPossible = m_bPartialDecodingPossible;
      26          88 :     return poClone;
      27             : }
      28             : 
      29             : /************************************************************************/
      30             : /*                 ZarrV3CodecSequence::InitFromJson()                  */
      31             : /************************************************************************/
      32             : 
      33        2792 : bool ZarrV3CodecSequence::InitFromJson(const std::string &osArrayName,
      34             :                                        const CPLJSONObject &oCodecs,
      35             :                                        ZarrArrayMetadata &oOutputArrayMetadata)
      36             : {
      37        2792 :     if (oCodecs.GetType() != CPLJSONObject::Type::Array)
      38             :     {
      39           0 :         CPLError(CE_Failure, CPLE_AppDefined, "%s: codecs is not an array",
      40             :                  osArrayName.c_str());
      41           0 :         return false;
      42             :     }
      43        5584 :     auto oCodecsArray = oCodecs.ToArray();
      44             : 
      45        5584 :     ZarrArrayMetadata oInputArrayMetadata = m_oInputArrayMetadata;
      46        2792 :     ZarrV3Codec::IOType eLastType = ZarrV3Codec::IOType::ARRAY;
      47        5584 :     std::string osLastCodec;
      48             : 
      49             :     const auto InsertImplicitEndianCodecIfNeeded =
      50        8453 :         [this, &oInputArrayMetadata, &eLastType, &osLastCodec]()
      51             :     {
      52        4223 :         CPL_IGNORE_RET_VAL(this);
      53        4223 :         if (eLastType == ZarrV3Codec::IOType::ARRAY &&
      54           1 :             oInputArrayMetadata.oElt.nativeSize > 1 &&
      55           1 :             oInputArrayMetadata.oElt.nativeType !=
      56           1 :                 DtypeElt::NativeType::STRING_ASCII &&
      57           1 :             oInputArrayMetadata.oElt.nativeType !=
      58             :                 DtypeElt::NativeType::STRING_UNICODE)
      59             :         {
      60           1 :             CPLError(CE_Warning, CPLE_AppDefined,
      61             :                      "'bytes' codec missing. Assuming little-endian storage, "
      62             :                      "but such tolerance may be removed in future versions");
      63           2 :             auto poEndianCodec = std::make_unique<ZarrV3CodecBytes>();
      64           2 :             ZarrArrayMetadata oTmpOutputArrayMetadata;
      65           2 :             poEndianCodec->InitFromConfiguration(
      66           2 :                 std::string(), ZarrV3CodecBytes::GetConfiguration(true),
      67             :                 oInputArrayMetadata, oTmpOutputArrayMetadata,
      68             :                 /* bEmitWarnings = */ true);
      69           1 :             oInputArrayMetadata = std::move(oTmpOutputArrayMetadata);
      70           1 :             eLastType = poEndianCodec->GetOutputType();
      71           1 :             osLastCodec = poEndianCodec->GetName();
      72             :             if constexpr (!CPL_IS_LSB)
      73             :             {
      74             :                 // Insert a little endian codec if we are on a big endian target
      75             :                 m_apoCodecs.emplace_back(std::move(poEndianCodec));
      76             :             }
      77             :         }
      78        4223 :     };
      79             : 
      80        2792 :     bool bShardingFound = false;
      81        5584 :     std::vector<size_t> anBlockSizesBeforeSharding;
      82        7097 :     for (const auto &oCodec : oCodecsArray)
      83             :     {
      84        4325 :         if (oCodec.GetType() != CPLJSONObject::Type::Object)
      85             :         {
      86           0 :             CPLError(CE_Failure, CPLE_AppDefined, "codecs[] is not an object");
      87          20 :             return false;
      88             :         }
      89        8650 :         const auto osName = oCodec["name"].ToString();
      90           0 :         std::unique_ptr<ZarrV3Codec> poCodec;
      91        4325 :         if (osName == ZarrV3CodecGZip::NAME)
      92          59 :             poCodec = std::make_unique<ZarrV3CodecGZip>();
      93        4266 :         else if (osName == ZarrV3CodecBlosc::NAME)
      94          28 :             poCodec = std::make_unique<ZarrV3CodecBlosc>();
      95        4238 :         else if (osName == ZarrV3CodecZstd::NAME)
      96         533 :             poCodec = std::make_unique<ZarrV3CodecZstd>();
      97        5345 :         else if (osName == ZarrV3CodecBytes::NAME ||
      98        1640 :                  osName == "endian" /* endian is the old name */)
      99        2065 :             poCodec = std::make_unique<ZarrV3CodecBytes>();
     100        1640 :         else if (osName == ZarrV3CodecTranspose::NAME)
     101          83 :             poCodec = std::make_unique<ZarrV3CodecTranspose>();
     102        1557 :         else if (osName == ZarrV3CodecCRC32C::NAME)
     103         831 :             poCodec = std::make_unique<ZarrV3CodecCRC32C>();
     104         726 :         else if (osName == ZarrV3CodecVLenUTF8::NAME)
     105          14 :             poCodec = std::make_unique<ZarrV3CodecVLenUTF8>();
     106         712 :         else if (osName == ZarrV3CodecShardingIndexed::NAME)
     107             :         {
     108         710 :             bShardingFound = true;
     109         710 :             poCodec = std::make_unique<ZarrV3CodecShardingIndexed>();
     110             :         }
     111           2 :         else if (osName == ZarrV3CodecPcodec::NAME)
     112             :         {
     113             : #ifdef HAVE_PCODEC
     114             :             constexpr bool bHasPcodec = true;
     115             : #else
     116           0 :             constexpr bool bHasPcodec = false;
     117             : #endif
     118             :             if constexpr (bHasPcodec)
     119             :             {
     120             :                 poCodec = std::make_unique<ZarrV3CodecPcodec>();
     121             :             }
     122             :             else
     123             :             {
     124           0 :                 CPLError(
     125             :                     CE_Failure, CPLE_NotSupported,
     126             :                     "%s: Codec '%s' is supported by GDAL, but not in this "
     127             :                     "particular build. Requires building GDAL with "
     128             :                     "-DGDAL_USE_PCODEC=ON and with a Rust toolchain available",
     129             :                     osArrayName.c_str(), osName.c_str());
     130           0 :                 return false;
     131             :             }
     132             :         }
     133             :         else
     134             :         {
     135           2 :             CPLError(CE_Failure, CPLE_NotSupported, "%s: Unsupported codec: %s",
     136             :                      osArrayName.c_str(), osName.c_str());
     137           2 :             return false;
     138             :         }
     139             : 
     140        4323 :         if (poCodec->GetInputType() == ZarrV3Codec::IOType::ARRAY)
     141             :         {
     142        2872 :             if (eLastType == ZarrV3Codec::IOType::BYTES)
     143             :             {
     144           0 :                 CPLError(CE_Failure, CPLE_AppDefined,
     145             :                          "Cannot chain codec %s with %s",
     146           0 :                          poCodec->GetName().c_str(), osLastCodec.c_str());
     147           0 :                 return false;
     148             :             }
     149             :         }
     150             :         else
     151             :         {
     152        1451 :             InsertImplicitEndianCodecIfNeeded();
     153             :         }
     154             : 
     155        4323 :         ZarrArrayMetadata oStepOutputArrayMetadata;
     156        4323 :         if (osName == ZarrV3CodecShardingIndexed::NAME)
     157             :         {
     158         710 :             anBlockSizesBeforeSharding = oInputArrayMetadata.anBlockSizes;
     159             :         }
     160        8646 :         if (!poCodec->InitFromConfiguration(
     161             :                 osArrayName, oCodec["configuration"], oInputArrayMetadata,
     162             :                 oStepOutputArrayMetadata,
     163        4323 :                 /* bEmitWarnings = */ true))
     164             :         {
     165          18 :             return false;
     166             :         }
     167        4305 :         oInputArrayMetadata = std::move(oStepOutputArrayMetadata);
     168        4305 :         eLastType = poCodec->GetOutputType();
     169        4305 :         osLastCodec = poCodec->GetName();
     170             : 
     171        4305 :         if (!poCodec->IsNoOp())
     172        2293 :             m_apoCodecs.emplace_back(std::move(poCodec));
     173             :     }
     174             : 
     175        2772 :     if (bShardingFound)
     176             :     {
     177         692 :         m_bPartialDecodingPossible =
     178         692 :             (m_apoCodecs.back()->GetName() == ZarrV3CodecShardingIndexed::NAME);
     179         692 :         if (!m_bPartialDecodingPossible)
     180             :         {
     181         139 :             m_bPartialDecodingPossible = false;
     182             :             // This is not an implementation limitation, but the result of a
     183             :             // badly thought dataset. Zarr-Python also emits a similar warning.
     184         139 :             CPLError(
     185             :                 CE_Warning, CPLE_AppDefined,
     186             :                 "Sharding codec found, but not in last position. Consequently "
     187             :                 "partial shard decoding will not be possible");
     188             :             oInputArrayMetadata.anBlockSizes =
     189         139 :                 std::move(anBlockSizesBeforeSharding);
     190             :         }
     191             :     }
     192             : 
     193        2772 :     InsertImplicitEndianCodecIfNeeded();
     194             : 
     195        2772 :     m_oCodecArray = oCodecs.Clone();
     196        2772 :     oOutputArrayMetadata = std::move(oInputArrayMetadata);
     197        2772 :     return true;
     198             : }
     199             : 
     200             : /************************************************************************/
     201             : /*                  ZarrV3CodecBytes::AllocateBuffer()                  */
     202             : /************************************************************************/
     203             : 
     204       44465 : bool ZarrV3CodecSequence::AllocateBuffer(ZarrByteVectorQuickResize &abyBuffer,
     205             :                                          size_t nEltCount)
     206             : {
     207       44465 :     if (!m_apoCodecs.empty())
     208             :     {
     209       33032 :         const size_t nRawSize =
     210       33032 :             nEltCount * m_oInputArrayMetadata.oElt.nativeSize;
     211             :         // Grow the temporary buffer a bit beyond the uncompressed size
     212       33032 :         const size_t nMaxSize = nRawSize + nRawSize / 3 + 64;
     213             :         try
     214             :         {
     215       33032 :             m_abyTmp.resize(nMaxSize);
     216             :         }
     217           0 :         catch (const std::exception &e)
     218             :         {
     219           0 :             CPLError(CE_Failure, CPLE_OutOfMemory, "%s", e.what());
     220           0 :             return false;
     221             :         }
     222       33032 :         m_abyTmp.resize(nRawSize);
     223             : 
     224             :         // Grow the input/output buffer too if we have several steps
     225       33032 :         if (m_apoCodecs.size() >= 2 && abyBuffer.capacity() < nMaxSize)
     226             :         {
     227         223 :             const size_t nSize = abyBuffer.size();
     228             :             try
     229             :             {
     230         223 :                 abyBuffer.resize(nMaxSize);
     231             :             }
     232           0 :             catch (const std::exception &e)
     233             :             {
     234           0 :                 CPLError(CE_Failure, CPLE_OutOfMemory, "%s", e.what());
     235           0 :                 return false;
     236             :             }
     237         223 :             abyBuffer.resize(nSize);
     238             :         }
     239             :     }
     240       44465 :     return true;
     241             : }
     242             : 
     243             : /************************************************************************/
     244             : /*                    ZarrV3CodecSequence::Encode()                     */
     245             : /************************************************************************/
     246             : 
     247       15239 : bool ZarrV3CodecSequence::Encode(ZarrByteVectorQuickResize &abyBuffer)
     248             : {
     249       15239 :     if (!AllocateBuffer(abyBuffer,
     250       15239 :                         MultiplyElements(m_oInputArrayMetadata.anBlockSizes)))
     251           0 :         return false;
     252       24985 :     for (const auto &poCodec : m_apoCodecs)
     253             :     {
     254        9746 :         if (!poCodec->Encode(abyBuffer, m_abyTmp))
     255           0 :             return false;
     256        9746 :         std::swap(abyBuffer, m_abyTmp);
     257             :     }
     258       15239 :     return true;
     259             : }
     260             : 
     261             : /************************************************************************/
     262             : /*                    ZarrV3CodecSequence::Decode()                     */
     263             : /************************************************************************/
     264             : 
     265       28198 : bool ZarrV3CodecSequence::Decode(ZarrByteVectorQuickResize &abyBuffer)
     266             : {
     267       28198 :     if (!AllocateBuffer(abyBuffer,
     268       28198 :                         MultiplyElements(m_oInputArrayMetadata.anBlockSizes)))
     269           0 :         return false;
     270       49795 :     for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
     271             :     {
     272       22491 :         const auto &poCodec = *iter;
     273       22491 :         if (!poCodec->Decode(abyBuffer, m_abyTmp))
     274         894 :             return false;
     275       21597 :         std::swap(abyBuffer, m_abyTmp);
     276             :     }
     277       27304 :     return true;
     278             : }
     279             : 
     280             : /************************************************************************/
     281             : /*                 ZarrV3CodecSequence::DecodePartial()                 */
     282             : /************************************************************************/
     283             : 
     284        1028 : bool ZarrV3CodecSequence::DecodePartial(VSIVirtualHandle *poFile,
     285             :                                         ZarrByteVectorQuickResize &abyBuffer,
     286             :                                         const std::vector<size_t> &anStartIdxIn,
     287             :                                         const std::vector<size_t> &anCountIn)
     288             : {
     289        1028 :     CPLAssert(anStartIdxIn.size() == m_oInputArrayMetadata.anBlockSizes.size());
     290        1028 :     CPLAssert(anStartIdxIn.size() == anCountIn.size());
     291             : 
     292        1028 :     if (!AllocateBuffer(abyBuffer, MultiplyElements(anCountIn)))
     293           0 :         return false;
     294             : 
     295             :     // anStartIdxIn and anCountIn are expressed in the shape *before* encoding
     296             :     // We need to apply the potential transpositions before submitting them
     297             :     // to the decoder of the Array->Bytes decoder
     298        2056 :     std::vector<size_t> anStartIdx(anStartIdxIn);
     299        2056 :     std::vector<size_t> anCount(anCountIn);
     300        2081 :     for (auto &poCodec : m_apoCodecs)
     301             :     {
     302        1053 :         poCodec->ChangeArrayShapeForward(anStartIdx, anCount);
     303             :     }
     304             : 
     305        1725 :     for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
     306             :     {
     307        1053 :         const auto &poCodec = *iter;
     308             : 
     309        2106 :         if (!poCodec->DecodePartial(poFile, abyBuffer, m_abyTmp, anStartIdx,
     310        1053 :                                     anCount))
     311         356 :             return false;
     312         697 :         std::swap(abyBuffer, m_abyTmp);
     313             :     }
     314         672 :     return true;
     315             : }
     316             : 
     317             : /************************************************************************/
     318             : /*              ZarrV3CodecSequence::BatchDecodePartial()               */
     319             : /************************************************************************/
     320             : 
     321        4031 : bool ZarrV3CodecSequence::BatchDecodePartial(
     322             :     VSIVirtualHandle *poFile, const char *pszFilename,
     323             :     const std::vector<std::pair<std::vector<size_t>, std::vector<size_t>>>
     324             :         &anRequests,
     325             :     std::vector<ZarrByteVectorQuickResize> &aResults)
     326             : {
     327             :     // Only batch-decode when sharding is the sole codec. If other codecs
     328             :     // (e.g. transpose) precede it, indices and output need codec-specific
     329             :     // transformations that BatchDecodePartial does not handle.
     330        4031 :     if (m_apoCodecs.size() == 1)
     331             :     {
     332        4023 :         auto *poSharding = dynamic_cast<ZarrV3CodecShardingIndexed *>(
     333        8046 :             m_apoCodecs.back().get());
     334        4023 :         if (poSharding)
     335             :         {
     336        4023 :             return poSharding->BatchDecodePartial(poFile, pszFilename,
     337        4023 :                                                   anRequests, aResults);
     338             :         }
     339             :     }
     340             : 
     341             :     // Fallback: sequential DecodePartial for non-sharding codec chains
     342           8 :     aResults.resize(anRequests.size());
     343          32 :     for (size_t i = 0; i < anRequests.size(); ++i)
     344             :     {
     345          24 :         if (!DecodePartial(poFile, aResults[i], anRequests[i].first,
     346          24 :                            anRequests[i].second))
     347           0 :             return false;
     348             :     }
     349           8 :     return true;
     350             : }
     351             : 
     352             : /************************************************************************/
     353             : /*             ZarrV3CodecSequence::GetInnerMostBlockSize()             */
     354             : /************************************************************************/
     355             : 
     356        1364 : std::vector<size_t> ZarrV3CodecSequence::GetInnerMostBlockSize(
     357             :     const std::vector<size_t> &anOuterBlockSize) const
     358             : {
     359        1364 :     auto chunkSize = anOuterBlockSize;
     360        2418 :     for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
     361             :     {
     362        1054 :         const auto &poCodec = *iter;
     363        1555 :         if (m_bPartialDecodingPossible ||
     364         501 :             poCodec->GetName() != ZarrV3CodecShardingIndexed::NAME)
     365             :         {
     366         915 :             chunkSize = poCodec->GetInnerMostBlockSize(chunkSize);
     367             :         }
     368             :     }
     369        1364 :     return chunkSize;
     370             : }