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