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 : #include "ograrrowarrayhelper.h"
14 : #include "ogr_p.h"
15 :
16 : #include <limits>
17 :
18 : //! @cond Doxygen_Suppress
19 :
20 : /************************************************************************/
21 : /* GetMemLimit() */
22 : /************************************************************************/
23 :
24 1689 : /*static*/ uint32_t OGRArrowArrayHelper::GetMemLimit()
25 : {
26 1689 : uint32_t nMemLimit =
27 : static_cast<uint32_t>(std::numeric_limits<int32_t>::max());
28 : // Just for tests
29 : const char *pszOGR_ARROW_MEM_LIMIT =
30 1689 : CPLGetConfigOption("OGR_ARROW_MEM_LIMIT", nullptr);
31 1689 : if (pszOGR_ARROW_MEM_LIMIT)
32 131 : nMemLimit = atoi(pszOGR_ARROW_MEM_LIMIT);
33 : else
34 : {
35 1558 : const auto nUsableRAM = CPLGetUsablePhysicalRAM();
36 1558 : if (nUsableRAM > 0 && static_cast<uint64_t>(nUsableRAM / 4) < nMemLimit)
37 0 : nMemLimit = static_cast<uint32_t>(nUsableRAM / 4);
38 : }
39 1689 : return nMemLimit;
40 : }
41 :
42 : /************************************************************************/
43 : /* GetMaxFeaturesInBatch() */
44 : /************************************************************************/
45 :
46 : /* static */
47 552 : int OGRArrowArrayHelper::GetMaxFeaturesInBatch(
48 : const CPLStringList &aosArrowArrayStreamOptions)
49 : {
50 552 : int l_nMaxBatchSize = atoi(aosArrowArrayStreamOptions.FetchNameValueDef(
51 : "MAX_FEATURES_IN_BATCH", "65536"));
52 552 : if (l_nMaxBatchSize <= 0)
53 0 : l_nMaxBatchSize = 1;
54 552 : if (l_nMaxBatchSize > INT_MAX - 1)
55 0 : l_nMaxBatchSize = INT_MAX - 1;
56 :
57 552 : return l_nMaxBatchSize;
58 : }
59 :
60 : /************************************************************************/
61 : /* OGRArrowArrayHelper() */
62 : /************************************************************************/
63 :
64 393 : OGRArrowArrayHelper::OGRArrowArrayHelper(
65 : GDALDataset *poDS, OGRFeatureDefn *poFeatureDefn,
66 : const CPLStringList &aosArrowArrayStreamOptions,
67 393 : struct ArrowArray *out_array)
68 393 : : m_bIncludeFID(CPLTestBool(
69 : aosArrowArrayStreamOptions.FetchNameValueDef("INCLUDE_FID", "YES"))),
70 786 : m_nMaxBatchSize(GetMaxFeaturesInBatch(aosArrowArrayStreamOptions)),
71 786 : m_nFieldCount(poFeatureDefn->GetFieldCount()),
72 786 : m_nGeomFieldCount(poFeatureDefn->GetGeomFieldCount()),
73 393 : m_out_array(out_array)
74 : {
75 393 : memset(out_array, 0, sizeof(*out_array));
76 :
77 393 : m_mapOGRFieldToArrowField.resize(m_nFieldCount, -1);
78 393 : m_mapOGRGeomFieldToArrowField.resize(m_nGeomFieldCount, -1);
79 393 : m_abNullableFields.resize(m_nFieldCount);
80 393 : m_anTZFlags.resize(m_nFieldCount);
81 393 : int nTZFlagOverride = -1;
82 : const char *pszTZOverride =
83 393 : aosArrowArrayStreamOptions.FetchNameValue("TIMEZONE");
84 393 : if (pszTZOverride)
85 : {
86 345 : if (EQUAL(pszTZOverride, "unknown") || EQUAL(pszTZOverride, ""))
87 : {
88 0 : nTZFlagOverride = OGR_TZFLAG_UNKNOWN;
89 : }
90 : else
91 : {
92 : // we don't really care about the actual timezone, since we
93 : // will convert OGRField::Date to UTC in all cases
94 345 : nTZFlagOverride = OGR_TZFLAG_UTC;
95 : }
96 : }
97 :
98 393 : if (m_bIncludeFID)
99 : {
100 385 : m_nChildren++;
101 : }
102 : // cppcheck-suppress knownConditionTrueFalse
103 3075 : for (int i = 0; i < m_nFieldCount; i++)
104 : {
105 2682 : const auto poFieldDefn = poFeatureDefn->GetFieldDefn(i);
106 2682 : m_abNullableFields[i] = CPL_TO_BOOL(poFieldDefn->IsNullable());
107 2682 : m_anTZFlags[i] =
108 2682 : nTZFlagOverride >= 0 ? nTZFlagOverride : poFieldDefn->GetTZFlag();
109 2682 : if (!poFieldDefn->IsIgnored())
110 : {
111 2647 : m_mapOGRFieldToArrowField[i] = m_nChildren;
112 2647 : m_nChildren++;
113 : }
114 : }
115 : // cppcheck-suppress knownConditionTrueFalse
116 786 : for (int i = 0; i < m_nGeomFieldCount; i++)
117 : {
118 393 : if (!poFeatureDefn->GetGeomFieldDefn(i)->IsIgnored())
119 : {
120 380 : m_mapOGRGeomFieldToArrowField[i] = m_nChildren;
121 380 : m_nChildren++;
122 : }
123 : }
124 :
125 393 : m_anArrowFieldMaxAlloc.resize(m_nChildren);
126 :
127 393 : out_array->release = OGRLayer::ReleaseArray;
128 :
129 393 : out_array->length = m_nMaxBatchSize;
130 393 : out_array->null_count = 0;
131 :
132 393 : out_array->n_children = m_nChildren;
133 393 : out_array->children = static_cast<struct ArrowArray **>(
134 393 : CPLCalloc(m_nChildren, sizeof(struct ArrowArray *)));
135 393 : out_array->release = OGRLayer::ReleaseArray;
136 393 : out_array->n_buffers = 1;
137 393 : out_array->buffers =
138 393 : static_cast<const void **>(CPLCalloc(1, sizeof(void *)));
139 :
140 : // Allocate buffers
141 :
142 393 : if (m_bIncludeFID)
143 : {
144 770 : out_array->children[0] = static_cast<struct ArrowArray *>(
145 385 : CPLCalloc(1, sizeof(struct ArrowArray)));
146 385 : auto psChild = out_array->children[0];
147 385 : psChild->release = OGRLayer::ReleaseArray;
148 385 : psChild->length = m_nMaxBatchSize;
149 385 : psChild->n_buffers = 2;
150 385 : psChild->buffers =
151 385 : static_cast<const void **>(CPLCalloc(2, sizeof(void *)));
152 385 : m_panFIDValues = static_cast<int64_t *>(
153 385 : VSI_MALLOC_ALIGNED_AUTO_VERBOSE(sizeof(int64_t) * m_nMaxBatchSize));
154 385 : if (m_panFIDValues == nullptr)
155 0 : goto error;
156 385 : psChild->buffers[1] = m_panFIDValues;
157 : }
158 :
159 : // cppcheck-suppress knownConditionTrueFalse
160 3075 : for (int i = 0; i < m_nFieldCount; i++)
161 : {
162 2682 : const int iArrowField = m_mapOGRFieldToArrowField[i];
163 2682 : if (iArrowField >= 0)
164 : {
165 2647 : const auto poFieldDefn = poFeatureDefn->GetFieldDefn(i);
166 5294 : out_array->children[iArrowField] = static_cast<struct ArrowArray *>(
167 2647 : CPLCalloc(1, sizeof(struct ArrowArray)));
168 2647 : auto psChild = out_array->children[iArrowField];
169 :
170 2647 : psChild->release = OGRLayer::ReleaseArray;
171 2647 : psChild->length = m_nMaxBatchSize;
172 2647 : const auto eSubType = poFieldDefn->GetSubType();
173 2647 : size_t nEltSize = 0;
174 2647 : switch (poFieldDefn->GetType())
175 : {
176 1705 : case OFTInteger:
177 : {
178 1705 : if (eSubType == OFSTBoolean)
179 : {
180 65 : nEltSize = sizeof(uint8_t);
181 : }
182 1640 : else if (eSubType == OFSTInt16)
183 : {
184 61 : nEltSize = sizeof(int16_t);
185 : }
186 : else
187 : {
188 1579 : nEltSize = sizeof(int32_t);
189 : }
190 :
191 1705 : const auto &osDomainName = poFieldDefn->GetDomainName();
192 1705 : if (!osDomainName.empty() && poDS != nullptr)
193 : {
194 : const auto poFieldDomain =
195 24 : poDS->GetFieldDomain(osDomainName);
196 48 : if (poFieldDomain &&
197 24 : poFieldDomain->GetDomainType() == OFDT_CODED)
198 : {
199 24 : const OGRCodedFieldDomain *poCodedDomain =
200 : static_cast<const OGRCodedFieldDomain *>(
201 : poFieldDomain);
202 24 : FillDict(psChild, poCodedDomain);
203 : }
204 : }
205 :
206 1705 : break;
207 : }
208 71 : case OFTInteger64:
209 : {
210 71 : nEltSize = sizeof(int64_t);
211 71 : break;
212 : }
213 141 : case OFTReal:
214 : {
215 141 : if (eSubType == OFSTFloat32)
216 : {
217 64 : nEltSize = sizeof(float);
218 : }
219 : else
220 : {
221 77 : nEltSize = sizeof(double);
222 : }
223 141 : break;
224 : }
225 628 : case OFTString:
226 : case OFTBinary:
227 : {
228 628 : psChild->n_buffers = 3;
229 628 : psChild->buffers = static_cast<const void **>(
230 628 : CPLCalloc(3, sizeof(void *)));
231 628 : psChild->buffers[1] = VSI_MALLOC_ALIGNED_AUTO_VERBOSE(
232 : sizeof(uint32_t) * (1 + m_nMaxBatchSize));
233 628 : if (psChild->buffers[1] == nullptr)
234 0 : goto error;
235 628 : memset(const_cast<void *>(psChild->buffers[1]), 0,
236 628 : sizeof(uint32_t) * (1 + m_nMaxBatchSize));
237 628 : constexpr size_t DEFAULT_STRING_SIZE = 10;
238 1256 : m_anArrowFieldMaxAlloc[iArrowField] =
239 628 : DEFAULT_STRING_SIZE * m_nMaxBatchSize;
240 628 : psChild->buffers[2] = VSI_MALLOC_ALIGNED_AUTO_VERBOSE(
241 : m_anArrowFieldMaxAlloc[iArrowField]);
242 628 : if (psChild->buffers[2] == nullptr)
243 0 : goto error;
244 628 : break;
245 : }
246 :
247 36 : case OFTDate:
248 : {
249 36 : nEltSize = sizeof(int32_t);
250 36 : break;
251 : }
252 :
253 0 : case OFTTime:
254 : {
255 0 : nEltSize = sizeof(int32_t);
256 0 : break;
257 : }
258 :
259 66 : case OFTDateTime:
260 : {
261 66 : nEltSize = sizeof(int64_t);
262 66 : break;
263 : }
264 :
265 0 : default:
266 0 : break;
267 : }
268 :
269 2647 : if (nEltSize != 0)
270 : {
271 2019 : psChild->n_buffers = 2;
272 2019 : psChild->buffers =
273 2019 : static_cast<const void **>(CPLCalloc(2, sizeof(void *)));
274 4038 : psChild->buffers[1] =
275 2019 : VSI_MALLOC_ALIGNED_AUTO_VERBOSE(nEltSize * m_nMaxBatchSize);
276 2019 : if (psChild->buffers[1] == nullptr)
277 0 : goto error;
278 2019 : memset(const_cast<void *>(psChild->buffers[1]), 0,
279 2019 : nEltSize * m_nMaxBatchSize);
280 : }
281 : }
282 : }
283 :
284 : // cppcheck-suppress knownConditionTrueFalse
285 786 : for (int i = 0; i < m_nGeomFieldCount; i++)
286 : {
287 393 : const int iArrowField = m_mapOGRGeomFieldToArrowField[i];
288 393 : if (iArrowField >= 0)
289 : {
290 760 : out_array->children[iArrowField] = static_cast<struct ArrowArray *>(
291 380 : CPLCalloc(1, sizeof(struct ArrowArray)));
292 380 : auto psChild = out_array->children[iArrowField];
293 :
294 380 : psChild->release = OGRLayer::ReleaseArray;
295 380 : psChild->length = m_nMaxBatchSize;
296 :
297 380 : psChild->n_buffers = 3;
298 380 : psChild->buffers =
299 380 : static_cast<const void **>(CPLCalloc(3, sizeof(void *)));
300 380 : psChild->buffers[1] = VSI_MALLOC_ALIGNED_AUTO_VERBOSE(
301 : sizeof(uint32_t) * (1 + m_nMaxBatchSize));
302 380 : if (psChild->buffers[1] == nullptr)
303 0 : goto error;
304 380 : memset(const_cast<void *>(psChild->buffers[1]), 0,
305 380 : sizeof(uint32_t) * (1 + m_nMaxBatchSize));
306 380 : constexpr size_t DEFAULT_WKB_SIZE = 100;
307 760 : m_anArrowFieldMaxAlloc[iArrowField] =
308 380 : DEFAULT_WKB_SIZE * m_nMaxBatchSize;
309 380 : psChild->buffers[2] = VSI_MALLOC_ALIGNED_AUTO_VERBOSE(
310 : m_anArrowFieldMaxAlloc[iArrowField]);
311 380 : if (psChild->buffers[2] == nullptr)
312 0 : goto error;
313 : }
314 : }
315 :
316 393 : return;
317 :
318 0 : error:
319 0 : out_array->release(out_array);
320 0 : memset(out_array, 0, sizeof(*out_array));
321 : }
322 :
323 : /************************************************************************/
324 : /* FillDict() */
325 : /************************************************************************/
326 :
327 : /* static */
328 36 : bool OGRArrowArrayHelper::FillDict(struct ArrowArray *psChild,
329 : const OGRCodedFieldDomain *poCodedDomain)
330 : {
331 36 : int nLastCode = -1;
332 36 : uint32_t nCountChars = 0;
333 36 : int nCountNull = 0;
334 36 : for (const OGRCodedValue *psIter = poCodedDomain->GetEnumeration();
335 117 : psIter->pszCode; ++psIter)
336 : {
337 81 : if (CPLGetValueType(psIter->pszCode) != CPL_VALUE_INTEGER)
338 : {
339 0 : return false;
340 : }
341 81 : int nCode = atoi(psIter->pszCode);
342 81 : if (nCode <= nLastCode || nCode - nLastCode > 100)
343 : {
344 0 : return false;
345 : }
346 112 : for (int i = nLastCode + 1; i < nCode; ++i)
347 : {
348 31 : nCountNull++;
349 : }
350 81 : if (psIter->pszValue)
351 : {
352 50 : const size_t nLen = strlen(psIter->pszValue);
353 50 : if (nLen > std::numeric_limits<uint32_t>::max() - nCountChars)
354 0 : return false;
355 50 : nCountChars += static_cast<uint32_t>(nLen);
356 : }
357 : else
358 : {
359 31 : nCountNull++;
360 : }
361 81 : nLastCode = nCode;
362 : }
363 36 : const int nLength = 1 + nLastCode;
364 :
365 : auto psDict = static_cast<struct ArrowArray *>(
366 36 : CPLCalloc(1, sizeof(struct ArrowArray)));
367 36 : psChild->dictionary = psDict;
368 :
369 36 : psDict->release = OGRLayer::ReleaseArray;
370 36 : psDict->length = nLength;
371 36 : psDict->n_buffers = 3;
372 36 : psDict->buffers = static_cast<const void **>(CPLCalloc(3, sizeof(void *)));
373 36 : psDict->null_count = nCountNull;
374 36 : uint8_t *pabyNull = nullptr;
375 36 : if (nCountNull)
376 : {
377 : pabyNull = static_cast<uint8_t *>(
378 31 : VSI_MALLOC_ALIGNED_AUTO_VERBOSE((nLength + 7) / 8));
379 31 : if (pabyNull == nullptr)
380 : {
381 0 : psDict->release(psDict);
382 0 : CPLFree(psDict);
383 0 : psChild->dictionary = nullptr;
384 0 : return false;
385 : }
386 31 : memset(pabyNull, 0xFF, (nLength + 7) / 8);
387 31 : psDict->buffers[0] = pabyNull;
388 : }
389 :
390 : uint32_t *panOffsets = static_cast<uint32_t *>(
391 36 : VSI_MALLOC_ALIGNED_AUTO_VERBOSE(sizeof(uint32_t) * (1 + nLength)));
392 36 : if (panOffsets == nullptr)
393 : {
394 0 : psDict->release(psDict);
395 0 : CPLFree(psDict);
396 0 : psChild->dictionary = nullptr;
397 0 : return false;
398 : }
399 36 : psDict->buffers[1] = panOffsets;
400 :
401 : char *pachValues =
402 36 : static_cast<char *>(VSI_MALLOC_ALIGNED_AUTO_VERBOSE(nCountChars));
403 36 : if (pachValues == nullptr)
404 : {
405 0 : psDict->release(psDict);
406 0 : CPLFree(psDict);
407 0 : psChild->dictionary = nullptr;
408 0 : return false;
409 : }
410 36 : psDict->buffers[2] = pachValues;
411 :
412 36 : nLastCode = -1;
413 36 : uint32_t nOffset = 0;
414 36 : for (const OGRCodedValue *psIter = poCodedDomain->GetEnumeration();
415 117 : psIter->pszCode; ++psIter)
416 : {
417 81 : if (CPLGetValueType(psIter->pszCode) != CPL_VALUE_INTEGER)
418 : {
419 0 : psDict->release(psDict);
420 0 : CPLFree(psDict);
421 0 : psChild->dictionary = nullptr;
422 0 : return false;
423 : }
424 81 : int nCode = atoi(psIter->pszCode);
425 81 : if (nCode <= nLastCode || nCode - nLastCode > 100)
426 : {
427 0 : psDict->release(psDict);
428 0 : CPLFree(psDict);
429 0 : psChild->dictionary = nullptr;
430 0 : return false;
431 : }
432 112 : for (int i = nLastCode + 1; i < nCode; ++i)
433 : {
434 31 : panOffsets[i] = nOffset;
435 31 : if (pabyNull)
436 31 : pabyNull[i / 8] &= static_cast<uint8_t>(~(1 << (i % 8)));
437 : }
438 81 : panOffsets[nCode] = nOffset;
439 81 : if (psIter->pszValue)
440 : {
441 50 : const size_t nLen = strlen(psIter->pszValue);
442 50 : memcpy(pachValues + nOffset, psIter->pszValue, nLen);
443 50 : nOffset += static_cast<uint32_t>(nLen);
444 : }
445 31 : else if (pabyNull)
446 : {
447 31 : pabyNull[nCode / 8] &= static_cast<uint8_t>(~(1 << (nCode % 8)));
448 : }
449 81 : nLastCode = nCode;
450 : }
451 36 : panOffsets[nLength] = nOffset;
452 :
453 36 : return true;
454 : }
455 :
456 : //! @endcond
|
