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