39 #include "ogr_recordbatch.h"
41 class CPL_DLL OGRArrowArrayHelper
43 OGRArrowArrayHelper(
const OGRArrowArrayHelper &) =
delete;
44 OGRArrowArrayHelper &operator=(
const OGRArrowArrayHelper &) =
delete;
47 bool bIncludeFID =
false;
48 int nMaxBatchSize = 0;
51 int nGeomFieldCount = 0;
52 std::vector<int> mapOGRFieldToArrowField{};
53 std::vector<int> mapOGRGeomFieldToArrowField{};
54 std::vector<bool> abNullableFields{};
55 std::vector<uint32_t> anArrowFieldMaxAlloc{};
56 int64_t *panFIDValues =
nullptr;
57 struct ArrowArray *m_out_array =
nullptr;
60 GetMaxFeaturesInBatch(
const CPLStringList &aosArrowArrayStreamOptions);
64 struct ArrowArray *out_array);
66 bool SetNull(
int iArrowField,
int iFeat)
68 auto psArray = m_out_array->children[iArrowField];
69 ++psArray->null_count;
71 static_cast<uint8_t *
>(
const_cast<void *
>(psArray->buffers[0]));
72 if (psArray->buffers[0] ==
nullptr)
74 pabyNull =
static_cast<uint8_t *
>(
76 if (pabyNull ==
nullptr)
80 memset(pabyNull, 0xFF, (nMaxBatchSize + 7) / 8);
81 psArray->buffers[0] = pabyNull;
83 pabyNull[iFeat / 8] &=
static_cast<uint8_t
>(~(1 << (iFeat % 8)));
85 if (psArray->n_buffers == 3)
88 static_cast<int32_t *
>(
const_cast<void *
>(psArray->buffers[1]));
89 panOffsets[iFeat + 1] = panOffsets[iFeat];
94 inline static void SetBoolOn(
struct ArrowArray *psArray,
int iFeat)
96 static_cast<uint8_t *
>(
97 const_cast<void *
>(psArray->buffers[1]))[iFeat / 8] |=
98 static_cast<uint8_t
>(1 << (iFeat / 8));
101 inline static void SetInt8(
struct ArrowArray *psArray,
int iFeat,
104 static_cast<int8_t *
>(
const_cast<void *
>(psArray->buffers[1]))[iFeat] =
108 inline static void SetUInt8(
struct ArrowArray *psArray,
int iFeat,
111 static_cast<uint8_t *
>(
const_cast<void *
>(psArray->buffers[1]))[iFeat] =
115 inline static void SetInt16(
struct ArrowArray *psArray,
int iFeat,
118 static_cast<int16_t *
>(
const_cast<void *
>(psArray->buffers[1]))[iFeat] =
122 inline static void SetUInt16(
struct ArrowArray *psArray,
int iFeat,
125 static_cast<uint16_t *
>(
126 const_cast<void *
>(psArray->buffers[1]))[iFeat] = nVal;
129 inline static void SetInt32(
struct ArrowArray *psArray,
int iFeat,
132 static_cast<int32_t *
>(
const_cast<void *
>(psArray->buffers[1]))[iFeat] =
136 inline static void SetUInt32(
struct ArrowArray *psArray,
int iFeat,
139 static_cast<uint32_t *
>(
140 const_cast<void *
>(psArray->buffers[1]))[iFeat] = nVal;
143 inline static void SetInt64(
struct ArrowArray *psArray,
int iFeat,
146 static_cast<int64_t *
>(
const_cast<void *
>(psArray->buffers[1]))[iFeat] =
150 inline static void SetUInt64(
struct ArrowArray *psArray,
int iFeat,
153 static_cast<uint64_t *
>(
154 const_cast<void *
>(psArray->buffers[1]))[iFeat] = nVal;
157 inline static void SetFloat(
struct ArrowArray *psArray,
int iFeat,
160 static_cast<float *
>(
const_cast<void *
>(psArray->buffers[1]))[iFeat] =
164 inline static void SetDouble(
struct ArrowArray *psArray,
int iFeat,
167 static_cast<double *
>(
const_cast<void *
>(psArray->buffers[1]))[iFeat] =
171 static void SetDate(
struct ArrowArray *psArray,
int iFeat,
172 struct tm &brokenDown,
const OGRField &ogrField)
174 brokenDown.tm_year = ogrField.Date.Year - 1900;
175 brokenDown.tm_mon = ogrField.Date.Month - 1;
176 brokenDown.tm_mday = ogrField.Date.Day;
177 brokenDown.tm_hour = 0;
178 brokenDown.tm_min = 0;
179 brokenDown.tm_sec = 0;
180 static_cast<int32_t *
>(
const_cast<void *
>(psArray->buffers[1]))[iFeat] =
181 static_cast<int>(CPLYMDHMSToUnixTime(&brokenDown) / 86400);
184 static void SetDateTime(
struct ArrowArray *psArray,
int iFeat,
185 struct tm &brokenDown,
const OGRField &ogrField)
187 brokenDown.tm_year = ogrField.Date.Year - 1900;
188 brokenDown.tm_mon = ogrField.Date.Month - 1;
189 brokenDown.tm_mday = ogrField.Date.Day;
190 brokenDown.tm_hour = ogrField.Date.Hour;
191 brokenDown.tm_min = ogrField.Date.Minute;
192 brokenDown.tm_sec =
static_cast<int>(ogrField.Date.Second);
193 static_cast<int64_t *
>(
const_cast<void *
>(psArray->buffers[1]))[iFeat] =
194 CPLYMDHMSToUnixTime(&brokenDown) * 1000 +
195 (
static_cast<int>(ogrField.Date.Second * 1000 + 0.5) % 1000);
198 GByte *GetPtrForStringOrBinary(
int iArrowField,
int iFeat,
size_t nLen)
200 auto psArray = m_out_array->children[iArrowField];
202 static_cast<int32_t *
>(
const_cast<void *
>(psArray->buffers[1]));
203 const uint32_t nCurLength =
static_cast<uint32_t
>(panOffsets[iFeat]);
204 if (nLen > anArrowFieldMaxAlloc[iArrowField] - nCurLength)
207 static_cast<uint32_t
>(std::numeric_limits<int32_t>::max()) -
211 "Too large string or binary content");
214 uint32_t nNewSize = nCurLength +
static_cast<uint32_t
>(nLen);
215 if ((anArrowFieldMaxAlloc[iArrowField] >> 31) == 0)
217 const uint32_t nDoubleSize =
218 2U * anArrowFieldMaxAlloc[iArrowField];
219 if (nNewSize < nDoubleSize)
220 nNewSize = nDoubleSize;
223 if (newBuffer ==
nullptr)
225 anArrowFieldMaxAlloc[iArrowField] = nNewSize;
226 memcpy(newBuffer, psArray->buffers[2], nCurLength);
228 psArray->buffers[2] = newBuffer;
231 static_cast<GByte *
>(
const_cast<void *
>(psArray->buffers[2])) +
233 panOffsets[iFeat + 1] = panOffsets[iFeat] +
static_cast<int32_t
>(nLen);
237 static void SetEmptyStringOrBinary(
struct ArrowArray *psArray,
int iFeat)
240 static_cast<int32_t *
>(
const_cast<void *
>(psArray->buffers[1]));
241 panOffsets[iFeat + 1] = panOffsets[iFeat];
244 void Shrink(
int nFeatures)
246 if (nFeatures < nMaxBatchSize)
248 m_out_array->length = nFeatures;
249 for (
int i = 0; i < nChildren; i++)
251 m_out_array->children[i]->length = nFeatures;
258 m_out_array->release(m_out_array);
259 memset(m_out_array, 0,
sizeof(*m_out_array));
262 static bool FillDict(
struct ArrowArray *psChild,