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          20 : std::unique_ptr<ZarrV3CodecSequence> ZarrV3CodecSequence::Clone() const
      20             : {
      21          20 :     auto poClone = std::make_unique<ZarrV3CodecSequence>(m_oInputArrayMetadata);
      22          32 :     for (const auto &poCodec : m_apoCodecs)
      23          12 :         poClone->m_apoCodecs.emplace_back(poCodec->Clone());
      24          20 :     poClone->m_oCodecArray = m_oCodecArray.Clone();
      25          20 :     poClone->m_bPartialDecodingPossible = m_bPartialDecodingPossible;
      26          20 :     return poClone;
      27             : }
      28             : 
      29             : /************************************************************************/
      30             : /*                 ZarrV3CodecSequence::InitFromJson()                  */
      31             : /************************************************************************/
      32             : 
      33        2458 : bool ZarrV3CodecSequence::InitFromJson(const CPLJSONObject &oCodecs,
      34             :                                        ZarrArrayMetadata &oOutputArrayMetadata)
      35             : {
      36        2458 :     if (oCodecs.GetType() != CPLJSONObject::Type::Array)
      37             :     {
      38           0 :         CPLError(CE_Failure, CPLE_AppDefined, "codecs is not an array");
      39           0 :         return false;
      40             :     }
      41        4916 :     auto oCodecsArray = oCodecs.ToArray();
      42             : 
      43        4916 :     ZarrArrayMetadata oInputArrayMetadata = m_oInputArrayMetadata;
      44        2458 :     ZarrV3Codec::IOType eLastType = ZarrV3Codec::IOType::ARRAY;
      45        4916 :     std::string osLastCodec;
      46             : 
      47             :     const auto InsertImplicitEndianCodecIfNeeded =
      48        7537 :         [this, &oInputArrayMetadata, &eLastType, &osLastCodec]()
      49             :     {
      50        3766 :         CPL_IGNORE_RET_VAL(this);
      51        3766 :         if (eLastType == ZarrV3Codec::IOType::ARRAY &&
      52           1 :             oInputArrayMetadata.oElt.nativeSize > 1)
      53             :         {
      54           1 :             CPLError(CE_Warning, CPLE_AppDefined,
      55             :                      "'bytes' codec missing. Assuming little-endian storage, "
      56             :                      "but such tolerance may be removed in future versions");
      57           2 :             auto poEndianCodec = std::make_unique<ZarrV3CodecBytes>();
      58           2 :             ZarrArrayMetadata oTmpOutputArrayMetadata;
      59           2 :             poEndianCodec->InitFromConfiguration(
      60           2 :                 ZarrV3CodecBytes::GetConfiguration(true), oInputArrayMetadata,
      61             :                 oTmpOutputArrayMetadata, /* bEmitWarnings = */ true);
      62           1 :             oInputArrayMetadata = std::move(oTmpOutputArrayMetadata);
      63           1 :             eLastType = poEndianCodec->GetOutputType();
      64           1 :             osLastCodec = poEndianCodec->GetName();
      65             :             if constexpr (!CPL_IS_LSB)
      66             :             {
      67             :                 // Insert a little endian codec if we are on a big endian target
      68             :                 m_apoCodecs.emplace_back(std::move(poEndianCodec));
      69             :             }
      70             :         }
      71        3766 :     };
      72             : 
      73        2458 :     bool bShardingFound = false;
      74        4916 :     std::vector<size_t> anBlockSizesBeforeSharding;
      75        6306 :     for (const auto &oCodec : oCodecsArray)
      76             :     {
      77        3868 :         if (oCodec.GetType() != CPLJSONObject::Type::Object)
      78             :         {
      79           0 :             CPLError(CE_Failure, CPLE_AppDefined, "codecs[] is not an object");
      80          20 :             return false;
      81             :         }
      82        7736 :         const auto osName = oCodec["name"].ToString();
      83           0 :         std::unique_ptr<ZarrV3Codec> poCodec;
      84        3868 :         if (osName == ZarrV3CodecGZip::NAME)
      85          48 :             poCodec = std::make_unique<ZarrV3CodecGZip>();
      86        3820 :         else if (osName == ZarrV3CodecBlosc::NAME)
      87           4 :             poCodec = std::make_unique<ZarrV3CodecBlosc>();
      88        3816 :         else if (osName == ZarrV3CodecZstd::NAME)
      89         503 :             poCodec = std::make_unique<ZarrV3CodecZstd>();
      90        4827 :         else if (osName == ZarrV3CodecBytes::NAME ||
      91        1514 :                  osName == "endian" /* endian is the old name */)
      92        1799 :             poCodec = std::make_unique<ZarrV3CodecBytes>();
      93        1514 :         else if (osName == ZarrV3CodecTranspose::NAME)
      94          83 :             poCodec = std::make_unique<ZarrV3CodecTranspose>();
      95        1431 :         else if (osName == ZarrV3CodecCRC32C::NAME)
      96         773 :             poCodec = std::make_unique<ZarrV3CodecCRC32C>();
      97         658 :         else if (osName == ZarrV3CodecShardingIndexed::NAME)
      98             :         {
      99         656 :             bShardingFound = true;
     100         656 :             poCodec = std::make_unique<ZarrV3CodecShardingIndexed>();
     101             :         }
     102             :         else
     103             :         {
     104           2 :             CPLError(CE_Failure, CPLE_NotSupported, "Unsupported codec: %s",
     105             :                      osName.c_str());
     106           2 :             return false;
     107             :         }
     108             : 
     109        3866 :         if (poCodec->GetInputType() == ZarrV3Codec::IOType::ARRAY)
     110             :         {
     111        2538 :             if (eLastType == ZarrV3Codec::IOType::BYTES)
     112             :             {
     113           0 :                 CPLError(CE_Failure, CPLE_AppDefined,
     114             :                          "Cannot chain codec %s with %s",
     115           0 :                          poCodec->GetName().c_str(), osLastCodec.c_str());
     116           0 :                 return false;
     117             :             }
     118             :         }
     119             :         else
     120             :         {
     121        1328 :             InsertImplicitEndianCodecIfNeeded();
     122             :         }
     123             : 
     124        3866 :         ZarrArrayMetadata oStepOutputArrayMetadata;
     125        3866 :         if (osName == ZarrV3CodecShardingIndexed::NAME)
     126             :         {
     127         656 :             anBlockSizesBeforeSharding = oInputArrayMetadata.anBlockSizes;
     128             :         }
     129        7732 :         if (!poCodec->InitFromConfiguration(oCodec["configuration"],
     130             :                                             oInputArrayMetadata,
     131             :                                             oStepOutputArrayMetadata,
     132        3866 :                                             /* bEmitWarnings = */ true))
     133             :         {
     134          18 :             return false;
     135             :         }
     136        3848 :         oInputArrayMetadata = std::move(oStepOutputArrayMetadata);
     137        3848 :         eLastType = poCodec->GetOutputType();
     138        3848 :         osLastCodec = poCodec->GetName();
     139             : 
     140        3848 :         if (!poCodec->IsNoOp())
     141        2102 :             m_apoCodecs.emplace_back(std::move(poCodec));
     142             :     }
     143             : 
     144        2438 :     if (bShardingFound)
     145             :     {
     146         638 :         m_bPartialDecodingPossible =
     147         638 :             (m_apoCodecs.back()->GetName() == ZarrV3CodecShardingIndexed::NAME);
     148         638 :         if (!m_bPartialDecodingPossible)
     149             :         {
     150         138 :             m_bPartialDecodingPossible = false;
     151             :             // This is not an implementation limitation, but the result of a
     152             :             // badly thought dataset. Zarr-Python also emits a similar warning.
     153         138 :             CPLError(
     154             :                 CE_Warning, CPLE_AppDefined,
     155             :                 "Sharding codec found, but not in last position. Consequently "
     156             :                 "partial shard decoding will not be possible");
     157             :             oInputArrayMetadata.anBlockSizes =
     158         138 :                 std::move(anBlockSizesBeforeSharding);
     159             :         }
     160             :     }
     161             : 
     162        2438 :     InsertImplicitEndianCodecIfNeeded();
     163             : 
     164        2438 :     m_oCodecArray = oCodecs.Clone();
     165        2438 :     oOutputArrayMetadata = std::move(oInputArrayMetadata);
     166        2438 :     return true;
     167             : }
     168             : 
     169             : /************************************************************************/
     170             : /*                  ZarrV3CodecBytes::AllocateBuffer()                  */
     171             : /************************************************************************/
     172             : 
     173       35188 : bool ZarrV3CodecSequence::AllocateBuffer(ZarrByteVectorQuickResize &abyBuffer,
     174             :                                          size_t nEltCount)
     175             : {
     176       35188 :     if (!m_apoCodecs.empty())
     177             :     {
     178       24279 :         const size_t nRawSize =
     179       24279 :             nEltCount * m_oInputArrayMetadata.oElt.nativeSize;
     180             :         // Grow the temporary buffer a bit beyond the uncompressed size
     181       24279 :         const size_t nMaxSize = nRawSize + nRawSize / 3 + 64;
     182             :         try
     183             :         {
     184       24279 :             m_abyTmp.resize(nMaxSize);
     185             :         }
     186           0 :         catch (const std::exception &e)
     187             :         {
     188           0 :             CPLError(CE_Failure, CPLE_OutOfMemory, "%s", e.what());
     189           0 :             return false;
     190             :         }
     191       24279 :         m_abyTmp.resize(nRawSize);
     192             : 
     193             :         // Grow the input/output buffer too if we have several steps
     194       24279 :         if (m_apoCodecs.size() >= 2 && abyBuffer.capacity() < nMaxSize)
     195             :         {
     196         219 :             const size_t nSize = abyBuffer.size();
     197             :             try
     198             :             {
     199         219 :                 abyBuffer.resize(nMaxSize);
     200             :             }
     201           0 :             catch (const std::exception &e)
     202             :             {
     203           0 :                 CPLError(CE_Failure, CPLE_OutOfMemory, "%s", e.what());
     204           0 :                 return false;
     205             :             }
     206         219 :             abyBuffer.resize(nSize);
     207             :         }
     208             :     }
     209       35188 :     return true;
     210             : }
     211             : 
     212             : /************************************************************************/
     213             : /*                    ZarrV3CodecSequence::Encode()                     */
     214             : /************************************************************************/
     215             : 
     216       10750 : bool ZarrV3CodecSequence::Encode(ZarrByteVectorQuickResize &abyBuffer)
     217             : {
     218       10750 :     if (!AllocateBuffer(abyBuffer,
     219       10750 :                         MultiplyElements(m_oInputArrayMetadata.anBlockSizes)))
     220           0 :         return false;
     221       16191 :     for (const auto &poCodec : m_apoCodecs)
     222             :     {
     223        5441 :         if (!poCodec->Encode(abyBuffer, m_abyTmp))
     224           0 :             return false;
     225        5441 :         std::swap(abyBuffer, m_abyTmp);
     226             :     }
     227       10750 :     return true;
     228             : }
     229             : 
     230             : /************************************************************************/
     231             : /*                    ZarrV3CodecSequence::Decode()                     */
     232             : /************************************************************************/
     233             : 
     234       23282 : bool ZarrV3CodecSequence::Decode(ZarrByteVectorQuickResize &abyBuffer)
     235             : {
     236       23282 :     if (!AllocateBuffer(abyBuffer,
     237       23282 :                         MultiplyElements(m_oInputArrayMetadata.anBlockSizes)))
     238           0 :         return false;
     239       40286 :     for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
     240             :     {
     241       17911 :         const auto &poCodec = *iter;
     242       17911 :         if (!poCodec->Decode(abyBuffer, m_abyTmp))
     243         907 :             return false;
     244       17004 :         std::swap(abyBuffer, m_abyTmp);
     245             :     }
     246       22375 :     return true;
     247             : }
     248             : 
     249             : /************************************************************************/
     250             : /*                 ZarrV3CodecSequence::DecodePartial()                 */
     251             : /************************************************************************/
     252             : 
     253        1156 : bool ZarrV3CodecSequence::DecodePartial(VSIVirtualHandle *poFile,
     254             :                                         ZarrByteVectorQuickResize &abyBuffer,
     255             :                                         const std::vector<size_t> &anStartIdxIn,
     256             :                                         const std::vector<size_t> &anCountIn)
     257             : {
     258        1156 :     CPLAssert(anStartIdxIn.size() == m_oInputArrayMetadata.anBlockSizes.size());
     259        1156 :     CPLAssert(anStartIdxIn.size() == anCountIn.size());
     260             : 
     261        1156 :     if (!AllocateBuffer(abyBuffer, MultiplyElements(anCountIn)))
     262           0 :         return false;
     263             : 
     264             :     // anStartIdxIn and anCountIn are expressed in the shape *before* encoding
     265             :     // We need to apply the potential transpositions before submitting them
     266             :     // to the decoder of the Array->Bytes decoder
     267        2312 :     std::vector<size_t> anStartIdx(anStartIdxIn);
     268        2312 :     std::vector<size_t> anCount(anCountIn);
     269        2337 :     for (auto &poCodec : m_apoCodecs)
     270             :     {
     271        1181 :         poCodec->ChangeArrayShapeForward(anStartIdx, anCount);
     272             :     }
     273             : 
     274        1901 :     for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
     275             :     {
     276        1181 :         const auto &poCodec = *iter;
     277             : 
     278        2362 :         if (!poCodec->DecodePartial(poFile, abyBuffer, m_abyTmp, anStartIdx,
     279        1181 :                                     anCount))
     280         436 :             return false;
     281         745 :         std::swap(abyBuffer, m_abyTmp);
     282             :     }
     283         720 :     return true;
     284             : }
     285             : 
     286             : /************************************************************************/
     287             : /*              ZarrV3CodecSequence::BatchDecodePartial()               */
     288             : /************************************************************************/
     289             : 
     290        3948 : bool ZarrV3CodecSequence::BatchDecodePartial(
     291             :     VSIVirtualHandle *poFile,
     292             :     const std::vector<std::pair<std::vector<size_t>, std::vector<size_t>>>
     293             :         &anRequests,
     294             :     std::vector<ZarrByteVectorQuickResize> &aResults)
     295             : {
     296             :     // Only batch-decode when sharding is the sole codec. If other codecs
     297             :     // (e.g. transpose) precede it, indices and output need codec-specific
     298             :     // transformations that BatchDecodePartial does not handle.
     299        3948 :     if (m_apoCodecs.size() == 1)
     300             :     {
     301        3940 :         auto *poSharding = dynamic_cast<ZarrV3CodecShardingIndexed *>(
     302        7880 :             m_apoCodecs.back().get());
     303        3940 :         if (poSharding)
     304             :         {
     305        3940 :             return poSharding->BatchDecodePartial(poFile, anRequests, aResults);
     306             :         }
     307             :     }
     308             : 
     309             :     // Fallback: sequential DecodePartial for non-sharding codec chains
     310           8 :     aResults.resize(anRequests.size());
     311          32 :     for (size_t i = 0; i < anRequests.size(); ++i)
     312             :     {
     313          24 :         if (!DecodePartial(poFile, aResults[i], anRequests[i].first,
     314          24 :                            anRequests[i].second))
     315           0 :             return false;
     316             :     }
     317           8 :     return true;
     318             : }
     319             : 
     320             : /************************************************************************/
     321             : /*             ZarrV3CodecSequence::GetInnerMostBlockSize()             */
     322             : /************************************************************************/
     323             : 
     324        1146 : std::vector<size_t> ZarrV3CodecSequence::GetInnerMostBlockSize(
     325             :     const std::vector<size_t> &anOuterBlockSize) const
     326             : {
     327        1146 :     auto chunkSize = anOuterBlockSize;
     328        2108 :     for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
     329             :     {
     330         962 :         const auto &poCodec = *iter;
     331        1424 :         if (m_bPartialDecodingPossible ||
     332         462 :             poCodec->GetName() != ZarrV3CodecShardingIndexed::NAME)
     333             :         {
     334         824 :             chunkSize = poCodec->GetInnerMostBlockSize(chunkSize);
     335             :         }
     336             :     }
     337        1146 :     return chunkSize;
     338             : }