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