Line data Source code
1 : /******************************************************************************
2 : *
3 : * Project: Parquet Translator
4 : * Purpose: Implements OGRParquetDriver.
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_PARQUET_H
14 : #define OGR_PARQUET_H
15 :
16 : #include "ogrsf_frmts.h"
17 :
18 : #include "cpl_json.h"
19 :
20 : #include <functional>
21 : #include <map>
22 : #include <set>
23 :
24 : #include "../arrow_common/ogr_arrow.h"
25 : #include "ogr_include_parquet.h"
26 :
27 : constexpr int DEFAULT_COMPRESSION_LEVEL = -1;
28 : constexpr int OGR_PARQUET_ZSTD_DEFAULT_COMPRESSION_LEVEL = 9;
29 :
30 : /************************************************************************/
31 : /* OGRParquetLayerBase */
32 : /************************************************************************/
33 :
34 : class OGRParquetDataset;
35 :
36 : class OGRParquetLayerBase CPL_NON_FINAL : public OGRArrowLayer
37 : {
38 : OGRParquetLayerBase(const OGRParquetLayerBase &) = delete;
39 : OGRParquetLayerBase &operator=(const OGRParquetLayerBase &) = delete;
40 :
41 : protected:
42 : OGRParquetLayerBase(OGRParquetDataset *poDS, const char *pszLayerName,
43 : CSLConstList papszOpenOptions);
44 :
45 : OGRParquetDataset *m_poDS = nullptr;
46 : std::shared_ptr<arrow::RecordBatchReader> m_poRecordBatchReader{};
47 : CPLStringList m_aosGeomPossibleNames{};
48 : std::string m_osCRS{};
49 :
50 : #if PARQUET_VERSION_MAJOR >= 21
51 : std::set<int>
52 : m_geoStatsWithBBOXAvailable{}; // key is index of OGR geometry column
53 : #endif
54 :
55 : //! Only for testing/validation purposes, with Arrow >= 21, value of 'crs'
56 : // field from GeometryLogicalType/GeographyLogicalType
57 : std::map<std::string, std::string> m_mapGeomFieldToParquetGeoCrs{};
58 :
59 : void LoadGeoMetadata(
60 : const std::shared_ptr<const arrow::KeyValueMetadata> &kv_metadata);
61 : bool DealWithGeometryColumn(
62 : int iFieldIdx, const std::shared_ptr<arrow::Field> &field,
63 : std::function<OGRwkbGeometryType(void)> computeGeometryTypeFun,
64 : const parquet::ColumnDescriptor *parquetColumn,
65 : const parquet::FileMetaData *fileMetadata, int iColumn);
66 :
67 : #if PARQUET_VERSION_MAJOR >= 21
68 : bool DealWithArrow21GeometryGeographyNativeTypes(
69 : int iFieldIdx, const std::shared_ptr<arrow::Field> &field,
70 : const parquet::ColumnDescriptor *parquetColumn,
71 : const parquet::FileMetaData *fileMetadata, int iColumn);
72 : #endif
73 :
74 : void InvalidateCachedBatches() override;
75 :
76 : static bool ParseGeometryColumnCovering(const CPLJSONObject &oJSONDef,
77 : std::string &osBBOXColumn,
78 : std::string &osXMin,
79 : std::string &osYMin,
80 : std::string &osXMax,
81 : std::string &osYMax);
82 :
83 : public:
84 : bool TestCapability(const char *) const override;
85 :
86 : void ResetReading() override;
87 :
88 : GDALDataset *GetDataset() override;
89 :
90 : static int GetNumCPUs();
91 : };
92 :
93 : /************************************************************************/
94 : /* OGRParquetLayer */
95 : /************************************************************************/
96 :
97 : class OGRParquetLayer final : public OGRParquetLayerBase
98 :
99 : {
100 : std::unique_ptr<parquet::arrow::FileReader> m_poArrowReader{};
101 : bool m_bSingleBatch = false;
102 : int m_iFIDParquetColumn = -1;
103 : std::shared_ptr<arrow::DataType> m_poFIDType{};
104 : std::vector<std::shared_ptr<arrow::DataType>>
105 : m_apoArrowDataTypes{}; // .size() == field ocunt
106 : std::vector<int> m_anMapFieldIndexToParquetColumn{};
107 : std::vector<std::vector<int>> m_anMapGeomFieldIndexToParquetColumns{};
108 : bool m_bHasMissingMappingToParquet = false;
109 :
110 : //! Contains pairs of (selected feature idx, total feature idx) break points.
111 : std::vector<std::pair<int64_t, int64_t>> m_asFeatureIdxRemapping{};
112 : //! Iterator over m_asFeatureIdxRemapping
113 : std::vector<std::pair<int64_t, int64_t>>::iterator
114 : m_oFeatureIdxRemappingIter{};
115 : //! Feature index among the potentially restricted set of selected row groups
116 : int64_t m_nFeatureIdxSelected = 0;
117 : std::vector<int> m_anRequestedParquetColumns{}; // only valid when
118 : // m_bIgnoredFields is set
119 : CPLStringList m_aosFeatherMetadata{};
120 :
121 : //! Describe the bbox column of a geometry column
122 : struct GeomColBBOXParquet
123 : {
124 : int iParquetXMin = -1;
125 : int iParquetYMin = -1;
126 : int iParquetXMax = -1;
127 : int iParquetYMax = -1;
128 : std::vector<int> anParquetCols{};
129 : };
130 :
131 : //! Map from OGR geometry field index to GeomColBBOXParquet
132 : std::map<int, GeomColBBOXParquet>
133 : m_oMapGeomFieldIndexToGeomColBBOXParquet{};
134 :
135 : //! GDAL creation options that were used to create the file (if done by GDAL)
136 : CPLStringList m_aosCreationOptions{};
137 :
138 : void EstablishFeatureDefn();
139 : void ProcessGeometryColumnCovering(
140 : const std::shared_ptr<arrow::Field> &field,
141 : const CPLJSONObject &oJSONGeometryColumn,
142 : const std::map<std::string, int> &oMapParquetColumnNameToIdx);
143 : bool CreateRecordBatchReader(int iStartingRowGroup);
144 : bool CreateRecordBatchReader(const std::vector<int> &anRowGroups);
145 : bool ReadNextBatch() override;
146 :
147 : void InvalidateCachedBatches() override;
148 :
149 : OGRwkbGeometryType ComputeGeometryColumnType(int iGeomCol,
150 : int iParquetCol) const;
151 : void CreateFieldFromSchema(
152 : const std::shared_ptr<arrow::Field> &field, bool bParquetColValid,
153 : int &iParquetCol, const std::vector<int> &path,
154 : const std::map<std::string, std::unique_ptr<OGRFieldDefn>>
155 : &oMapFieldNameToGDALSchemaFieldDefn);
156 : bool CheckMatchArrowParquetColumnNames(
157 : int &iParquetCol, const std::shared_ptr<arrow::Field> &field) const;
158 : OGRFeature *GetFeatureExplicitFID(GIntBig nFID);
159 : OGRFeature *GetFeatureByIndex(GIntBig nFID);
160 :
161 28780 : std::string GetDriverUCName() const override
162 : {
163 28780 : return "PARQUET";
164 : }
165 :
166 : bool FastGetExtent(int iGeomField, OGREnvelope *psExtent) const override;
167 :
168 : void IncrFeatureIdx() override;
169 :
170 : public:
171 : OGRParquetLayer(OGRParquetDataset *poDS, const char *pszLayerName,
172 : std::unique_ptr<parquet::arrow::FileReader> &&arrow_reader,
173 : CSLConstList papszOpenOptions);
174 :
175 : void ResetReading() override;
176 : OGRFeature *GetFeature(GIntBig nFID) override;
177 : GIntBig GetFeatureCount(int bForce) override;
178 : bool TestCapability(const char *pszCap) const override;
179 : OGRErr SetIgnoredFields(CSLConstList papszFields) override;
180 : const char *GetMetadataItem(const char *pszName,
181 : const char *pszDomain = "") override;
182 : CSLConstList GetMetadata(const char *pszDomain = "") override;
183 : OGRErr SetNextByIndex(GIntBig nIndex) override;
184 :
185 : bool GetArrowStream(struct ArrowArrayStream *out_stream,
186 : CSLConstList papszOptions = nullptr) override;
187 :
188 : std::unique_ptr<OGRFieldDomain> BuildDomain(const std::string &osDomainName,
189 : int iFieldIndex) const override;
190 :
191 1546 : parquet::arrow::FileReader *GetReader() const
192 : {
193 1546 : return m_poArrowReader.get();
194 : }
195 :
196 : std::vector<int>
197 : GetParquetColumnIndicesForArrowField(const std::string &field_name) const;
198 :
199 : const std::vector<std::shared_ptr<arrow::DataType>> &
200 249 : GetArrowFieldTypes() const
201 : {
202 249 : return m_apoArrowDataTypes;
203 : }
204 :
205 2 : int GetFIDParquetColumn() const
206 : {
207 2 : return m_iFIDParquetColumn;
208 : }
209 :
210 4 : const CPLStringList &GetCreationOptions() const
211 : {
212 4 : return m_aosCreationOptions;
213 : }
214 :
215 : static constexpr int OGR_FID_INDEX = -2;
216 : bool GetMinMaxForOGRField(int iRowGroup, // -1 for all
217 : int iOGRField, // or OGR_FID_INDEX
218 : bool bComputeMin, OGRField &sMin, bool &bFoundMin,
219 : bool bComputeMax, OGRField &sMax, bool &bFoundMax,
220 : OGRFieldType &eType, OGRFieldSubType &eSubType,
221 : std::string &osMinTmp,
222 : std::string &osMaxTmp) const;
223 :
224 : bool GetMinMaxForParquetCol(int iRowGroup, // -1 for all
225 : int iCol,
226 : const std::shared_ptr<arrow::DataType>
227 : &arrowType, // potentially nullptr
228 : bool bComputeMin, OGRField &sMin,
229 : bool &bFoundMin, bool bComputeMax,
230 : OGRField &sMax, bool &bFoundMax,
231 : OGRFieldType &eType, OGRFieldSubType &eSubType,
232 : std::string &osMinTmp,
233 : std::string &osMaxTmp) const;
234 :
235 : bool GeomColsBBOXParquet(int iGeom, int &iParquetXMin, int &iParquetYMin,
236 : int &iParquetXMax, int &iParquetYMax) const;
237 : };
238 :
239 : /************************************************************************/
240 : /* OGRParquetDatasetLayer */
241 : /************************************************************************/
242 :
243 : #ifdef GDAL_USE_ARROWDATASET
244 :
245 : class OGRParquetDatasetLayer final : public OGRParquetLayerBase
246 : {
247 : bool m_bIsVSI = false;
248 : bool m_bRebuildScanner = true;
249 : bool m_bSkipFilterGeometry = false;
250 : std::shared_ptr<arrow::dataset::Dataset> m_poDataset{};
251 : std::shared_ptr<arrow::dataset::Scanner> m_poScanner{};
252 : std::vector<std::string> m_aosProjectedFields{};
253 :
254 : void EstablishFeatureDefn();
255 : void
256 : ProcessGeometryColumnCovering(const std::shared_ptr<arrow::Field> &field,
257 : const CPLJSONObject &oJSONGeometryColumn);
258 :
259 : void BuildScanner();
260 :
261 : //! Translate a OGR SQL expression into an Arrow one
262 : // bFullyTranslated should be set to true before calling this method.
263 : arrow::compute::Expression BuildArrowFilter(const swq_expr_node *poNode,
264 : bool &bFullyTranslated);
265 :
266 : protected:
267 7547 : std::string GetDriverUCName() const override
268 : {
269 7547 : return "PARQUET";
270 : }
271 :
272 : bool ReadNextBatch() override;
273 :
274 : bool FastGetExtent(int iGeomField, OGREnvelope *psExtent) const override;
275 :
276 : public:
277 : OGRParquetDatasetLayer(
278 : OGRParquetDataset *poDS, const char *pszLayerName, bool bIsVSI,
279 : const std::shared_ptr<arrow::dataset::Dataset> &dataset,
280 : CSLConstList papszOpenOptions);
281 :
282 : OGRFeature *GetNextFeature() override;
283 :
284 : GIntBig GetFeatureCount(int bForce) override;
285 : OGRErr IGetExtent(int iGeomField, OGREnvelope *psExtent,
286 : bool bForce) override;
287 :
288 : OGRErr ISetSpatialFilter(int iGeomField,
289 : const OGRGeometry *poGeom) override;
290 :
291 : OGRErr SetAttributeFilter(const char *pszFilter) override;
292 :
293 : OGRErr SetIgnoredFields(CSLConstList papszFields) override;
294 :
295 : bool TestCapability(const char *) const override;
296 :
297 : // TODO
298 : std::unique_ptr<OGRFieldDomain>
299 0 : BuildDomain(const std::string & /*osDomainName*/,
300 : int /*iFieldIndex*/) const override
301 : {
302 0 : return nullptr;
303 : }
304 : };
305 :
306 : #endif
307 :
308 : /************************************************************************/
309 : /* OGRParquetDataset */
310 : /************************************************************************/
311 :
312 : class OGRParquetDataset final : public OGRArrowDataset
313 : {
314 : std::shared_ptr<arrow::fs::FileSystem> m_poFS{};
315 :
316 : public:
317 : explicit OGRParquetDataset();
318 : ~OGRParquetDataset() override;
319 :
320 : CPLErr Close(GDALProgressFunc = nullptr, void * = nullptr) override;
321 :
322 : OGRLayer *ExecuteSQL(const char *pszSQLCommand,
323 : OGRGeometry *poSpatialFilter,
324 : const char *pszDialect) override;
325 : void ReleaseResultSet(OGRLayer *poResultsSet) override;
326 :
327 : bool TestCapability(const char *) const override;
328 :
329 360 : void SetFileSystem(const std::shared_ptr<arrow::fs::FileSystem> &fs)
330 : {
331 360 : m_poFS = fs;
332 360 : }
333 :
334 : std::unique_ptr<OGRParquetLayer>
335 : CreateReaderLayer(const std::string &osFilename, VSILFILE *&fpIn,
336 : CSLConstList papszOpenOptionsIn);
337 : };
338 :
339 : /************************************************************************/
340 : /* OGRGeoParquetVersion */
341 : /************************************************************************/
342 :
343 : enum class OGRGeoParquetVersion
344 : {
345 : VERSION_1_1,
346 : VERSION_2_0,
347 : };
348 :
349 : /************************************************************************/
350 : /* OGRParquetWriterLayer */
351 : /************************************************************************/
352 :
353 : class OGRParquetWriterDataset;
354 :
355 : class OGRParquetWriterLayer final : public OGRArrowWriterLayer
356 : {
357 : OGRParquetWriterLayer(const OGRParquetWriterLayer &) = delete;
358 : OGRParquetWriterLayer &operator=(const OGRParquetWriterLayer &) = delete;
359 :
360 : OGRParquetWriterDataset *m_poDataset = nullptr;
361 : std::unique_ptr<parquet::arrow::FileWriter> m_poFileWriter{};
362 : std::shared_ptr<const arrow::KeyValueMetadata> m_poKeyValueMetadata{};
363 : bool m_bForceCounterClockwiseOrientation = false;
364 : parquet::WriterProperties::Builder m_oWriterPropertiesBuilder{};
365 :
366 : //! Temporary GeoPackage dataset. Only used in SORT_BY_BBOX mode
367 : std::unique_ptr<GDALDataset> m_poTmpGPKG{};
368 : //! Temporary GeoPackage layer. Only used in SORT_BY_BBOX mode
369 : OGRLayer *m_poTmpGPKGLayer = nullptr;
370 : //! Number of features written by ICreateFeature(). Only used in SORT_BY_BBOX mode
371 : GIntBig m_nTmpFeatureCount = 0;
372 :
373 : //! Whether to write "geo" footer metadata;
374 : bool m_bWriteGeoMetadata = true;
375 :
376 : OGRGeoParquetVersion m_nGeoParquetVersion =
377 : OGRGeoParquetVersion::VERSION_1_1;
378 :
379 895 : bool IsFileWriterCreated() const override
380 : {
381 895 : return m_poFileWriter != nullptr;
382 : }
383 :
384 : void CreateWriter() override;
385 : bool CloseFileWriter() override;
386 :
387 : void CreateSchema() override;
388 : void PerformStepsBeforeFinalFlushGroup() override;
389 :
390 : bool FlushGroup() override;
391 :
392 383 : std::string GetDriverUCName() const override
393 : {
394 383 : return "PARQUET";
395 : }
396 :
397 : virtual bool
398 : IsSupportedGeometryType(OGRwkbGeometryType eGType) const override;
399 :
400 : virtual void FixupWKBGeometryBeforeWriting(GByte *pabyWKB,
401 : size_t nLen) override;
402 : void FixupGeometryBeforeWriting(OGRGeometry *poGeom) override;
403 :
404 26 : bool IsSRSRequired() const override
405 : {
406 26 : return false;
407 : }
408 :
409 : std::string GetGeoMetadata() const;
410 :
411 : //! Copy temporary GeoPackage layer to final Parquet file
412 : bool CopyTmpGpkgLayerToFinalFile();
413 :
414 : public:
415 : OGRParquetWriterLayer(
416 : OGRParquetWriterDataset *poDS, arrow::MemoryPool *poMemoryPool,
417 : const std::shared_ptr<arrow::io::OutputStream> &poOutputStream,
418 : const char *pszLayerName);
419 :
420 : CPLErr SetMetadata(CSLConstList papszMetadata,
421 : const char *pszDomain) override;
422 :
423 : bool SetOptions(const OGRGeomFieldDefn *poSrcGeomFieldDefn,
424 : CSLConstList papszOptions);
425 :
426 : OGRErr CreateGeomField(const OGRGeomFieldDefn *poField,
427 : int bApproxOK = TRUE) override;
428 :
429 : bool TestCapability(const char *pszCap) const override;
430 : #if PARQUET_VERSION_MAJOR <= 10
431 : // Parquet <= 10 doesn't support the WriteRecordBatch() API
432 : bool IsArrowSchemaSupported(const struct ArrowSchema *schema,
433 : CSLConstList papszOptions,
434 : std::string &osErrorMsg) const override
435 : {
436 : return OGRLayer::IsArrowSchemaSupported(schema, papszOptions,
437 : osErrorMsg);
438 : }
439 :
440 : bool
441 : CreateFieldFromArrowSchema(const struct ArrowSchema *schema,
442 : CSLConstList papszOptions = nullptr) override
443 : {
444 : return OGRLayer::CreateFieldFromArrowSchema(schema, papszOptions);
445 : }
446 :
447 : bool WriteArrowBatch(const struct ArrowSchema *schema,
448 : struct ArrowArray *array,
449 : CSLConstList papszOptions = nullptr) override
450 : {
451 : return OGRLayer::WriteArrowBatch(schema, array, papszOptions);
452 : }
453 : #else
454 : bool IsArrowSchemaSupported(const struct ArrowSchema *schema,
455 : CSLConstList papszOptions,
456 : std::string &osErrorMsg) const override;
457 : bool
458 : CreateFieldFromArrowSchema(const struct ArrowSchema *schema,
459 : CSLConstList papszOptions = nullptr) override;
460 : bool WriteArrowBatch(const struct ArrowSchema *schema,
461 : struct ArrowArray *array,
462 : CSLConstList papszOptions = nullptr) override;
463 : #endif
464 :
465 : GDALDataset *GetDataset() override;
466 :
467 : protected:
468 : OGRErr ICreateFeature(OGRFeature *poFeature) override;
469 :
470 : friend class OGRParquetWriterDataset;
471 : bool Close();
472 : };
473 :
474 : /************************************************************************/
475 : /* OGRParquetWriterDataset */
476 : /************************************************************************/
477 :
478 : class OGRParquetWriterDataset final : public GDALPamDataset
479 : {
480 : std::unique_ptr<arrow::MemoryPool> m_poMemoryPool{};
481 : std::unique_ptr<OGRParquetWriterLayer> m_poLayer{};
482 : std::shared_ptr<arrow::io::OutputStream> m_poOutputStream{};
483 :
484 : public:
485 : explicit OGRParquetWriterDataset(
486 : const std::shared_ptr<arrow::io::OutputStream> &poOutputStream);
487 :
488 : ~OGRParquetWriterDataset() override;
489 :
490 : arrow::MemoryPool *GetMemoryPool() const
491 : {
492 : return m_poMemoryPool.get();
493 : }
494 :
495 : CPLErr Close(GDALProgressFunc = nullptr, void * = nullptr) override;
496 :
497 : int GetLayerCount() const override;
498 : const OGRLayer *GetLayer(int idx) const override;
499 : bool TestCapability(const char *pszCap) const override;
500 : std::vector<std::string> GetFieldDomainNames(
501 : CSLConstList /*papszOptions*/ = nullptr) const override;
502 : const OGRFieldDomain *
503 : GetFieldDomain(const std::string &name) const override;
504 : bool AddFieldDomain(std::unique_ptr<OGRFieldDomain> &&domain,
505 : std::string &failureReason) override;
506 :
507 353 : GDALMultiDomainMetadata &GetMultiDomainMetadata()
508 : {
509 353 : return oMDMD;
510 : }
511 :
512 : protected:
513 : OGRLayer *ICreateLayer(const char *pszName,
514 : const OGRGeomFieldDefn *poGeomFieldDefn,
515 : CSLConstList papszOptions) override;
516 : };
517 :
518 : #endif // OGR_PARQUET_H
|
