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