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