Line data Source code
1 : /******************************************************************************
2 : *
3 : * Project: Arrow generic code
4 : * Purpose: Arrow generic code
5 : * Author: Even Rouault, <even.rouault at spatialys.com>
6 : *
7 : ******************************************************************************
8 : * Copyright (c) 2022, Planet Labs
9 : *
10 : * SPDX-License-Identifier: MIT
11 : ****************************************************************************/
12 :
13 : #ifndef OGR_ARROW_H
14 : #define OGR_ARROW_H
15 :
16 : #include "gdal_pam.h"
17 : #include "ogrsf_frmts.h"
18 : #include "ogrlayerarrow.h"
19 :
20 : #include <map>
21 : #include <set>
22 :
23 : #include "ogr_include_arrow.h"
24 :
25 : #if defined(__clang__)
26 : #pragma clang diagnostic push
27 : #pragma clang diagnostic ignored "-Wweak-vtables"
28 : #endif
29 :
30 : enum class OGRArrowGeomEncoding
31 : {
32 : WKB,
33 : WKT,
34 :
35 : // F(ixed) S(ize) L(ist) of (x,y[,z][,m]) values / Interleaved layout
36 : GEOARROW_FSL_GENERIC, // only used by OGRArrowWriterLayer::m_eGeomEncoding
37 : GEOARROW_FSL_POINT,
38 : GEOARROW_FSL_LINESTRING,
39 : GEOARROW_FSL_POLYGON,
40 : GEOARROW_FSL_MULTIPOINT,
41 : GEOARROW_FSL_MULTILINESTRING,
42 : GEOARROW_FSL_MULTIPOLYGON,
43 :
44 : // Struct of (x,y,[,z][,m])
45 : GEOARROW_STRUCT_GENERIC, // only used by OGRArrowWriterLayer::m_eGeomEncoding
46 : GEOARROW_STRUCT_POINT,
47 : GEOARROW_STRUCT_LINESTRING,
48 : GEOARROW_STRUCT_POLYGON,
49 : GEOARROW_STRUCT_MULTIPOINT,
50 : GEOARROW_STRUCT_MULTILINESTRING,
51 : GEOARROW_STRUCT_MULTIPOLYGON,
52 : };
53 :
54 : /************************************************************************/
55 : /* OGRArrowIsGeoArrowStruct() */
56 : /************************************************************************/
57 :
58 260 : inline bool OGRArrowIsGeoArrowStruct(OGRArrowGeomEncoding eEncoding)
59 : {
60 260 : switch (eEncoding)
61 : {
62 260 : case OGRArrowGeomEncoding::GEOARROW_STRUCT_GENERIC:
63 : case OGRArrowGeomEncoding::GEOARROW_STRUCT_POINT:
64 : case OGRArrowGeomEncoding::GEOARROW_STRUCT_LINESTRING:
65 : case OGRArrowGeomEncoding::GEOARROW_STRUCT_POLYGON:
66 : case OGRArrowGeomEncoding::GEOARROW_STRUCT_MULTIPOINT:
67 : case OGRArrowGeomEncoding::GEOARROW_STRUCT_MULTILINESTRING:
68 : case OGRArrowGeomEncoding::GEOARROW_STRUCT_MULTIPOLYGON:
69 260 : return true;
70 :
71 0 : default:
72 0 : return false;
73 : }
74 : }
75 :
76 : /************************************************************************/
77 : /* OGRArrowLayer */
78 : /************************************************************************/
79 :
80 : class OGRArrowDataset;
81 :
82 : class OGRArrowLayer CPL_NON_FINAL
83 : : public OGRLayer,
84 : public OGRGetNextFeatureThroughRaw<OGRArrowLayer>
85 : {
86 : public:
87 : struct Constraint
88 : {
89 : enum class Type
90 : {
91 : Integer,
92 : Integer64,
93 : Real,
94 : String,
95 : };
96 : int iField = -1; // index to a OGRFeatureDefn OGRField
97 : int iArrayIdx = -1; // index to m_poBatchColumns
98 : int nOperation = -1; // SWQ_xxxx
99 : Type eType{};
100 : OGRField sValue{};
101 : std::string osValue{};
102 : };
103 :
104 : private:
105 : OGRArrowLayer(const OGRArrowLayer &) = delete;
106 : OGRArrowLayer &operator=(const OGRArrowLayer &) = delete;
107 :
108 : int m_nUseOptimizedAttributeFilter = -1;
109 : bool m_bSpatialFilterIntersectsLayerExtent = true;
110 : bool m_bUseRecordBatchBaseImplementation = false;
111 :
112 : // Modified by UseRecordBatchBaseImplementation()
113 : mutable struct ArrowSchema m_sCachedSchema = {};
114 :
115 : bool SkipToNextFeatureDueToAttributeFilter() const;
116 : void ExploreExprNode(const swq_expr_node *poNode);
117 : bool UseRecordBatchBaseImplementation() const;
118 :
119 : template <typename SourceOffset>
120 : static struct ArrowArray *
121 : CreateWKBArrayFromWKTArray(const struct ArrowArray *sourceArray);
122 :
123 : int GetArrowSchemaInternal(struct ArrowSchema *out) const;
124 :
125 : protected:
126 : OGRArrowDataset *m_poArrowDS = nullptr;
127 : arrow::MemoryPool *m_poMemoryPool = nullptr;
128 : OGRFeatureDefn *m_poFeatureDefn = nullptr;
129 : std::shared_ptr<arrow::Schema> m_poSchema{};
130 : std::string m_osFIDColumn{};
131 : int m_iFIDArrowColumn = -1;
132 : std::vector<std::vector<int>> m_anMapFieldIndexToArrowColumn{};
133 : std::vector<int> m_anMapGeomFieldIndexToArrowColumn{};
134 : std::vector<OGRArrowGeomEncoding> m_aeGeomEncoding{};
135 :
136 : //! Whether bounding box based spatial filter should be skipped.
137 : // This is set to true by OGRParquetDatasetLayer when there is a bounding
138 : // box field, as an optimization.
139 : bool m_bBaseArrowIgnoreSpatialFilterRect = false;
140 :
141 : //! Whether spatial filter should be skipped (by GetNextArrowArray())
142 : // This is set to true by OGRParquetDatasetLayer when filtering points in
143 : // a rectangle.
144 : bool m_bBaseArrowIgnoreSpatialFilter = false;
145 :
146 : //! Describe the bbox column of a geometry column
147 : struct GeomColBBOX
148 : {
149 : bool bIsFloat = false;
150 : int iArrowCol = -1;
151 : int iArrayIdx = -1; // only valid when m_bIgnoredFields == true
152 : int iArrowSubfieldXMin = -1;
153 : int iArrowSubfieldYMin = -1;
154 : int iArrowSubfieldXMax = -1;
155 : int iArrowSubfieldYMax = -1;
156 : };
157 :
158 : //! Map from OGR geometry field index to GeomColBBOX
159 : std::map<int, GeomColBBOX> m_oMapGeomFieldIndexToGeomColBBOX{};
160 :
161 : const arrow::BinaryArray *m_poArrayWKB = nullptr;
162 : const arrow::LargeBinaryArray *m_poArrayWKBLarge = nullptr;
163 : const arrow::Array *m_poArrayBBOX = nullptr;
164 : const arrow::DoubleArray *m_poArrayXMinDouble = nullptr;
165 : const arrow::DoubleArray *m_poArrayYMinDouble = nullptr;
166 : const arrow::DoubleArray *m_poArrayXMaxDouble = nullptr;
167 : const arrow::DoubleArray *m_poArrayYMaxDouble = nullptr;
168 : const arrow::FloatArray *m_poArrayXMinFloat = nullptr;
169 : const arrow::FloatArray *m_poArrayYMinFloat = nullptr;
170 : const arrow::FloatArray *m_poArrayXMaxFloat = nullptr;
171 : const arrow::FloatArray *m_poArrayYMaxFloat = nullptr;
172 :
173 : //! References values in range [0, m_poSchema->field_count()-1]
174 : std::set<int> m_oSetBBoxArrowColumns{};
175 :
176 : bool m_bIgnoredFields = false;
177 : std::vector<int>
178 : m_anMapFieldIndexToArrayIndex{}; // only valid when m_bIgnoredFields is
179 : // set
180 : std::vector<int> m_anMapGeomFieldIndexToArrayIndex{}; // only valid when
181 : // m_bIgnoredFields is set
182 : int m_nRequestedFIDColumn = -1; // only valid when m_bIgnoredFields is set
183 :
184 : int m_nExpectedBatchColumns =
185 : -1; // Should be equal to m_poBatch->num_columns() (when
186 : // m_bIgnoredFields is set)
187 :
188 : bool m_bEOF = false;
189 : int64_t m_nFeatureIdx = 0;
190 : int64_t m_nIdxInBatch = 0;
191 : std::map<std::string, CPLJSONObject> m_oMapGeometryColumns{};
192 : mutable std::map<int, OGREnvelope> m_oMapExtents{};
193 : mutable std::map<int, OGREnvelope3D> m_oMapExtents3D{};
194 : int m_iRecordBatch = -1;
195 : std::shared_ptr<arrow::RecordBatch> m_poBatch{};
196 : // m_poBatch->columns() is a relatively costly operation, so cache its
197 : // result
198 : std::vector<std::shared_ptr<arrow::Array>>
199 : m_poBatchColumns{}; // must always be == m_poBatch->columns()
200 : mutable std::shared_ptr<arrow::Array> m_poReadFeatureTmpArray{};
201 :
202 : std::vector<Constraint> m_asAttributeFilterConstraints{};
203 :
204 : //! Whether attribute filter should be skipped.
205 : // This is set to true by OGRParquetDatasetLayer when it can fully translate
206 : // a filter, as an optimization.
207 : bool m_bBaseArrowIgnoreAttributeFilter = false;
208 :
209 : std::map<std::string, std::unique_ptr<OGRFieldDefn>>
210 : LoadGDALSchema(const arrow::KeyValueMetadata *kv_metadata);
211 :
212 : void LoadGDALMetadata(const arrow::KeyValueMetadata *kv_metadata);
213 :
214 : OGRArrowLayer(OGRArrowDataset *poDS, const char *pszLayerName);
215 :
216 : virtual std::string GetDriverUCName() const = 0;
217 : static bool IsIntegerArrowType(arrow::Type::type typeId);
218 : static bool
219 : IsHandledListOrMapType(const std::shared_ptr<arrow::DataType> &valueType);
220 : static bool
221 : IsHandledListType(const std::shared_ptr<arrow::BaseListType> &listType);
222 : static bool
223 : IsHandledMapType(const std::shared_ptr<arrow::MapType> &mapType);
224 : static bool
225 : IsValidGeometryEncoding(const std::shared_ptr<arrow::Field> &field,
226 : const std::string &osEncoding,
227 : bool bWarnIfUnknownEncoding,
228 : OGRwkbGeometryType &eGeomTypeOut,
229 : OGRArrowGeomEncoding &eGeomEncodingOut);
230 : static OGRwkbGeometryType
231 : GetGeometryTypeFromString(const std::string &osType);
232 : bool
233 : MapArrowTypeToOGR(const std::shared_ptr<arrow::DataType> &type,
234 : const std::shared_ptr<arrow::Field> &field,
235 : OGRFieldDefn &oField, OGRFieldType &eType,
236 : OGRFieldSubType &eSubType, const std::vector<int> &path,
237 : const std::map<std::string, std::unique_ptr<OGRFieldDefn>>
238 : &oMapFieldNameToGDALSchemaFieldDefn);
239 : void CreateFieldFromSchema(
240 : const std::shared_ptr<arrow::Field> &field,
241 : const std::vector<int> &path,
242 : const std::map<std::string, std::unique_ptr<OGRFieldDefn>>
243 : &oMapFieldNameToGDALSchemaFieldDefn);
244 : std::unique_ptr<OGRFieldDomain>
245 : BuildDomainFromBatch(const std::string &osDomainName,
246 : const std::shared_ptr<arrow::RecordBatch> &poBatch,
247 : int iCol) const;
248 : OGRwkbGeometryType ComputeGeometryColumnTypeProcessBatch(
249 : const std::shared_ptr<arrow::RecordBatch> &poBatch, int iGeomCol,
250 : int iBatchCol, OGRwkbGeometryType eGeomType) const;
251 : static bool ReadWKBBoundingBox(const uint8_t *data, size_t size,
252 : OGREnvelope &sEnvelope);
253 : OGRFeature *ReadFeature(
254 : int64_t nIdxInBatch,
255 : const std::vector<std::shared_ptr<arrow::Array>> &poColumnArrays) const;
256 : OGRGeometry *ReadGeometry(int iGeomField, const arrow::Array *array,
257 : int64_t nIdxInBatch) const;
258 : virtual bool ReadNextBatch() = 0;
259 : virtual void InvalidateCachedBatches() = 0;
260 : OGRFeature *GetNextRawFeature();
261 :
262 0 : virtual bool CanRunNonForcedGetExtent()
263 : {
264 0 : return true;
265 : }
266 :
267 : void SetBatch(const std::shared_ptr<arrow::RecordBatch> &poBatch);
268 :
269 : // Refreshes Constraint.iArrayIdx from iField. To be called by SetIgnoredFields()
270 : void ComputeConstraintsArrayIdx();
271 :
272 : static const swq_expr_node *GetColumnSubNode(const swq_expr_node *poNode);
273 : static const swq_expr_node *GetConstantSubNode(const swq_expr_node *poNode);
274 : static bool IsComparisonOp(int op);
275 :
276 : virtual bool FastGetExtent(int iGeomField, OGREnvelope *psExtent) const;
277 : bool FastGetExtent3D(int iGeomField, OGREnvelope3D *psExtent) const;
278 : static OGRErr GetExtentFromMetadata(const CPLJSONObject &oJSONDef,
279 : OGREnvelope3D *psExtent);
280 :
281 : int GetArrowSchema(struct ArrowArrayStream *,
282 : struct ArrowSchema *out) override;
283 : int GetNextArrowArray(struct ArrowArrayStream *,
284 : struct ArrowArray *out) override;
285 :
286 4695 : virtual void IncrFeatureIdx()
287 : {
288 4695 : ++m_nFeatureIdx;
289 4695 : }
290 :
291 : void SanityCheckOfSetBatch() const;
292 :
293 : public:
294 : virtual ~OGRArrowLayer() override;
295 :
296 13550 : OGRFeatureDefn *GetLayerDefn() override
297 : {
298 13550 : return m_poFeatureDefn;
299 : }
300 :
301 : void ResetReading() override;
302 :
303 1875 : const char *GetFIDColumn() override
304 : {
305 1875 : return m_osFIDColumn.c_str();
306 : }
307 9056 : DEFINE_GET_NEXT_FEATURE_THROUGH_RAW(OGRArrowLayer)
308 : OGRErr IGetExtent(int iGeomField, OGREnvelope *psExtent,
309 : bool bForce) override;
310 : OGRErr IGetExtent3D(int iGeomField, OGREnvelope3D *psExtent,
311 : bool bForce) override;
312 : OGRErr SetAttributeFilter(const char *pszFilter) override;
313 :
314 : OGRErr ISetSpatialFilter(int iGeomField,
315 : const OGRGeometry *poGeom) override;
316 :
317 : int TestCapability(const char *pszCap) override;
318 :
319 : bool GetArrowStream(struct ArrowArrayStream *out_stream,
320 : CSLConstList papszOptions = nullptr) override;
321 :
322 : virtual std::unique_ptr<OGRFieldDomain>
323 : BuildDomain(const std::string &osDomainName, int iFieldIndex) const = 0;
324 :
325 : static void TimestampToOGR(int64_t timestamp,
326 : const arrow::TimestampType *timestampType,
327 : int nTZFlag, OGRField *psField);
328 : };
329 :
330 : /************************************************************************/
331 : /* OGRArrowDataset */
332 : /************************************************************************/
333 :
334 : class OGRArrowDataset CPL_NON_FINAL : public GDALPamDataset
335 : {
336 : std::shared_ptr<arrow::MemoryPool> m_poMemoryPool{};
337 : std::unique_ptr<OGRArrowLayer> m_poLayer{};
338 : std::vector<std::string> m_aosDomainNames{};
339 : std::map<std::string, int> m_oMapDomainNameToCol{};
340 :
341 : protected:
342 1136 : void close()
343 : {
344 1136 : m_poLayer.reset();
345 1136 : m_poMemoryPool.reset();
346 1136 : }
347 :
348 : public:
349 : explicit OGRArrowDataset(
350 : const std::shared_ptr<arrow::MemoryPool> &poMemoryPool);
351 :
352 1681 : inline arrow::MemoryPool *GetMemoryPool() const
353 : {
354 1681 : return m_poMemoryPool.get();
355 : }
356 :
357 609 : inline const std::shared_ptr<arrow::MemoryPool> &GetSharedMemoryPool() const
358 : {
359 609 : return m_poMemoryPool;
360 : }
361 :
362 : void SetLayer(std::unique_ptr<OGRArrowLayer> &&poLayer);
363 :
364 : void RegisterDomainName(const std::string &osDomainName, int iFieldIndex);
365 :
366 : std::vector<std::string> GetFieldDomainNames(
367 : CSLConstList /*papszOptions*/ = nullptr) const override;
368 : const OGRFieldDomain *
369 : GetFieldDomain(const std::string &name) const override;
370 :
371 : int GetLayerCount() override;
372 : OGRLayer *GetLayer(int idx) override;
373 : };
374 :
375 : /************************************************************************/
376 : /* OGRArrowWriterLayer */
377 : /************************************************************************/
378 :
379 : class OGRArrowWriterLayer CPL_NON_FINAL : public OGRLayer
380 :
381 : {
382 : protected:
383 : OGRArrowWriterLayer(const OGRArrowWriterLayer &) = delete;
384 : OGRArrowWriterLayer &operator=(const OGRArrowWriterLayer &) = delete;
385 :
386 : arrow::MemoryPool *m_poMemoryPool = nullptr;
387 : bool m_bInitializationOK = false;
388 : std::shared_ptr<arrow::io::OutputStream> m_poOutputStream{};
389 : std::shared_ptr<arrow::Schema> m_poSchema{};
390 : OGRFeatureDefn *m_poFeatureDefn = nullptr;
391 : std::map<std::string, std::unique_ptr<OGRFieldDomain>> m_oMapFieldDomains{};
392 : std::map<std::string, std::shared_ptr<arrow::Array>>
393 : m_oMapFieldDomainToStringArray{};
394 :
395 : bool m_bWriteFieldArrowExtensionName = false;
396 : OGRArrowGeomEncoding m_eGeomEncoding = OGRArrowGeomEncoding::WKB;
397 : std::vector<OGRArrowGeomEncoding> m_aeGeomEncoding{};
398 : int m_nWKTCoordinatePrecision = -1;
399 :
400 : //! Base struct data type for GeoArrow struct geometry columns.
401 : // Constraint: if not empty, m_apoBaseStructGeomType.size() == m_poFeatureDefn->GetGeomFieldCount()
402 : std::vector<std::shared_ptr<arrow::DataType>> m_apoBaseStructGeomType{};
403 :
404 : //! Whether to use a struct field with the values of the bounding box
405 : // of the geometries. Used by Parquet.
406 : bool m_bWriteBBoxStruct = false;
407 :
408 : //! Schema fields for bounding box of geometry columns.
409 : // Constraint: if not empty, m_apoFieldsBBOX.size() == m_poFeatureDefn->GetGeomFieldCount()
410 : std::vector<std::shared_ptr<arrow::Field>> m_apoFieldsBBOX{};
411 :
412 : //! Array builers for bounding box of geometry columns.
413 : // m_apoBuildersBBOXStruct is for the top-level field of type struct.
414 : // m_apoBuildersBBOX{XMin|YMin|XMax|YMax} are for the floating-point values
415 : // Constraint: if not empty, m_apoBuildersBBOX{Struct|XMin|YMin|XMax|YMax}.size() == m_poFeatureDefn->GetGeomFieldCount()
416 : std::vector<std::shared_ptr<arrow::StructBuilder>>
417 : m_apoBuildersBBOXStruct{};
418 : std::vector<std::shared_ptr<arrow::FloatBuilder>> m_apoBuildersBBOXXMin{};
419 : std::vector<std::shared_ptr<arrow::FloatBuilder>> m_apoBuildersBBOXYMin{};
420 : std::vector<std::shared_ptr<arrow::FloatBuilder>> m_apoBuildersBBOXXMax{};
421 : std::vector<std::shared_ptr<arrow::FloatBuilder>> m_apoBuildersBBOXYMax{};
422 :
423 : std::string m_osFIDColumn{};
424 : int64_t m_nFeatureCount = 0;
425 :
426 : int64_t m_nRowGroupSize = 64 * 1024;
427 : arrow::Compression::type m_eCompression = arrow::Compression::UNCOMPRESSED;
428 :
429 : std::vector<std::shared_ptr<arrow::Field>> m_apoFieldsFromArrowSchema{};
430 : std::vector<std::shared_ptr<arrow::ArrayBuilder>> m_apoBuilders{};
431 :
432 : std::vector<uint8_t> m_abyBuffer{};
433 :
434 : std::vector<int> m_anTZFlag{}; // size: GetFieldCount()
435 : std::vector<OGREnvelope3D> m_aoEnvelopes{}; // size: GetGeomFieldCount()
436 : std::vector<std::set<OGRwkbGeometryType>>
437 : m_oSetWrittenGeometryTypes{}; // size: GetGeomFieldCount()
438 :
439 : bool m_bEdgesSpherical = false;
440 : #if ARROW_VERSION_MAJOR >= 21
441 : bool m_bUseArrowWKBExtension = false;
442 : #endif
443 :
444 : static OGRArrowGeomEncoding
445 : GetPreciseArrowGeomEncoding(OGRArrowGeomEncoding eEncodingType,
446 : OGRwkbGeometryType eGType);
447 : static const char *
448 : GetGeomEncodingAsString(OGRArrowGeomEncoding eGeomEncoding,
449 : bool bForParquetGeo);
450 :
451 : virtual bool IsSupportedGeometryType(OGRwkbGeometryType eGType) const = 0;
452 :
453 : virtual std::string GetDriverUCName() const = 0;
454 :
455 : virtual bool IsFileWriterCreated() const = 0;
456 : virtual void CreateWriter() = 0;
457 : virtual bool CloseFileWriter() = 0;
458 :
459 : void CreateSchemaCommon();
460 : void FinalizeSchema();
461 : virtual void CreateSchema() = 0;
462 :
463 0 : virtual void PerformStepsBeforeFinalFlushGroup()
464 : {
465 0 : }
466 :
467 : void CreateArrayBuilders();
468 :
469 : //! Clear array builders
470 : void ClearArrayBuilers();
471 :
472 : virtual bool FlushGroup() = 0;
473 : bool FinalizeWriting();
474 : bool WriteArrays(std::function<bool(const std::shared_ptr<arrow::Field> &,
475 : const std::shared_ptr<arrow::Array> &)>
476 : postProcessArray);
477 :
478 128 : virtual void FixupWKBGeometryBeforeWriting(GByte * /*pabyWKB*/,
479 : size_t /*nLen*/)
480 : {
481 128 : }
482 :
483 0 : virtual void FixupGeometryBeforeWriting(OGRGeometry * /* poGeom */)
484 : {
485 0 : }
486 :
487 : virtual bool IsSRSRequired() const = 0;
488 : bool WriteArrowBatchInternal(
489 : const struct ArrowSchema *schema, struct ArrowArray *array,
490 : CSLConstList papszOptions,
491 : std::function<bool(const std::shared_ptr<arrow::RecordBatch> &)>
492 : writeBatch);
493 :
494 : OGRErr BuildGeometry(OGRGeometry *poGeom, int iGeomField,
495 : arrow::ArrayBuilder *poBuilder);
496 :
497 : public:
498 : OGRArrowWriterLayer(
499 : arrow::MemoryPool *poMemoryPool,
500 : const std::shared_ptr<arrow::io::OutputStream> &poOutputStream,
501 : const char *pszLayerName);
502 :
503 : ~OGRArrowWriterLayer() override;
504 :
505 : bool AddFieldDomain(std::unique_ptr<OGRFieldDomain> &&domain,
506 : std::string &failureReason);
507 : std::vector<std::string> GetFieldDomainNames() const;
508 : const OGRFieldDomain *GetFieldDomain(const std::string &name) const;
509 :
510 9 : const char *GetFIDColumn() override
511 : {
512 9 : return m_osFIDColumn.c_str();
513 : }
514 :
515 6105 : OGRFeatureDefn *GetLayerDefn() override
516 : {
517 6105 : return m_poFeatureDefn;
518 : }
519 :
520 23 : void ResetReading() override
521 : {
522 23 : }
523 :
524 23 : OGRFeature *GetNextFeature() override
525 : {
526 23 : return nullptr;
527 : }
528 :
529 : int TestCapability(const char *pszCap) override;
530 : OGRErr CreateField(const OGRFieldDefn *poField,
531 : int bApproxOK = TRUE) override;
532 : OGRErr CreateGeomField(const OGRGeomFieldDefn *poField,
533 : int bApproxOK = TRUE) override;
534 : GIntBig GetFeatureCount(int bForce) override;
535 :
536 112 : bool IsArrowSchemaSupported(const struct ArrowSchema * /*schema*/,
537 : CSLConstList /* papszOptions */,
538 : std::string & /*osErrorMsg */) const override
539 : {
540 112 : return true;
541 : }
542 :
543 : bool
544 : CreateFieldFromArrowSchema(const struct ArrowSchema *schema,
545 : CSLConstList papszOptions = nullptr) override;
546 : bool WriteArrowBatch(const struct ArrowSchema *schema,
547 : struct ArrowArray *array,
548 : CSLConstList papszOptions = nullptr) override = 0;
549 :
550 : protected:
551 : OGRErr ICreateFeature(OGRFeature *poFeature) override;
552 :
553 : bool FlushFeatures();
554 :
555 : static void RemoveIDFromMemberOfEnsembles(CPLJSONObject &obj);
556 : static OGRSpatialReference IdentifyCRS(const OGRSpatialReference *poSRS);
557 : };
558 :
559 : /************************************************************************/
560 : /* OGRGeoArrowWkbExtensionType */
561 : /************************************************************************/
562 :
563 : #if ARROW_VERSION_MAJOR >= 21
564 :
565 : class OGRGeoArrowWkbExtensionType final : public arrow::ExtensionType
566 : {
567 : public:
568 : explicit OGRGeoArrowWkbExtensionType(
569 : const std::shared_ptr<arrow::DataType> &storage_type,
570 : const std::string &metadata)
571 : : arrow::ExtensionType(storage_type), metadata_(metadata)
572 : {
573 : }
574 :
575 : std::string extension_name() const override
576 : {
577 : return EXTENSION_NAME_GEOARROW_WKB;
578 : }
579 :
580 : bool ExtensionEquals(const arrow::ExtensionType &other) const override
581 : {
582 : return extension_name() == other.extension_name() &&
583 : storage_type_->Equals(other.storage_type()) &&
584 : Serialize() == other.Serialize();
585 : }
586 :
587 : arrow::Result<std::shared_ptr<arrow::DataType>>
588 : Deserialize(std::shared_ptr<arrow::DataType> storage_type,
589 : const std::string &serialized) const override
590 : {
591 : return Make(std::move(storage_type), serialized);
592 : }
593 :
594 : std::string Serialize() const override
595 : {
596 : return metadata_;
597 : }
598 :
599 : std::shared_ptr<arrow::Array>
600 : MakeArray(std::shared_ptr<arrow::ArrayData> data) const override
601 : {
602 : CPLAssert(data->type->id() == arrow::Type::EXTENSION);
603 : CPLAssert(EXTENSION_NAME_GEOARROW_WKB ==
604 : static_cast<const arrow::ExtensionType &>(*data->type)
605 : .extension_name());
606 : return std::make_shared<arrow::ExtensionArray>(data);
607 : }
608 :
609 : static bool IsSupportedStorageType(arrow::Type::type type_id)
610 : {
611 : // TODO: also add BINARY_VIEW if we support it some day.
612 : return type_id == arrow::Type::BINARY ||
613 : type_id == arrow::Type::LARGE_BINARY;
614 : }
615 :
616 : static arrow::Result<std::shared_ptr<arrow::DataType>>
617 : Make(std::shared_ptr<arrow::DataType> storage_type,
618 : const std::string &metadata)
619 : {
620 : if (!IsSupportedStorageType(storage_type->id()))
621 : {
622 : return arrow::Status::Invalid(
623 : "Invalid storage type for OGRGeoArrowWkbExtensionType: ",
624 : storage_type->ToString());
625 : }
626 : return std::make_shared<OGRGeoArrowWkbExtensionType>(
627 : std::move(storage_type), metadata);
628 : }
629 :
630 : private:
631 : std::string metadata_{};
632 : };
633 :
634 : #endif // ARROW_VERSION_MAJOR >= 21
635 :
636 : #if defined(__clang__)
637 : #pragma clang diagnostic pop
638 : #endif
639 :
640 : #endif // OGR_ARROW_H
|
