Line data    Source code

       1             : /******************************************************************************
       2             :  *
       3             :  * Project:  OpenGIS Simple Features Reference Implementation
       4             :  * Purpose:  Helper to fill ArrowArray
       5             :  * Author:   Even Rouault <even dot rouault at spatialys.com>
       6             :  *
       7             :  ******************************************************************************
       8             :  * Copyright (c) 2022, Even Rouault <even dot rouault at spatialys.com>
       9             :  *
      10             :  * SPDX-License-Identifier: MIT
      11             :  ****************************************************************************/
      12             : 
      13             : #pragma once
      14             : 
      15             : //! @cond Doxygen_Suppress
      16             : 
      17             : #include <algorithm>
      18             : #include <limits>
      19             : 
      20             : #include "cpl_time.h"
      21             : 
      22             : #include "ogrsf_frmts.h"
      23             : #include "ogr_recordbatch.h"
      24             : 
      25             : class CPL_DLL OGRArrowArrayHelper
      26             : {
      27             :     OGRArrowArrayHelper(const OGRArrowArrayHelper &) = delete;
      28             :     OGRArrowArrayHelper &operator=(const OGRArrowArrayHelper &) = delete;
      29             : 
      30             :   public:
      31             :     bool m_bIncludeFID = false;
      32             :     int m_nMaxBatchSize = 0;
      33             :     int m_nChildren = 0;
      34             :     const int m_nFieldCount = 0;
      35             :     const int m_nGeomFieldCount = 0;
      36             :     std::vector<int> m_mapOGRFieldToArrowField{};
      37             :     std::vector<int> m_mapOGRGeomFieldToArrowField{};
      38             :     std::vector<bool> m_abNullableFields{};
      39             :     std::vector<uint32_t> m_anArrowFieldMaxAlloc{};
      40             :     std::vector<int> m_anTZFlags{};
      41             :     int64_t *m_panFIDValues = nullptr;
      42             :     struct ArrowArray *m_out_array = nullptr;
      43             : 
      44             :     static uint32_t GetMemLimit();
      45             : 
      46             :     static int
      47             :     GetMaxFeaturesInBatch(const CPLStringList &aosArrowArrayStreamOptions);
      48             : 
      49             :     OGRArrowArrayHelper(GDALDataset *poDS, OGRFeatureDefn *poFeatureDefn,
      50             :                         const CPLStringList &aosArrowArrayStreamOptions,
      51             :                         struct ArrowArray *out_array);
      52             : 
      53             :     //! Construct an helper from an already initialized array
      54             :     OGRArrowArrayHelper(struct ArrowArray *out_array, int nMaxBatchSize);
      55             : 
      56        1751 :     static bool SetNull(struct ArrowArray *psArray, int iFeat,
      57             :                         int nMaxBatchSize, bool bAlignedMalloc)
      58             :     {
      59        1751 :         ++psArray->null_count;
      60        1751 :         uint8_t *pabyNull =
      61        1751 :             static_cast<uint8_t *>(const_cast<void *>(psArray->buffers[0]));
      62        1751 :         if (psArray->buffers[0] == nullptr)
      63             :         {
      64         802 :             pabyNull = static_cast<uint8_t *>(
      65             :                 bAlignedMalloc
      66         743 :                     ? VSI_MALLOC_ALIGNED_AUTO_VERBOSE((nMaxBatchSize + 7) / 8)
      67          59 :                     : VSI_MALLOC_VERBOSE((nMaxBatchSize + 7) / 8));
      68         802 :             if (pabyNull == nullptr)
      69             :             {
      70           0 :                 return false;
      71             :             }
      72         802 :             memset(pabyNull, 0xFF, (nMaxBatchSize + 7) / 8);
      73         802 :             psArray->buffers[0] = pabyNull;
      74             :         }
      75        1751 :         pabyNull[iFeat / 8] &= static_cast<uint8_t>(~(1 << (iFeat % 8)));
      76             : 
      77        1751 :         if (psArray->n_buffers == 3)
      78             :         {
      79        1493 :             auto panOffsets =
      80        1493 :                 static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]));
      81        1493 :             panOffsets[iFeat + 1] = panOffsets[iFeat];
      82             :         }
      83        1751 :         return true;
      84             :     }
      85             : 
      86        1684 :     bool SetNull(int iArrowField, int iFeat)
      87             :     {
      88        1684 :         return SetNull(m_out_array->children[iArrowField], iFeat,
      89        1684 :                        m_nMaxBatchSize, true);
      90             :     }
      91             : 
      92         151 :     inline static void SetBoolOn(struct ArrowArray *psArray, int iFeat)
      93             :     {
      94             :         static_cast<uint8_t *>(
      95         151 :             const_cast<void *>(psArray->buffers[1]))[iFeat / 8] |=
      96         151 :             static_cast<uint8_t>(1 << (iFeat % 8));
      97         151 :     }
      98             : 
      99           0 :     inline static void SetInt8(struct ArrowArray *psArray, int iFeat,
     100             :                                int8_t nVal)
     101             :     {
     102           0 :         static_cast<int8_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
     103             :             nVal;
     104           0 :     }
     105             : 
     106           0 :     inline static void SetUInt8(struct ArrowArray *psArray, int iFeat,
     107             :                                 uint8_t nVal)
     108             :     {
     109           0 :         static_cast<uint8_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
     110             :             nVal;
     111           0 :     }
     112             : 
     113         132 :     inline static void SetInt16(struct ArrowArray *psArray, int iFeat,
     114             :                                 int16_t nVal)
     115             :     {
     116         132 :         static_cast<int16_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
     117             :             nVal;
     118         132 :     }
     119             : 
     120             :     inline static void SetUInt16(struct ArrowArray *psArray, int iFeat,
     121             :                                  uint16_t nVal)
     122             :     {
     123             :         static_cast<uint16_t *>(
     124             :             const_cast<void *>(psArray->buffers[1]))[iFeat] = nVal;
     125             :     }
     126             : 
     127        1878 :     inline static void SetInt32(struct ArrowArray *psArray, int iFeat,
     128             :                                 int32_t nVal)
     129             :     {
     130        1878 :         static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
     131             :             nVal;
     132        1878 :     }
     133             : 
     134             :     inline static void SetUInt32(struct ArrowArray *psArray, int iFeat,
     135             :                                  uint32_t nVal)
     136             :     {
     137             :         static_cast<uint32_t *>(
     138             :             const_cast<void *>(psArray->buffers[1]))[iFeat] = nVal;
     139             :     }
     140             : 
     141         201 :     inline static void SetInt64(struct ArrowArray *psArray, int iFeat,
     142             :                                 int64_t nVal)
     143             :     {
     144         201 :         static_cast<int64_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
     145             :             nVal;
     146         201 :     }
     147             : 
     148             :     inline static void SetUInt64(struct ArrowArray *psArray, int iFeat,
     149             :                                  uint64_t nVal)
     150             :     {
     151             :         static_cast<uint64_t *>(
     152             :             const_cast<void *>(psArray->buffers[1]))[iFeat] = nVal;
     153             :     }
     154             : 
     155         111 :     inline static void SetFloat(struct ArrowArray *psArray, int iFeat,
     156             :                                 float fVal)
     157             :     {
     158         111 :         static_cast<float *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
     159             :             fVal;
     160         111 :     }
     161             : 
     162         231 :     inline static void SetDouble(struct ArrowArray *psArray, int iFeat,
     163             :                                  double dfVal)
     164             :     {
     165         231 :         static_cast<double *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
     166             :             dfVal;
     167         231 :     }
     168             : 
     169          91 :     static void SetDate(struct ArrowArray *psArray, int iFeat,
     170             :                         struct tm &brokenDown, const OGRField &ogrField)
     171             :     {
     172          91 :         brokenDown.tm_year = ogrField.Date.Year - 1900;
     173          91 :         brokenDown.tm_mon = ogrField.Date.Month - 1;
     174          91 :         brokenDown.tm_mday = ogrField.Date.Day;
     175          91 :         brokenDown.tm_hour = 0;
     176          91 :         brokenDown.tm_min = 0;
     177          91 :         brokenDown.tm_sec = 0;
     178          91 :         static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
     179          91 :             static_cast<int>(CPLYMDHMSToUnixTime(&brokenDown) / 86400);
     180          91 :     }
     181             : 
     182         115 :     static void SetDateTime(struct ArrowArray *psArray, int iFeat,
     183             :                             struct tm &brokenDown, int nFieldTZFlag,
     184             :                             const OGRField &ogrField)
     185             :     {
     186         115 :         brokenDown.tm_year = ogrField.Date.Year - 1900;
     187         115 :         brokenDown.tm_mon = ogrField.Date.Month - 1;
     188         115 :         brokenDown.tm_mday = ogrField.Date.Day;
     189         115 :         brokenDown.tm_hour = ogrField.Date.Hour;
     190         115 :         brokenDown.tm_min = ogrField.Date.Minute;
     191         115 :         brokenDown.tm_sec = static_cast<int>(ogrField.Date.Second);
     192             :         auto nVal =
     193         115 :             CPLYMDHMSToUnixTime(&brokenDown) * 1000 +
     194         115 :             (static_cast<int>(ogrField.Date.Second * 1000 + 0.5f) % 1000);
     195         115 :         if (nFieldTZFlag >= OGR_TZFLAG_MIXED_TZ &&
     196          95 :             ogrField.Date.TZFlag > OGR_TZFLAG_MIXED_TZ)
     197             :         {
     198             :             // Convert for ogrField.Date.TZFlag to UTC
     199          95 :             const int TZOffset = (ogrField.Date.TZFlag - OGR_TZFLAG_UTC) * 15;
     200          95 :             const int TZOffsetMS = TZOffset * 60 * 1000;
     201          95 :             nVal -= TZOffsetMS;
     202             :         }
     203         115 :         static_cast<int64_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
     204             :             nVal;
     205         115 :     }
     206             : 
     207     1265420 :     static GByte *GetPtrForStringOrBinary(struct ArrowArray *psArray, int iFeat,
     208             :                                           size_t nLen, uint32_t &nMaxAlloc,
     209             :                                           bool bAlignedMalloc)
     210             :     {
     211     1265420 :         auto panOffsets =
     212     1265420 :             static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]));
     213     1265420 :         const uint32_t nCurLength = static_cast<uint32_t>(panOffsets[iFeat]);
     214     1265420 :         if (nLen > nMaxAlloc - nCurLength)
     215             :         {
     216         493 :             constexpr uint32_t INT32_MAX_AS_UINT32 =
     217             :                 static_cast<uint32_t>(std::numeric_limits<int32_t>::max());
     218         493 :             if (nCurLength > INT32_MAX_AS_UINT32 ||
     219         493 :                 nLen > INT32_MAX_AS_UINT32 - nCurLength)
     220             :             {
     221           0 :                 CPLError(CE_Failure, CPLE_AppDefined,
     222             :                          "Too large string or binary content");
     223           0 :                 return nullptr;
     224             :             }
     225         493 :             uint32_t nNewSize = nCurLength + static_cast<uint32_t>(nLen);
     226         493 :             if ((nMaxAlloc >> 31) == 0)
     227             :             {
     228         493 :                 const uint32_t nDoubleSize = 2U * nMaxAlloc;
     229         493 :                 if (nNewSize < nDoubleSize)
     230          11 :                     nNewSize = nDoubleSize;
     231             :             }
     232             :             void *newBuffer;
     233         493 :             if (bAlignedMalloc)
     234             :             {
     235         404 :                 newBuffer = VSI_MALLOC_ALIGNED_AUTO_VERBOSE(nNewSize);
     236         404 :                 if (newBuffer == nullptr)
     237           0 :                     return nullptr;
     238         404 :                 nMaxAlloc = nNewSize;
     239         404 :                 memcpy(newBuffer, psArray->buffers[2], nCurLength);
     240         404 :                 VSIFreeAligned(const_cast<void *>(psArray->buffers[2]));
     241             :             }
     242             :             else
     243             :             {
     244             :                 // coverity[overflow_sink]
     245          89 :                 newBuffer = VSI_REALLOC_VERBOSE(
     246             :                     const_cast<void *>(psArray->buffers[2]), nNewSize);
     247          89 :                 if (newBuffer == nullptr)
     248           0 :                     return nullptr;
     249          89 :                 nMaxAlloc = nNewSize;
     250             :             }
     251         493 :             psArray->buffers[2] = newBuffer;
     252             :         }
     253     1265420 :         GByte *paby =
     254     1265420 :             static_cast<GByte *>(const_cast<void *>(psArray->buffers[2])) +
     255     1265420 :             nCurLength;
     256     1265420 :         panOffsets[iFeat + 1] = panOffsets[iFeat] + static_cast<int32_t>(nLen);
     257     1265420 :         return paby;
     258             :     }
     259             : 
     260     1265360 :     GByte *GetPtrForStringOrBinary(int iArrowField, int iFeat, size_t nLen,
     261             :                                    bool bAlignedMalloc = true)
     262             :     {
     263     1265360 :         auto psArray = m_out_array->children[iArrowField];
     264     1265350 :         return GetPtrForStringOrBinary(psArray, iFeat, nLen,
     265     1265360 :                                        m_anArrowFieldMaxAlloc[iArrowField],
     266     1265360 :                                        bAlignedMalloc);
     267             :     }
     268             : 
     269          31 :     static void SetEmptyStringOrBinary(struct ArrowArray *psArray, int iFeat)
     270             :     {
     271          31 :         auto panOffsets =
     272          31 :             static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]));
     273          31 :         panOffsets[iFeat + 1] = panOffsets[iFeat];
     274          31 :     }
     275             : 
     276         410 :     void Shrink(int nFeatures)
     277             :     {
     278         410 :         if (nFeatures < m_nMaxBatchSize)
     279             :         {
     280         339 :             m_out_array->length = nFeatures;
     281        2666 :             for (int i = 0; i < m_nChildren; i++)
     282             :             {
     283        2327 :                 m_out_array->children[i]->length = nFeatures;
     284             :             }
     285             :         }
     286         410 :     }
     287             : 
     288          26 :     void ClearArray()
     289             :     {
     290          26 :         if (m_out_array->release)
     291          25 :             m_out_array->release(m_out_array);
     292          26 :         memset(m_out_array, 0, sizeof(*m_out_array));
     293          26 :     }
     294             : 
     295             :     static bool FillDict(struct ArrowArray *psChild,
     296             :                          const OGRCodedFieldDomain *poCodedDomain);
     297             : };
     298             : 
     299             : //! @endcond