38 #include "ogr_recordbatch.h"
40 class CPL_DLL OGRArrowArrayHelper
42 OGRArrowArrayHelper(
const OGRArrowArrayHelper&) =
delete;
43 OGRArrowArrayHelper& operator= (
const OGRArrowArrayHelper&) =
delete;
46 bool bIncludeFID =
false;
47 int nMaxBatchSize = 0;
50 int nGeomFieldCount = 0;
51 std::vector<int> mapOGRFieldToArrowField{};
52 std::vector<int> mapOGRGeomFieldToArrowField{};
53 std::vector<bool> abNullableFields{};
54 std::vector<uint32_t> anArrowFieldMaxAlloc{};
55 int64_t* panFIDValues =
nullptr;
56 struct ArrowArray* m_out_array =
nullptr;
58 static int GetMaxFeaturesInBatch(
const CPLStringList& aosArrowArrayStreamOptions);
63 struct ArrowArray* out_array);
65 bool SetNull(
int iArrowField,
int iFeat)
67 auto psArray = m_out_array->children[iArrowField];
68 ++psArray->null_count;
69 uint8_t* pabyNull =
static_cast<uint8_t*
>(
const_cast<void*
>(psArray->buffers[0]));
70 if( psArray->buffers[0] ==
nullptr )
73 if( pabyNull ==
nullptr )
77 memset(pabyNull, 0xFF, (nMaxBatchSize + 7) / 8);
78 psArray->buffers[0] = pabyNull;
80 pabyNull[iFeat / 8] &=
static_cast<uint8_t
>(~(1 << (iFeat % 8)));
82 if( psArray->n_buffers == 3 )
84 auto panOffsets =
static_cast<int32_t*
>(
const_cast<void*
>(psArray->buffers[1]));
85 panOffsets[iFeat+1] = panOffsets[iFeat];
90 inline static void SetBoolOn(
struct ArrowArray* psArray,
int iFeat)
92 static_cast<uint8_t*
>(
const_cast<void*
>(
93 psArray->buffers[1]))[iFeat / 8] |=
static_cast<uint8_t
>(1 << (iFeat / 8));
96 inline static void SetInt8(
struct ArrowArray* psArray,
int iFeat, int8_t nVal)
98 static_cast<int8_t*
>(
const_cast<void*
>(psArray->buffers[1]))[iFeat] = nVal;
101 inline static void SetUInt8(
struct ArrowArray* psArray,
int iFeat, uint8_t nVal)
103 static_cast<uint8_t*
>(
const_cast<void*
>(psArray->buffers[1]))[iFeat] = nVal;
106 inline static void SetInt16(
struct ArrowArray* psArray,
int iFeat, int16_t nVal)
108 static_cast<int16_t*
>(
const_cast<void*
>(psArray->buffers[1]))[iFeat] = nVal;
111 inline static void SetUInt16(
struct ArrowArray* psArray,
int iFeat, uint16_t nVal)
113 static_cast<uint16_t*
>(
const_cast<void*
>(psArray->buffers[1]))[iFeat] = nVal;
116 inline static void SetInt32(
struct ArrowArray* psArray,
int iFeat, int32_t nVal)
118 static_cast<int32_t*
>(
const_cast<void*
>(psArray->buffers[1]))[iFeat] = nVal;
121 inline static void SetUInt32(
struct ArrowArray* psArray,
int iFeat, uint32_t nVal)
123 static_cast<uint32_t*
>(
const_cast<void*
>(psArray->buffers[1]))[iFeat] = nVal;
126 inline static void SetInt64(
struct ArrowArray* psArray,
int iFeat, int64_t nVal)
128 static_cast<int64_t*
>(
const_cast<void*
>(psArray->buffers[1]))[iFeat] = nVal;
131 inline static void SetUInt64(
struct ArrowArray* psArray,
int iFeat, uint64_t nVal)
133 static_cast<uint64_t*
>(
const_cast<void*
>(psArray->buffers[1]))[iFeat] = nVal;
136 inline static void SetFloat(
struct ArrowArray* psArray,
int iFeat,
float fVal)
138 static_cast<float*
>(
const_cast<void*
>(psArray->buffers[1]))[iFeat] = fVal;
141 inline static void SetDouble(
struct ArrowArray* psArray,
int iFeat,
double dfVal)
143 static_cast<double*
>(
const_cast<void*
>(psArray->buffers[1]))[iFeat] = dfVal;
147 void SetDate(
struct ArrowArray* psArray,
int iFeat,
148 struct tm& brokenDown,
const OGRField& ogrField)
150 brokenDown.tm_year = ogrField.Date.Year - 1900;
151 brokenDown.tm_mon = ogrField.Date.Month - 1;
152 brokenDown.tm_mday = ogrField.Date.Day;
153 brokenDown.tm_hour = ogrField.Date.Hour;
154 brokenDown.tm_min = ogrField.Date.Minute;
155 brokenDown.tm_sec =
static_cast<int>(ogrField.Date.Second);
156 static_cast<int32_t*
>(
const_cast<void*
>(
157 psArray->buffers[1]))[iFeat] =
158 static_cast<int>((CPLYMDHMSToUnixTime(&brokenDown) + 36200) / 86400);
162 void SetDateTime(
struct ArrowArray* psArray,
int iFeat,
163 struct tm& brokenDown,
const OGRField& ogrField)
165 brokenDown.tm_year = ogrField.Date.Year - 1900;
166 brokenDown.tm_mon = ogrField.Date.Month - 1;
167 brokenDown.tm_mday = ogrField.Date.Day;
168 brokenDown.tm_hour = ogrField.Date.Hour;
169 brokenDown.tm_min = ogrField.Date.Minute;
170 brokenDown.tm_sec =
static_cast<int>(ogrField.Date.Second);
171 static_cast<int64_t*
>(
const_cast<void*
>(
172 psArray->buffers[1]))[iFeat] =
173 CPLYMDHMSToUnixTime(&brokenDown) * 1000 +
174 (
static_cast<int>(ogrField.Date.Second * 1000 + 0.5) % 1000);
177 GByte* GetPtrForStringOrBinary(
int iArrowField,
int iFeat,
size_t nLen)
179 auto psArray = m_out_array->children[iArrowField];
180 auto panOffsets =
static_cast<int32_t*
>(
const_cast<void*
>(psArray->buffers[1]));
181 const uint32_t nCurLength =
static_cast<uint32_t
>(panOffsets[iFeat]);
182 if( nLen > anArrowFieldMaxAlloc[iArrowField] - nCurLength )
184 if( nLen >
static_cast<uint32_t
>(std::numeric_limits<int32_t>::max()) - nCurLength )
189 uint32_t nNewSize = nCurLength +
static_cast<uint32_t
>(nLen);
190 if( (anArrowFieldMaxAlloc[iArrowField] >> 31) == 0 )
192 const uint32_t nDoubleSize = 2U * anArrowFieldMaxAlloc[iArrowField];
193 if( nNewSize < nDoubleSize )
194 nNewSize = nDoubleSize;
197 if( newBuffer ==
nullptr )
199 anArrowFieldMaxAlloc[iArrowField] = nNewSize;
200 memcpy(newBuffer, psArray->buffers[2], nCurLength);
202 psArray->buffers[2] = newBuffer;
204 GByte* paby =
static_cast<GByte*
>(
const_cast<void*
>(psArray->buffers[2])) + nCurLength;
205 panOffsets[iFeat+1] = panOffsets[iFeat] +
static_cast<int32_t
>(nLen);
210 void SetEmptyStringOrBinary(
struct ArrowArray* psArray,
int iFeat)
212 auto panOffsets =
static_cast<int32_t*
>(
const_cast<void*
>(psArray->buffers[1]));
213 panOffsets[iFeat+1] = panOffsets[iFeat];
216 void Shrink(
int nFeatures)
218 if( nFeatures < nMaxBatchSize )
220 m_out_array->length = nFeatures;
221 for(
int i = 0; i < nChildren; i++ )
223 m_out_array->children[i]->length = nFeatures;
230 m_out_array->release(m_out_array);
231 memset(m_out_array, 0,
sizeof(*m_out_array));
234 static bool FillDict(
struct ArrowArray* psChild,