GDAL

ograrrowarrayhelper.h
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
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
54 OGRArrowArrayHelper(struct ArrowArray *out_array, int nMaxBatchSize);
55
56 static bool SetNull(struct ArrowArray *psArray, int iFeat,
57 int nMaxBatchSize, bool bAlignedMalloc)
58 {
59 ++psArray->null_count;
60 uint8_t *pabyNull =
61 static_cast<uint8_t *>(const_cast<void *>(psArray->buffers[0]));
62 if (psArray->buffers[0] == nullptr)
63 {
64 pabyNull = static_cast<uint8_t *>(
65 bAlignedMalloc
66 ? VSI_MALLOC_ALIGNED_AUTO_VERBOSE((nMaxBatchSize + 7) / 8)
67 : VSI_MALLOC_VERBOSE((nMaxBatchSize + 7) / 8));
68 if (pabyNull == nullptr)
69 {
70 return false;
71 }
72 memset(pabyNull, 0xFF, (nMaxBatchSize + 7) / 8);
73 psArray->buffers[0] = pabyNull;
74 }
75 pabyNull[iFeat / 8] &= static_cast<uint8_t>(~(1 << (iFeat % 8)));
76
77 if (psArray->n_buffers == 3)
78 {
79 auto panOffsets =
80 static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]));
81 panOffsets[iFeat + 1] = panOffsets[iFeat];
82 }
83 return true;
84 }
85
86 bool SetNull(int iArrowField, int iFeat)
87 {
88 return SetNull(m_out_array->children[iArrowField], iFeat,
89 m_nMaxBatchSize, true);
90 }
91
92 inline static void SetBoolOn(struct ArrowArray *psArray, int iFeat)
93 {
94 static_cast<uint8_t *>(
95 const_cast<void *>(psArray->buffers[1]))[iFeat / 8] |=
96 static_cast<uint8_t>(1 << (iFeat % 8));
97 }
98
99 inline static void SetInt8(struct ArrowArray *psArray, int iFeat,
100 int8_t nVal)
101 {
102 static_cast<int8_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
103 nVal;
104 }
105
106 inline static void SetUInt8(struct ArrowArray *psArray, int iFeat,
107 uint8_t nVal)
108 {
109 static_cast<uint8_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
110 nVal;
111 }
112
113 inline static void SetInt16(struct ArrowArray *psArray, int iFeat,
114 int16_t nVal)
115 {
116 static_cast<int16_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
117 nVal;
118 }
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 inline static void SetInt32(struct ArrowArray *psArray, int iFeat,
128 int32_t nVal)
129 {
130 static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
131 nVal;
132 }
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 inline static void SetInt64(struct ArrowArray *psArray, int iFeat,
142 int64_t nVal)
143 {
144 static_cast<int64_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
145 nVal;
146 }
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 inline static void SetFloat(struct ArrowArray *psArray, int iFeat,
156 float fVal)
157 {
158 static_cast<float *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
159 fVal;
160 }
161
162 inline static void SetDouble(struct ArrowArray *psArray, int iFeat,
163 double dfVal)
164 {
165 static_cast<double *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
166 dfVal;
167 }
168
169 static void SetDate(struct ArrowArray *psArray, int iFeat,
170 struct tm &brokenDown, const OGRField &ogrField)
171 {
172 brokenDown.tm_year = ogrField.Date.Year - 1900;
173 brokenDown.tm_mon = ogrField.Date.Month - 1;
174 brokenDown.tm_mday = ogrField.Date.Day;
175 brokenDown.tm_hour = 0;
176 brokenDown.tm_min = 0;
177 brokenDown.tm_sec = 0;
178 static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
179 static_cast<int>(CPLYMDHMSToUnixTime(&brokenDown) / 86400);
180 }
181
182 static void SetDateTime(struct ArrowArray *psArray, int iFeat,
183 struct tm &brokenDown, int nFieldTZFlag,
184 const OGRField &ogrField)
185 {
186 brokenDown.tm_year = ogrField.Date.Year - 1900;
187 brokenDown.tm_mon = ogrField.Date.Month - 1;
188 brokenDown.tm_mday = ogrField.Date.Day;
189 brokenDown.tm_hour = ogrField.Date.Hour;
190 brokenDown.tm_min = ogrField.Date.Minute;
191 brokenDown.tm_sec = static_cast<int>(ogrField.Date.Second);
192 auto nVal =
193 CPLYMDHMSToUnixTime(&brokenDown) * 1000 +
194 (static_cast<int>(ogrField.Date.Second * 1000 + 0.5f) % 1000);
195 if (nFieldTZFlag >= OGR_TZFLAG_MIXED_TZ &&
196 ogrField.Date.TZFlag > OGR_TZFLAG_MIXED_TZ)
197 {
198 // Convert for ogrField.Date.TZFlag to UTC
199 const int TZOffset = (ogrField.Date.TZFlag - OGR_TZFLAG_UTC) * 15;
200 const int TZOffsetMS = TZOffset * 60 * 1000;
201 nVal -= TZOffsetMS;
202 }
203 static_cast<int64_t *>(const_cast<void *>(psArray->buffers[1]))[iFeat] =
204 nVal;
205 }
206
207 static GByte *GetPtrForStringOrBinary(struct ArrowArray *psArray, int iFeat,
208 size_t nLen, uint32_t &nMaxAlloc,
209 bool bAlignedMalloc)
210 {
211 auto panOffsets =
212 static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]));
213 const uint32_t nCurLength = static_cast<uint32_t>(panOffsets[iFeat]);
214 if (nLen > nMaxAlloc - nCurLength)
215 {
216 constexpr uint32_t INT32_MAX_AS_UINT32 =
217 static_cast<uint32_t>(std::numeric_limits<int32_t>::max());
218 if (nCurLength > INT32_MAX_AS_UINT32 ||
219 nLen > INT32_MAX_AS_UINT32 - nCurLength)
220 {
221 CPLError(CE_Failure, CPLE_AppDefined,
222 "Too large string or binary content");
223 return nullptr;
224 }
225 uint32_t nNewSize = nCurLength + static_cast<uint32_t>(nLen);
226 if ((nMaxAlloc >> 31) == 0)
227 {
228 const uint32_t nDoubleSize = 2U * nMaxAlloc;
229 if (nNewSize < nDoubleSize)
230 nNewSize = nDoubleSize;
231 }
232 void *newBuffer;
233 if (bAlignedMalloc)
234 {
235 newBuffer = VSI_MALLOC_ALIGNED_AUTO_VERBOSE(nNewSize);
236 if (newBuffer == nullptr)
237 return nullptr;
238 nMaxAlloc = nNewSize;
239 memcpy(newBuffer, psArray->buffers[2], nCurLength);
240 VSIFreeAligned(const_cast<void *>(psArray->buffers[2]));
241 }
242 else
243 {
244 // coverity[overflow_sink]
245 newBuffer = VSI_REALLOC_VERBOSE(
246 const_cast<void *>(psArray->buffers[2]), nNewSize);
247 if (newBuffer == nullptr)
248 return nullptr;
249 nMaxAlloc = nNewSize;
250 }
251 psArray->buffers[2] = newBuffer;
252 }
253 GByte *paby =
254 static_cast<GByte *>(const_cast<void *>(psArray->buffers[2])) +
255 nCurLength;
256 panOffsets[iFeat + 1] = panOffsets[iFeat] + static_cast<int32_t>(nLen);
257 return paby;
258 }
259
260 GByte *GetPtrForStringOrBinary(int iArrowField, int iFeat, size_t nLen,
261 bool bAlignedMalloc = true)
262 {
263 auto psArray = m_out_array->children[iArrowField];
264 return GetPtrForStringOrBinary(psArray, iFeat, nLen,
265 m_anArrowFieldMaxAlloc[iArrowField],
266 bAlignedMalloc);
267 }
268
269 static void SetEmptyStringOrBinary(struct ArrowArray *psArray, int iFeat)
270 {
271 auto panOffsets =
272 static_cast<int32_t *>(const_cast<void *>(psArray->buffers[1]));
273 panOffsets[iFeat + 1] = panOffsets[iFeat];
274 }
275
276 void Shrink(int nFeatures)
277 {
278 if (nFeatures < m_nMaxBatchSize)
279 {
280 m_out_array->length = nFeatures;
281 for (int i = 0; i < m_nChildren; i++)
282 {
283 m_out_array->children[i]->length = nFeatures;
284 }
285 }
286 }
287
288 void ClearArray()
289 {
290 if (m_out_array->release)
291 m_out_array->release(m_out_array);
292 memset(m_out_array, 0, sizeof(*m_out_array));
293 }
294
295 static bool FillDict(struct ArrowArray *psChild,
296 const OGRCodedFieldDomain *poCodedDomain);
297};
298
CPLStringList
String list class designed around our use of C "char**" string lists.
Definition cpl_string.h:458
GDALDataset
A set of associated raster bands, usually from one file.
Definition gdal_dataset.h:76
OGRCodedFieldDomain
Definition of a coded / enumerated field domain.
Definition ogr_feature.h:1792
OGRFeatureDefn
Definition of a feature class or feature layer.
Definition ogr_feature.h:504
CPLE_AppDefined
#define CPLE_AppDefined
Application defined error.
Definition cpl_error.h:85
GByte
unsigned char GByte
Unsigned byte type.
Definition cpl_port.h:175
VSI_REALLOC_VERBOSE
#define VSI_REALLOC_VERBOSE(pOldPtr, nNewSize)
VSI_REALLOC_VERBOSE.
Definition cpl_vsi.h:368
VSI_MALLOC_ALIGNED_AUTO_VERBOSE
#define VSI_MALLOC_ALIGNED_AUTO_VERBOSE(size)
VSIMallocAlignedAutoVerbose() with FILE and LINE reporting.
Definition cpl_vsi.h:311
VSI_MALLOC_VERBOSE
#define VSI_MALLOC_VERBOSE(size)
VSI_MALLOC_VERBOSE.
Definition cpl_vsi.h:338
VSIFreeAligned
void VSIFreeAligned(void *ptr)
Free a buffer allocated with VSIMallocAligned().
Definition cpl_vsisimple.cpp:981
OGR_TZFLAG_UTC
#define OGR_TZFLAG_UTC
Time zone flag indicating UTC.
Definition ogr_core.h:875
OGR_TZFLAG_MIXED_TZ
#define OGR_TZFLAG_MIXED_TZ
Time zone flag only returned by OGRFieldDefn::GetTZFlag() to indicate that all values in the field ha...
Definition ogr_core.h:867
ogrsf_frmts.h
Classes related to registration of format support, and opening datasets.
OGRField
OGRFeature field attribute value union.
Definition ogr_core.h:886