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          96 : std::unique_ptr<ZarrV3CodecSequence> ZarrV3CodecSequence::Clone() const
      20             : {
      21          96 :     auto poClone = std::make_unique<ZarrV3CodecSequence>(m_oInputArrayMetadata);
      22         164 :     for (const auto &poCodec : m_apoCodecs)
      23          68 :         poClone->m_apoCodecs.emplace_back(poCodec->Clone());
      24          96 :     poClone->m_oCodecArray = m_oCodecArray.Clone();
      25          96 :     poClone->m_bPartialDecodingPossible = m_bPartialDecodingPossible;
      26          96 :     return poClone;
      27             : }
      28             : 
      29             : /************************************************************************/
      30             : /*                 ZarrV3CodecSequence::InitFromJson()                  */
      31             : /************************************************************************/
      32             : 
      33        2665 : bool ZarrV3CodecSequence::InitFromJson(const CPLJSONObject &oCodecs,
      34             :                                        ZarrArrayMetadata &oOutputArrayMetadata)
      35             : {
      36        2665 :     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        5330 :     auto oCodecsArray = oCodecs.ToArray();
      42             : 
      43        5330 :     ZarrArrayMetadata oInputArrayMetadata = m_oInputArrayMetadata;
      44        2665 :     ZarrV3Codec::IOType eLastType = ZarrV3Codec::IOType::ARRAY;
      45        5330 :     std::string osLastCodec;
      46             : 
      47             :     const auto InsertImplicitEndianCodecIfNeeded =
      48        8087 :         [this, &oInputArrayMetadata, &eLastType, &osLastCodec]()
      49             :     {
      50        4041 :         CPL_IGNORE_RET_VAL(this);
      51        4041 :         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        4041 :     };
      72             : 
      73        2665 :     bool bShardingFound = false;
      74        5330 :     std::vector<size_t> anBlockSizesBeforeSharding;
      75        6788 :     for (const auto &oCodec : oCodecsArray)
      76             :     {
      77        4143 :         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        8286 :         const auto osName = oCodec["name"].ToString();
      83           0 :         std::unique_ptr<ZarrV3Codec> poCodec;
      84        4143 :         if (osName == ZarrV3CodecGZip::NAME)
      85          59 :             poCodec = std::make_unique<ZarrV3CodecGZip>();
      86        4084 :         else if (osName == ZarrV3CodecBlosc::NAME)
      87           4 :             poCodec = std::make_unique<ZarrV3CodecBlosc>();
      88        4080 :         else if (osName == ZarrV3CodecZstd::NAME)
      89         524 :             poCodec = std::make_unique<ZarrV3CodecZstd>();
      90        5138 :         else if (osName == ZarrV3CodecBytes::NAME ||
      91        1582 :                  osName == "endian" /* endian is the old name */)
      92        1974 :             poCodec = std::make_unique<ZarrV3CodecBytes>();
      93        1582 :         else if (osName == ZarrV3CodecTranspose::NAME)
      94          83 :             poCodec = std::make_unique<ZarrV3CodecTranspose>();
      95        1499 :         else if (osName == ZarrV3CodecCRC32C::NAME)
      96         809 :             poCodec = std::make_unique<ZarrV3CodecCRC32C>();
      97         690 :         else if (osName == ZarrV3CodecShardingIndexed::NAME)
      98             :         {
      99         688 :             bShardingFound = true;
     100         688 :             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        4141 :         if (poCodec->GetInputType() == ZarrV3Codec::IOType::ARRAY)
     110             :         {
     111        2745 :             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        1396 :             InsertImplicitEndianCodecIfNeeded();
     122             :         }
     123             : 
     124        4141 :         ZarrArrayMetadata oStepOutputArrayMetadata;
     125        4141 :         if (osName == ZarrV3CodecShardingIndexed::NAME)
     126             :         {
     127         688 :             anBlockSizesBeforeSharding = oInputArrayMetadata.anBlockSizes;
     128             :         }
     129        8282 :         if (!poCodec->InitFromConfiguration(oCodec["configuration"],
     130             :                                             oInputArrayMetadata,
     131             :                                             oStepOutputArrayMetadata,
     132        4141 :                                             /* bEmitWarnings = */ true))
     133             :         {
     134          18 :             return false;
     135             :         }
     136        4123 :         oInputArrayMetadata = std::move(oStepOutputArrayMetadata);
     137        4123 :         eLastType = poCodec->GetOutputType();
     138        4123 :         osLastCodec = poCodec->GetName();
     139             : 
     140        4123 :         if (!poCodec->IsNoOp())
     141        2202 :             m_apoCodecs.emplace_back(std::move(poCodec));
     142             :     }
     143             : 
     144        2645 :     if (bShardingFound)
     145             :     {
     146         670 :         m_bPartialDecodingPossible =
     147         670 :             (m_apoCodecs.back()->GetName() == ZarrV3CodecShardingIndexed::NAME);
     148         670 :         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        2645 :     InsertImplicitEndianCodecIfNeeded();
     163             : 
     164        2645 :     m_oCodecArray = oCodecs.Clone();
     165        2645 :     oOutputArrayMetadata = std::move(oInputArrayMetadata);
     166        2645 :     return true;
     167             : }
     168             : 
     169             : /************************************************************************/
     170             : /*                  ZarrV3CodecBytes::AllocateBuffer()                  */
     171             : /************************************************************************/
     172             : 
     173       44396 : bool ZarrV3CodecSequence::AllocateBuffer(ZarrByteVectorQuickResize &abyBuffer,
     174             :                                          size_t nEltCount)
     175             : {
     176       44396 :     if (!m_apoCodecs.empty())
     177             :     {
     178       33019 :         const size_t nRawSize =
     179       33019 :             nEltCount * m_oInputArrayMetadata.oElt.nativeSize;
     180             :         // Grow the temporary buffer a bit beyond the uncompressed size
     181       33019 :         const size_t nMaxSize = nRawSize + nRawSize / 3 + 64;
     182             :         try
     183             :         {
     184       33019 :             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       33019 :         m_abyTmp.resize(nRawSize);
     192             : 
     193             :         // Grow the input/output buffer too if we have several steps
     194       33019 :         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       44396 :     return true;
     210             : }
     211             : 
     212             : /************************************************************************/
     213             : /*                    ZarrV3CodecSequence::Encode()                     */
     214             : /************************************************************************/
     215             : 
     216       15232 : bool ZarrV3CodecSequence::Encode(ZarrByteVectorQuickResize &abyBuffer)
     217             : {
     218       15232 :     if (!AllocateBuffer(abyBuffer,
     219       15232 :                         MultiplyElements(m_oInputArrayMetadata.anBlockSizes)))
     220           0 :         return false;
     221       24973 :     for (const auto &poCodec : m_apoCodecs)
     222             :     {
     223        9741 :         if (!poCodec->Encode(abyBuffer, m_abyTmp))
     224           0 :             return false;
     225        9741 :         std::swap(abyBuffer, m_abyTmp);
     226             :     }
     227       15232 :     return true;
     228             : }
     229             : 
     230             : /************************************************************************/
     231             : /*                    ZarrV3CodecSequence::Decode()                     */
     232             : /************************************************************************/
     233             : 
     234       28137 : bool ZarrV3CodecSequence::Decode(ZarrByteVectorQuickResize &abyBuffer)
     235             : {
     236       28137 :     if (!AllocateBuffer(abyBuffer,
     237       28137 :                         MultiplyElements(m_oInputArrayMetadata.anBlockSizes)))
     238           0 :         return false;
     239       49724 :     for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
     240             :     {
     241       22480 :         const auto &poCodec = *iter;
     242       22480 :         if (!poCodec->Decode(abyBuffer, m_abyTmp))
     243         893 :             return false;
     244       21587 :         std::swap(abyBuffer, m_abyTmp);
     245             :     }
     246       27244 :     return true;
     247             : }
     248             : 
     249             : /************************************************************************/
     250             : /*                 ZarrV3CodecSequence::DecodePartial()                 */
     251             : /************************************************************************/
     252             : 
     253        1027 : bool ZarrV3CodecSequence::DecodePartial(VSIVirtualHandle *poFile,
     254             :                                         ZarrByteVectorQuickResize &abyBuffer,
     255             :                                         const std::vector<size_t> &anStartIdxIn,
     256             :                                         const std::vector<size_t> &anCountIn)
     257             : {
     258        1027 :     CPLAssert(anStartIdxIn.size() == m_oInputArrayMetadata.anBlockSizes.size());
     259        1027 :     CPLAssert(anStartIdxIn.size() == anCountIn.size());
     260             : 
     261        1027 :     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        2054 :     std::vector<size_t> anStartIdx(anStartIdxIn);
     268        2054 :     std::vector<size_t> anCount(anCountIn);
     269        2079 :     for (auto &poCodec : m_apoCodecs)
     270             :     {
     271        1052 :         poCodec->ChangeArrayShapeForward(anStartIdx, anCount);
     272             :     }
     273             : 
     274        1723 :     for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
     275             :     {
     276        1052 :         const auto &poCodec = *iter;
     277             : 
     278        2104 :         if (!poCodec->DecodePartial(poFile, abyBuffer, m_abyTmp, anStartIdx,
     279        1052 :                                     anCount))
     280         356 :             return false;
     281         696 :         std::swap(abyBuffer, m_abyTmp);
     282             :     }
     283         671 :     return true;
     284             : }
     285             : 
     286             : /************************************************************************/
     287             : /*              ZarrV3CodecSequence::BatchDecodePartial()               */
     288             : /************************************************************************/
     289             : 
     290        4031 : bool ZarrV3CodecSequence::BatchDecodePartial(
     291             :     VSIVirtualHandle *poFile, const char *pszFilename,
     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        4031 :     if (m_apoCodecs.size() == 1)
     300             :     {
     301        4023 :         auto *poSharding = dynamic_cast<ZarrV3CodecShardingIndexed *>(
     302        8046 :             m_apoCodecs.back().get());
     303        4023 :         if (poSharding)
     304             :         {
     305        4023 :             return poSharding->BatchDecodePartial(poFile, pszFilename,
     306        4023 :                                                   anRequests, aResults);
     307             :         }
     308             :     }
     309             : 
     310             :     // Fallback: sequential DecodePartial for non-sharding codec chains
     311           8 :     aResults.resize(anRequests.size());
     312          32 :     for (size_t i = 0; i < anRequests.size(); ++i)
     313             :     {
     314          24 :         if (!DecodePartial(poFile, aResults[i], anRequests[i].first,
     315          24 :                            anRequests[i].second))
     316           0 :             return false;
     317             :     }
     318           8 :     return true;
     319             : }
     320             : 
     321             : /************************************************************************/
     322             : /*             ZarrV3CodecSequence::GetInnerMostBlockSize()             */
     323             : /************************************************************************/
     324             : 
     325        1281 : std::vector<size_t> ZarrV3CodecSequence::GetInnerMostBlockSize(
     326             :     const std::vector<size_t> &anOuterBlockSize) const
     327             : {
     328        1281 :     auto chunkSize = anOuterBlockSize;
     329        2286 :     for (auto iter = m_apoCodecs.rbegin(); iter != m_apoCodecs.rend(); ++iter)
     330             :     {
     331        1005 :         const auto &poCodec = *iter;
     332        1478 :         if (m_bPartialDecodingPossible ||
     333         473 :             poCodec->GetName() != ZarrV3CodecShardingIndexed::NAME)
     334             :         {
     335         867 :             chunkSize = poCodec->GetInnerMostBlockSize(chunkSize);
     336             :         }
     337             :     }
     338        1281 :     return chunkSize;
     339             : }