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 4737 : virtual void IncrFeatureIdx()
287 : {
288 4737 : ++m_nFeatureIdx;
289 4737 : }
290 :
291 : void SanityCheckOfSetBatch() const;
292 :
293 : public:
294 : ~OGRArrowLayer() override;
295 :
296 13059 : const OGRFeatureDefn *GetLayerDefn() const override
297 : {
298 13059 : return m_poFeatureDefn;
299 : }
300 :
301 : void ResetReading() override;
302 :
303 1875 : const char *GetFIDColumn() const override
304 : {
305 1875 : return m_osFIDColumn.c_str();
306 : }
307 9103 : 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) const 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 : public:
342 : explicit OGRArrowDataset(
343 : const std::shared_ptr<arrow::MemoryPool> &poMemoryPool);
344 :
345 1694 : ~OGRArrowDataset() override
346 1694 : {
347 1694 : OGRArrowDataset::Close();
348 1694 : }
349 :
350 3386 : CPLErr Close() override
351 : {
352 3386 : m_poLayer.reset();
353 3386 : m_poMemoryPool.reset();
354 3386 : return GDALPamDataset::Close();
355 : }
356 :
357 1692 : inline arrow::MemoryPool *GetMemoryPool() const
358 : {
359 1692 : return m_poMemoryPool.get();
360 : }
361 :
362 609 : inline const std::shared_ptr<arrow::MemoryPool> &GetSharedMemoryPool() const
363 : {
364 609 : return m_poMemoryPool;
365 : }
366 :
367 : void SetLayer(std::unique_ptr<OGRArrowLayer> &&poLayer);
368 :
369 : void RegisterDomainName(const std::string &osDomainName, int iFieldIndex);
370 :
371 : std::vector<std::string> GetFieldDomainNames(
372 : CSLConstList /*papszOptions*/ = nullptr) const override;
373 : const OGRFieldDomain *
374 : GetFieldDomain(const std::string &name) const override;
375 :
376 : int GetLayerCount() const override;
377 : using GDALDataset::GetLayer;
378 : const OGRLayer *GetLayer(int idx) const override;
379 : };
380 :
381 : /************************************************************************/
382 : /* OGRArrowWriterLayer */
383 : /************************************************************************/
384 :
385 : class OGRArrowWriterLayer CPL_NON_FINAL : public OGRLayer
386 :
387 : {
388 : protected:
389 : OGRArrowWriterLayer(const OGRArrowWriterLayer &) = delete;
390 : OGRArrowWriterLayer &operator=(const OGRArrowWriterLayer &) = delete;
391 :
392 : arrow::MemoryPool *m_poMemoryPool = nullptr;
393 : bool m_bInitializationOK = false;
394 : std::shared_ptr<arrow::io::OutputStream> m_poOutputStream{};
395 : std::shared_ptr<arrow::Schema> m_poSchema{};
396 : OGRFeatureDefn *m_poFeatureDefn = nullptr;
397 : std::map<std::string, std::unique_ptr<OGRFieldDomain>> m_oMapFieldDomains{};
398 : std::map<std::string, std::shared_ptr<arrow::Array>>
399 : m_oMapFieldDomainToStringArray{};
400 :
401 : bool m_bWriteFieldArrowExtensionName = false;
402 : OGRArrowGeomEncoding m_eGeomEncoding = OGRArrowGeomEncoding::WKB;
403 : std::vector<OGRArrowGeomEncoding> m_aeGeomEncoding{};
404 : int m_nWKTCoordinatePrecision = -1;
405 :
406 : //! Base struct data type for GeoArrow struct geometry columns.
407 : // Constraint: if not empty, m_apoBaseStructGeomType.size() == m_poFeatureDefn->GetGeomFieldCount()
408 : std::vector<std::shared_ptr<arrow::DataType>> m_apoBaseStructGeomType{};
409 :
410 : //! Whether to use a struct field with the values of the bounding box
411 : // of the geometries. Used by Parquet.
412 : bool m_bWriteBBoxStruct = false;
413 :
414 : //! Schema fields for bounding box of geometry columns.
415 : // Constraint: if not empty, m_apoFieldsBBOX.size() == m_poFeatureDefn->GetGeomFieldCount()
416 : std::vector<std::shared_ptr<arrow::Field>> m_apoFieldsBBOX{};
417 :
418 : //! Array builers for bounding box of geometry columns.
419 : // m_apoBuildersBBOXStruct is for the top-level field of type struct.
420 : // m_apoBuildersBBOX{XMin|YMin|XMax|YMax} are for the floating-point values
421 : // Constraint: if not empty, m_apoBuildersBBOX{Struct|XMin|YMin|XMax|YMax}.size() == m_poFeatureDefn->GetGeomFieldCount()
422 : std::vector<std::shared_ptr<arrow::StructBuilder>>
423 : m_apoBuildersBBOXStruct{};
424 : std::vector<std::shared_ptr<arrow::FloatBuilder>> m_apoBuildersBBOXXMin{};
425 : std::vector<std::shared_ptr<arrow::FloatBuilder>> m_apoBuildersBBOXYMin{};
426 : std::vector<std::shared_ptr<arrow::FloatBuilder>> m_apoBuildersBBOXXMax{};
427 : std::vector<std::shared_ptr<arrow::FloatBuilder>> m_apoBuildersBBOXYMax{};
428 :
429 : std::string m_osFIDColumn{};
430 : int64_t m_nFeatureCount = 0;
431 :
432 : int64_t m_nRowGroupSize = 64 * 1024;
433 : arrow::Compression::type m_eCompression = arrow::Compression::UNCOMPRESSED;
434 :
435 : std::vector<std::shared_ptr<arrow::Field>> m_apoFieldsFromArrowSchema{};
436 : std::vector<std::shared_ptr<arrow::ArrayBuilder>> m_apoBuilders{};
437 :
438 : std::vector<uint8_t> m_abyBuffer{};
439 :
440 : std::vector<int> m_anTZFlag{}; // size: GetFieldCount()
441 : std::vector<OGREnvelope3D> m_aoEnvelopes{}; // size: GetGeomFieldCount()
442 : std::vector<std::set<OGRwkbGeometryType>>
443 : m_oSetWrittenGeometryTypes{}; // size: GetGeomFieldCount()
444 :
445 : bool m_bEdgesSpherical = false;
446 : #if ARROW_VERSION_MAJOR >= 21
447 : bool m_bUseArrowWKBExtension = false;
448 : #endif
449 :
450 : static OGRArrowGeomEncoding
451 : GetPreciseArrowGeomEncoding(OGRArrowGeomEncoding eEncodingType,
452 : OGRwkbGeometryType eGType);
453 : static const char *
454 : GetGeomEncodingAsString(OGRArrowGeomEncoding eGeomEncoding,
455 : bool bForParquetGeo);
456 :
457 : virtual bool IsSupportedGeometryType(OGRwkbGeometryType eGType) const = 0;
458 :
459 : virtual std::string GetDriverUCName() const = 0;
460 :
461 : virtual bool IsFileWriterCreated() const = 0;
462 : virtual void CreateWriter() = 0;
463 : virtual bool CloseFileWriter() = 0;
464 :
465 : void CreateSchemaCommon();
466 : void FinalizeSchema();
467 : virtual void CreateSchema() = 0;
468 :
469 0 : virtual void PerformStepsBeforeFinalFlushGroup()
470 : {
471 0 : }
472 :
473 : void CreateArrayBuilders();
474 :
475 : //! Clear array builders
476 : void ClearArrayBuilers();
477 :
478 : virtual bool FlushGroup() = 0;
479 : bool FinalizeWriting();
480 : bool WriteArrays(std::function<bool(const std::shared_ptr<arrow::Field> &,
481 : const std::shared_ptr<arrow::Array> &)>
482 : postProcessArray);
483 :
484 128 : virtual void FixupWKBGeometryBeforeWriting(GByte * /*pabyWKB*/,
485 : size_t /*nLen*/)
486 : {
487 128 : }
488 :
489 0 : virtual void FixupGeometryBeforeWriting(OGRGeometry * /* poGeom */)
490 : {
491 0 : }
492 :
493 : virtual bool IsSRSRequired() const = 0;
494 : bool WriteArrowBatchInternal(
495 : const struct ArrowSchema *schema, struct ArrowArray *array,
496 : CSLConstList papszOptions,
497 : std::function<bool(const std::shared_ptr<arrow::RecordBatch> &)>
498 : writeBatch);
499 :
500 : OGRErr BuildGeometry(OGRGeometry *poGeom, int iGeomField,
501 : arrow::ArrayBuilder *poBuilder);
502 :
503 : public:
504 : OGRArrowWriterLayer(
505 : arrow::MemoryPool *poMemoryPool,
506 : const std::shared_ptr<arrow::io::OutputStream> &poOutputStream,
507 : const char *pszLayerName);
508 :
509 : ~OGRArrowWriterLayer() override;
510 :
511 : bool AddFieldDomain(std::unique_ptr<OGRFieldDomain> &&domain,
512 : std::string &failureReason);
513 : std::vector<std::string> GetFieldDomainNames() const;
514 : const OGRFieldDomain *GetFieldDomain(const std::string &name) const;
515 :
516 10 : const char *GetFIDColumn() const override
517 : {
518 10 : return m_osFIDColumn.c_str();
519 : }
520 :
521 : using OGRLayer::GetLayerDefn;
522 :
523 6161 : const OGRFeatureDefn *GetLayerDefn() const override
524 : {
525 6161 : return m_poFeatureDefn;
526 : }
527 :
528 23 : void ResetReading() override
529 : {
530 23 : }
531 :
532 23 : OGRFeature *GetNextFeature() override
533 : {
534 23 : return nullptr;
535 : }
536 :
537 : int TestCapability(const char *pszCap) const override;
538 : OGRErr CreateField(const OGRFieldDefn *poField,
539 : int bApproxOK = TRUE) override;
540 : OGRErr CreateGeomField(const OGRGeomFieldDefn *poField,
541 : int bApproxOK = TRUE) override;
542 : GIntBig GetFeatureCount(int bForce) override;
543 :
544 112 : bool IsArrowSchemaSupported(const struct ArrowSchema * /*schema*/,
545 : CSLConstList /* papszOptions */,
546 : std::string & /*osErrorMsg */) const override
547 : {
548 112 : return true;
549 : }
550 :
551 : bool
552 : CreateFieldFromArrowSchema(const struct ArrowSchema *schema,
553 : CSLConstList papszOptions = nullptr) override;
554 : bool WriteArrowBatch(const struct ArrowSchema *schema,
555 : struct ArrowArray *array,
556 : CSLConstList papszOptions = nullptr) override = 0;
557 :
558 : protected:
559 : OGRErr ICreateFeature(OGRFeature *poFeature) override;
560 :
561 : bool FlushFeatures();
562 :
563 : static void RemoveIDFromMemberOfEnsembles(CPLJSONObject &obj);
564 : static OGRSpatialReference IdentifyCRS(const OGRSpatialReference *poSRS);
565 : };
566 :
567 : /************************************************************************/
568 : /* OGRGeoArrowWkbExtensionType */
569 : /************************************************************************/
570 :
571 : #if ARROW_VERSION_MAJOR >= 21
572 :
573 : class OGRGeoArrowWkbExtensionType final : public arrow::ExtensionType
574 : {
575 : public:
576 : explicit OGRGeoArrowWkbExtensionType(
577 : const std::shared_ptr<arrow::DataType> &storage_type,
578 : const std::string &metadata)
579 : : arrow::ExtensionType(storage_type), metadata_(metadata)
580 : {
581 : }
582 :
583 : std::string extension_name() const override
584 : {
585 : return EXTENSION_NAME_GEOARROW_WKB;
586 : }
587 :
588 : bool ExtensionEquals(const arrow::ExtensionType &other) const override
589 : {
590 : return extension_name() == other.extension_name() &&
591 : storage_type_->Equals(other.storage_type()) &&
592 : Serialize() == other.Serialize();
593 : }
594 :
595 : arrow::Result<std::shared_ptr<arrow::DataType>>
596 : Deserialize(std::shared_ptr<arrow::DataType> storage_type,
597 : const std::string &serialized) const override
598 : {
599 : return Make(std::move(storage_type), serialized);
600 : }
601 :
602 : std::string Serialize() const override
603 : {
604 : return metadata_;
605 : }
606 :
607 : std::shared_ptr<arrow::Array>
608 : MakeArray(std::shared_ptr<arrow::ArrayData> data) const override
609 : {
610 : CPLAssert(data->type->id() == arrow::Type::EXTENSION);
611 : CPLAssert(EXTENSION_NAME_GEOARROW_WKB ==
612 : static_cast<const arrow::ExtensionType &>(*data->type)
613 : .extension_name());
614 : return std::make_shared<arrow::ExtensionArray>(data);
615 : }
616 :
617 : static bool IsSupportedStorageType(arrow::Type::type typeId)
618 : {
619 : // TODO: also add BINARY_VIEW if we support it some day.
620 : return typeId == arrow::Type::BINARY ||
621 : typeId == arrow::Type::LARGE_BINARY;
622 : }
623 :
624 : static arrow::Result<std::shared_ptr<arrow::DataType>>
625 : Make(std::shared_ptr<arrow::DataType> storage_type,
626 : const std::string &metadata)
627 : {
628 : if (!IsSupportedStorageType(storage_type->id()))
629 : {
630 : return arrow::Status::Invalid(
631 : "Invalid storage type for OGRGeoArrowWkbExtensionType: ",
632 : storage_type->ToString());
633 : }
634 : return std::make_shared<OGRGeoArrowWkbExtensionType>(
635 : std::move(storage_type), metadata);
636 : }
637 :
638 : private:
639 : std::string metadata_{};
640 : };
641 :
642 : #endif // ARROW_VERSION_MAJOR >= 21
643 :
644 : #if defined(__clang__)
645 : #pragma clang diagnostic pop
646 : #endif
647 :
648 : #endif // OGR_ARROW_H
|
