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-2024, Planet Labs
9 : *
10 : * SPDX-License-Identifier: MIT
11 : ****************************************************************************/
12 :
13 : #include "ogrsf_frmts.h"
14 :
15 : #include <algorithm>
16 : #include <cassert>
17 : #include <map>
18 : #include <set>
19 : #include <utility>
20 :
21 : #include "cpl_time.h"
22 : #include "ogr_api.h"
23 :
24 : #include "ogr_parquet.h"
25 :
26 : #include "../arrow_common/ograrrowlayer.hpp"
27 : #include "../arrow_common/ograrrowdataset.hpp"
28 :
29 : #if PARQUET_VERSION_MAJOR >= 13
30 : // Using field indices for FieldRef is only supported since
31 : // https://github.com/apache/arrow/commit/10eedbe63c71f4cf8f0621f3a2304ab3168a2ae5
32 : #define SUPPORTS_INDICES_IN_FIELD_REF
33 : #endif
34 :
35 : namespace cp = ::arrow::compute;
36 :
37 : /************************************************************************/
38 : /* OGRParquetLayer() */
39 : /************************************************************************/
40 :
41 360 : OGRParquetDatasetLayer::OGRParquetDatasetLayer(
42 : OGRParquetDataset *poDS, const char *pszLayerName, bool bIsVSI,
43 : const std::shared_ptr<arrow::dataset::Dataset> &dataset,
44 360 : CSLConstList papszOpenOptions)
45 : : OGRParquetLayerBase(poDS, pszLayerName, papszOpenOptions),
46 360 : m_bIsVSI(bIsVSI), m_poDataset(dataset)
47 : {
48 360 : m_poSchema = m_poDataset->schema();
49 360 : EstablishFeatureDefn();
50 360 : CPLAssert(static_cast<int>(m_aeGeomEncoding.size()) ==
51 : m_poFeatureDefn->GetGeomFieldCount());
52 360 : }
53 :
54 : /************************************************************************/
55 : /* ProcessGeometryColumnCovering() */
56 : /************************************************************************/
57 :
58 : /** Process GeoParquet JSON geometry field object to extract information about
59 : * its bounding box column, and appropriately fill m_oMapGeomFieldIndexToGeomColBBOX
60 : * member with information on that bounding box column.
61 : */
62 337 : void OGRParquetDatasetLayer::ProcessGeometryColumnCovering(
63 : const std::shared_ptr<arrow::Field> &field,
64 : const CPLJSONObject &oJSONGeometryColumn)
65 : {
66 674 : std::string osBBOXColumn;
67 674 : std::string osXMin, osYMin, osXMax, osYMax;
68 337 : if (ParseGeometryColumnCovering(oJSONGeometryColumn, osBBOXColumn, osXMin,
69 : osYMin, osXMax, osYMax))
70 : {
71 173 : OGRArrowLayer::GeomColBBOX sDesc;
72 173 : sDesc.iArrowCol = m_poSchema->GetFieldIndex(osBBOXColumn);
73 346 : const auto fieldBBOX = m_poSchema->GetFieldByName(osBBOXColumn);
74 346 : if (sDesc.iArrowCol >= 0 && fieldBBOX &&
75 173 : fieldBBOX->type()->id() == arrow::Type::STRUCT)
76 : {
77 : const auto fieldBBOXStruct =
78 346 : std::static_pointer_cast<arrow::StructType>(fieldBBOX->type());
79 346 : const auto fieldXMin = fieldBBOXStruct->GetFieldByName(osXMin);
80 346 : const auto fieldYMin = fieldBBOXStruct->GetFieldByName(osYMin);
81 346 : const auto fieldXMax = fieldBBOXStruct->GetFieldByName(osXMax);
82 346 : const auto fieldYMax = fieldBBOXStruct->GetFieldByName(osYMax);
83 173 : const int nXMinIdx = fieldBBOXStruct->GetFieldIndex(osXMin);
84 173 : const int nYMinIdx = fieldBBOXStruct->GetFieldIndex(osYMin);
85 173 : const int nXMaxIdx = fieldBBOXStruct->GetFieldIndex(osXMax);
86 173 : const int nYMaxIdx = fieldBBOXStruct->GetFieldIndex(osYMax);
87 173 : if (nXMinIdx >= 0 && nYMinIdx >= 0 && nXMaxIdx >= 0 &&
88 346 : nYMaxIdx >= 0 && fieldXMin && fieldYMin && fieldXMax &&
89 346 : fieldYMax &&
90 173 : (fieldXMin->type()->id() == arrow::Type::FLOAT ||
91 0 : fieldXMin->type()->id() == arrow::Type::DOUBLE) &&
92 173 : fieldXMin->type()->id() == fieldYMin->type()->id() &&
93 519 : fieldXMin->type()->id() == fieldXMax->type()->id() &&
94 173 : fieldXMin->type()->id() == fieldYMax->type()->id())
95 : {
96 173 : CPLDebug("PARQUET",
97 : "Bounding box column '%s' detected for "
98 : "geometry column '%s'",
99 173 : osBBOXColumn.c_str(), field->name().c_str());
100 173 : sDesc.iArrowSubfieldXMin = nXMinIdx;
101 173 : sDesc.iArrowSubfieldYMin = nYMinIdx;
102 173 : sDesc.iArrowSubfieldXMax = nXMaxIdx;
103 173 : sDesc.iArrowSubfieldYMax = nYMaxIdx;
104 173 : sDesc.bIsFloat =
105 173 : (fieldXMin->type()->id() == arrow::Type::FLOAT);
106 :
107 : m_oMapGeomFieldIndexToGeomColBBOX
108 173 : [m_poFeatureDefn->GetGeomFieldCount() - 1] =
109 173 : std::move(sDesc);
110 : }
111 : }
112 : }
113 337 : }
114 :
115 : /************************************************************************/
116 : /* EstablishFeatureDefn() */
117 : /************************************************************************/
118 :
119 360 : void OGRParquetDatasetLayer::EstablishFeatureDefn()
120 : {
121 360 : const auto &kv_metadata = m_poSchema->metadata();
122 :
123 360 : LoadGeoMetadata(kv_metadata);
124 : const auto oMapFieldNameToGDALSchemaFieldDefn =
125 720 : LoadGDALSchema(kv_metadata.get());
126 :
127 360 : LoadGDALMetadata(kv_metadata.get());
128 :
129 : const bool bUseBBOX =
130 360 : CPLTestBool(CPLGetConfigOption("OGR_PARQUET_USE_BBOX", "YES"));
131 :
132 : // Keep track of declared bounding box columns in GeoParquet JSON metadata,
133 : // in order not to expose them as regular fields.
134 720 : std::set<std::string> oSetBBOXColumns;
135 360 : if (bUseBBOX)
136 : {
137 696 : for (const auto &iter : m_oMapGeometryColumns)
138 : {
139 674 : std::string osBBOXColumn;
140 674 : std::string osXMin, osYMin, osXMax, osYMax;
141 337 : if (ParseGeometryColumnCovering(iter.second, osBBOXColumn, osXMin,
142 : osYMin, osXMax, osYMax))
143 : {
144 172 : oSetBBOXColumns.insert(std::move(osBBOXColumn));
145 : }
146 : }
147 : }
148 :
149 360 : const auto &fields = m_poSchema->fields();
150 :
151 : // Overture Maps 2024-04-16-beta.0 almost follows GeoParquet 1.1, except
152 : // they don't declare the "covering" element in the GeoParquet JSON metadata
153 1058 : if (m_oMapGeometryColumns.find("geometry") != m_oMapGeometryColumns.end() &&
154 674 : bUseBBOX &&
155 1583 : !m_oMapGeometryColumns["geometry"].GetObj("covering").IsValid() &&
156 525 : m_oMapGeometryColumns["geometry"].GetString("encoding") == "WKB")
157 : {
158 6327 : for (int i = 0; i < m_poSchema->num_fields(); ++i)
159 : {
160 6247 : const auto &field = fields[i];
161 6248 : if (field->name() == "bbox" &&
162 1 : field->type()->id() == arrow::Type::STRUCT)
163 : {
164 1 : bool bXMin = false;
165 1 : bool bXMax = false;
166 1 : bool bYMin = false;
167 1 : bool bYMax = false;
168 2 : const auto subfields = field->Flatten();
169 1 : if (subfields.size() == 4)
170 : {
171 5 : for (int j = 0; j < static_cast<int>(subfields.size()); j++)
172 : {
173 4 : const auto &subfield = subfields[j];
174 4 : if (subfield->name() == "bbox.xmin")
175 1 : bXMin = true;
176 3 : else if (subfield->name() == "bbox.xmax")
177 1 : bXMax = true;
178 2 : else if (subfield->name() == "bbox.ymin")
179 1 : bYMin = true;
180 1 : else if (subfield->name() == "bbox.ymax")
181 1 : bYMax = true;
182 : }
183 : }
184 1 : if (bXMin && bXMax && bYMin && bYMax)
185 : {
186 3 : CPLJSONObject oDef = m_oMapGeometryColumns["geometry"];
187 2 : CPLJSONObject oCovering;
188 1 : oDef.Add("covering", oCovering);
189 1 : CPLJSONObject oBBOX;
190 1 : oCovering.Add("bbox", oBBOX);
191 : {
192 1 : CPLJSONArray oArray;
193 1 : oArray.Add("bbox");
194 1 : oArray.Add("xmin");
195 1 : oBBOX.Add("xmin", oArray);
196 : }
197 : {
198 1 : CPLJSONArray oArray;
199 1 : oArray.Add("bbox");
200 1 : oArray.Add("ymin");
201 1 : oBBOX.Add("ymin", oArray);
202 : }
203 : {
204 1 : CPLJSONArray oArray;
205 1 : oArray.Add("bbox");
206 1 : oArray.Add("xmax");
207 1 : oBBOX.Add("xmax", oArray);
208 : }
209 : {
210 1 : CPLJSONArray oArray;
211 1 : oArray.Add("bbox");
212 1 : oArray.Add("ymax");
213 1 : oBBOX.Add("ymax", oArray);
214 : }
215 1 : oSetBBOXColumns.insert("bbox");
216 1 : m_oMapGeometryColumns["geometry"] = std::move(oDef);
217 : }
218 1 : break;
219 : }
220 : }
221 : }
222 :
223 7510 : for (int i = 0; i < m_poSchema->num_fields(); ++i)
224 : {
225 7150 : const auto &field = fields[i];
226 :
227 7150 : if (!m_osFIDColumn.empty() && field->name() == m_osFIDColumn)
228 : {
229 2 : m_iFIDArrowColumn = i;
230 2 : continue;
231 : }
232 :
233 7148 : if (oSetBBOXColumns.find(field->name()) != oSetBBOXColumns.end())
234 : {
235 173 : m_oSetBBoxArrowColumns.insert(i);
236 173 : continue;
237 : }
238 :
239 6975 : const bool bGeometryField = DealWithGeometryColumn(
240 79 : i, field, []() { return wkbUnknown; }, nullptr, nullptr, -1);
241 6975 : if (bGeometryField)
242 : {
243 339 : const auto oIter = m_oMapGeometryColumns.find(field->name());
244 339 : if (bUseBBOX && oIter != m_oMapGeometryColumns.end())
245 : {
246 337 : ProcessGeometryColumnCovering(field, oIter->second);
247 : }
248 : }
249 : else
250 : {
251 6636 : CreateFieldFromSchema(field, {i},
252 : oMapFieldNameToGDALSchemaFieldDefn);
253 : }
254 : }
255 :
256 360 : CPLAssert(static_cast<int>(m_anMapFieldIndexToArrowColumn.size()) ==
257 : m_poFeatureDefn->GetFieldCount());
258 360 : CPLAssert(static_cast<int>(m_anMapGeomFieldIndexToArrowColumn.size()) ==
259 : m_poFeatureDefn->GetGeomFieldCount());
260 360 : }
261 :
262 : namespace
263 : {
264 :
265 : /************************************************************************/
266 : /* WKBGeometryOptionsType */
267 : /************************************************************************/
268 :
269 : class WKBGeometryOptions;
270 :
271 : class WKBGeometryOptionsType : public cp::FunctionOptionsType
272 : {
273 : WKBGeometryOptionsType() = default;
274 :
275 : static const WKBGeometryOptions &Cast(const cp::FunctionOptions &opts);
276 :
277 : public:
278 23 : const char *type_name() const override
279 : {
280 23 : return "WKBGeometryOptionsType";
281 : }
282 :
283 : std::string Stringify(const cp::FunctionOptions &) const override;
284 : bool Compare(const cp::FunctionOptions &,
285 : const cp::FunctionOptions &) const override;
286 : std::unique_ptr<cp::FunctionOptions>
287 : Copy(const cp::FunctionOptions &) const override;
288 :
289 28 : static WKBGeometryOptionsType *GetSingleton()
290 : {
291 28 : static WKBGeometryOptionsType instance;
292 28 : return &instance;
293 : }
294 : };
295 :
296 : /************************************************************************/
297 : /* WKBGeometryOptions */
298 : /************************************************************************/
299 :
300 : class WKBGeometryOptions : public cp::FunctionOptions
301 : {
302 :
303 : public:
304 22 : explicit WKBGeometryOptions(
305 : const std::vector<GByte> &abyFilterGeomWkbIn = std::vector<GByte>())
306 22 : : cp::FunctionOptions(WKBGeometryOptionsType::GetSingleton()),
307 22 : abyFilterGeomWkb(abyFilterGeomWkbIn)
308 : {
309 22 : }
310 :
311 3 : bool operator==(const WKBGeometryOptions &other) const
312 : {
313 3 : return abyFilterGeomWkb == other.abyFilterGeomWkb;
314 : }
315 :
316 : std::vector<GByte> abyFilterGeomWkb;
317 : };
318 :
319 : const WKBGeometryOptions &
320 32 : WKBGeometryOptionsType::Cast(const cp::FunctionOptions &opts)
321 : {
322 32 : return *cpl::down_cast<const WKBGeometryOptions *>(&opts);
323 : }
324 :
325 3 : bool WKBGeometryOptionsType::Compare(const cp::FunctionOptions &optsA,
326 : const cp::FunctionOptions &optsB) const
327 : {
328 3 : return Cast(optsA) == Cast(optsB);
329 : }
330 :
331 : std::string
332 23 : WKBGeometryOptionsType::Stringify(const cp::FunctionOptions &opts) const
333 : {
334 23 : const auto &bboxOptions = Cast(opts);
335 23 : std::string osRet(type_name());
336 23 : osRet += '-';
337 2162 : for (GByte byVal : bboxOptions.abyFilterGeomWkb)
338 2139 : osRet += CPLSPrintf("%02X", byVal);
339 23 : return osRet;
340 : }
341 :
342 : std::unique_ptr<cp::FunctionOptions>
343 3 : WKBGeometryOptionsType::Copy(const cp::FunctionOptions &opts) const
344 : {
345 3 : return std::make_unique<WKBGeometryOptions>(Cast(opts));
346 : }
347 :
348 : /************************************************************************/
349 : /* OptionsWrapper */
350 : /************************************************************************/
351 :
352 : /// KernelState adapter for the common case of kernels whose only
353 : /// state is an instance of a subclass of FunctionOptions.
354 : template <typename OptionsType> struct OptionsWrapper : public cp::KernelState
355 : {
356 38 : explicit OptionsWrapper(OptionsType optionsIn)
357 38 : : options(std::move(optionsIn))
358 : {
359 38 : }
360 :
361 : static arrow::Result<std::unique_ptr<cp::KernelState>>
362 38 : Init(cp::KernelContext *, const cp::KernelInitArgs &args)
363 : {
364 38 : auto options = cpl::down_cast<const OptionsType *>(args.options);
365 38 : CPLAssert(options);
366 38 : return std::make_unique<OptionsWrapper>(*options);
367 : }
368 :
369 32 : static const OptionsType &Get(cp::KernelContext *ctx)
370 : {
371 32 : return cpl::down_cast<const OptionsWrapper *>(ctx->state())->options;
372 : }
373 :
374 : OptionsType options;
375 : };
376 : } // namespace
377 :
378 : /************************************************************************/
379 : /* ExecOGRWKBIntersects() */
380 : /************************************************************************/
381 :
382 32 : static arrow::Status ExecOGRWKBIntersects(cp::KernelContext *ctx,
383 : const cp::ExecSpan &batch,
384 : cp::ExecResult *out)
385 : {
386 : // Get filter geometry
387 32 : const auto &opts = OptionsWrapper<WKBGeometryOptions>::Get(ctx);
388 32 : OGRGeometry *poGeomTmp = nullptr;
389 64 : OGRErr eErr = OGRGeometryFactory::createFromWkb(
390 32 : opts.abyFilterGeomWkb.data(), nullptr, &poGeomTmp,
391 32 : opts.abyFilterGeomWkb.size());
392 32 : CPL_IGNORE_RET_VAL(eErr);
393 32 : CPLAssert(eErr == OGRERR_NONE);
394 32 : CPLAssert(poGeomTmp != nullptr);
395 64 : std::unique_ptr<OGRGeometry> poFilterGeom(poGeomTmp);
396 32 : OGREnvelope sFilterEnvelope;
397 32 : poFilterGeom->getEnvelope(&sFilterEnvelope);
398 32 : const bool bFilterIsEnvelope = poFilterGeom->IsRectangle();
399 :
400 : // Deal with input array
401 32 : CPLAssert(batch.num_values() == 1);
402 32 : const arrow::ArraySpan &input = batch[0].array;
403 32 : CPLAssert(input.type->id() == arrow::Type::BINARY);
404 : // Packed array of bits
405 32 : const auto pabyInputValidity = input.buffers[0].data;
406 32 : const auto nInputOffsets = input.offset;
407 32 : const auto panWkbOffsets = input.GetValues<int32_t>(1);
408 32 : const auto pabyWkbArray = input.buffers[2].data;
409 :
410 : // Deal with output array
411 32 : CPLAssert(out->type()->id() == arrow::Type::BOOL);
412 32 : auto out_span = out->array_span();
413 : // Below array holds 8 bits per uint8_t
414 32 : uint8_t *pabitsOutValues = out_span->buffers[1].data;
415 32 : const auto nOutOffset = out_span->offset;
416 :
417 : // Iterate over WKB geometries
418 32 : OGRPreparedGeometry *pPreparedFilterGeom = nullptr;
419 32 : OGREnvelope sEnvelope;
420 148 : for (int64_t i = 0; i < batch.length; ++i)
421 : {
422 : const bool bInputIsNull =
423 152 : (pabyInputValidity &&
424 36 : arrow::bit_util::GetBit(pabyInputValidity, i + nInputOffsets) ==
425 116 : 0);
426 116 : bool bOutputVal = false;
427 116 : if (!bInputIsNull)
428 : {
429 104 : const GByte *pabyWkb = pabyWkbArray + panWkbOffsets[i];
430 104 : const size_t nWkbSize = panWkbOffsets[i + 1] - panWkbOffsets[i];
431 104 : bOutputVal = OGRLayer::FilterWKBGeometry(
432 : pabyWkb, nWkbSize,
433 : /* bEnvelopeAlreadySet = */ false, sEnvelope,
434 104 : poFilterGeom.get(), bFilterIsEnvelope, sFilterEnvelope,
435 : pPreparedFilterGeom);
436 : }
437 116 : if (bOutputVal)
438 56 : arrow::bit_util::SetBit(pabitsOutValues, i + nOutOffset);
439 : else
440 60 : arrow::bit_util::ClearBit(pabitsOutValues, i + nOutOffset);
441 : }
442 :
443 : // Cleanup
444 32 : if (pPreparedFilterGeom)
445 0 : OGRDestroyPreparedGeometry(pPreparedFilterGeom);
446 :
447 64 : return arrow::Status::OK();
448 : }
449 :
450 : /************************************************************************/
451 : /* RegisterOGRWKBIntersectsIfNeeded() */
452 : /************************************************************************/
453 :
454 19 : static bool RegisterOGRWKBIntersectsIfNeeded()
455 : {
456 19 : auto registry = cp::GetFunctionRegistry();
457 : bool bRet =
458 19 : registry->GetFunction("OGRWKBIntersects").ValueOr(nullptr) != nullptr;
459 19 : if (!bRet)
460 : {
461 3 : static const WKBGeometryOptions defaultOpts;
462 :
463 : // Below assert is completely useless but helps improve test coverage
464 3 : CPLAssert(WKBGeometryOptionsType::GetSingleton()->Compare(
465 : defaultOpts, *(WKBGeometryOptionsType::GetSingleton()
466 : ->Copy(defaultOpts)
467 : .get())));
468 :
469 : auto func = std::make_shared<cp::ScalarFunction>(
470 6 : "OGRWKBIntersects", cp::Arity::Unary(), cp::FunctionDoc(),
471 9 : &defaultOpts);
472 : cp::ScalarKernel kernel({arrow::binary()}, arrow::boolean(),
473 : ExecOGRWKBIntersects,
474 12 : OptionsWrapper<WKBGeometryOptions>::Init);
475 3 : kernel.null_handling = cp::NullHandling::OUTPUT_NOT_NULL;
476 9 : bRet = func->AddKernel(std::move(kernel)).ok() &&
477 6 : registry->AddFunction(std::move(func)).ok();
478 : }
479 19 : return bRet;
480 : }
481 :
482 : /************************************************************************/
483 : /* BuildScanner() */
484 : /************************************************************************/
485 :
486 941 : void OGRParquetDatasetLayer::BuildScanner()
487 : {
488 941 : m_bRebuildScanner = false;
489 941 : m_bSkipFilterGeometry = false;
490 941 : m_bBaseArrowIgnoreSpatialFilterRect = false;
491 941 : m_bBaseArrowIgnoreSpatialFilter = false;
492 941 : m_bBaseArrowIgnoreAttributeFilter = false;
493 :
494 : try
495 : {
496 941 : std::shared_ptr<arrow::dataset::ScannerBuilder> scannerBuilder;
497 1882 : PARQUET_ASSIGN_OR_THROW(scannerBuilder, m_poDataset->NewScan());
498 941 : assert(scannerBuilder);
499 :
500 : // We cannot use the shared memory pool. Otherwise we get random
501 : // crashes in multi-threaded arrow code (apparently some cleanup code),
502 : // that may used the memory pool after it has been destroyed.
503 : // At least this was true with some older libarrow version
504 : // PARQUET_THROW_NOT_OK(scannerBuilder->Pool(m_poMemoryPool));
505 :
506 941 : if (m_bIsVSI)
507 : {
508 377 : const int nFragmentReadAhead = atoi(
509 : CPLGetConfigOption("OGR_PARQUET_FRAGMENT_READ_AHEAD", "2"));
510 754 : PARQUET_THROW_NOT_OK(
511 : scannerBuilder->FragmentReadahead(nFragmentReadAhead));
512 : }
513 :
514 : const char *pszBatchSize =
515 941 : CPLGetConfigOption("OGR_PARQUET_BATCH_SIZE", nullptr);
516 941 : if (pszBatchSize)
517 : {
518 0 : PARQUET_THROW_NOT_OK(
519 : scannerBuilder->BatchSize(CPLAtoGIntBig(pszBatchSize)));
520 : }
521 :
522 941 : const int nNumCPUs = GetNumCPUs();
523 : const char *pszUseThreads =
524 941 : CPLGetConfigOption("OGR_PARQUET_USE_THREADS", nullptr);
525 941 : if (!pszUseThreads && nNumCPUs > 1)
526 : {
527 941 : pszUseThreads = "YES";
528 : }
529 941 : if (pszUseThreads && CPLTestBool(pszUseThreads))
530 : {
531 1882 : PARQUET_THROW_NOT_OK(scannerBuilder->UseThreads(true));
532 : }
533 :
534 : #if PARQUET_VERSION_MAJOR >= 10
535 : const char *pszBatchReadAhead =
536 941 : CPLGetConfigOption("OGR_PARQUET_BATCH_READ_AHEAD", nullptr);
537 941 : if (pszBatchReadAhead)
538 : {
539 0 : PARQUET_THROW_NOT_OK(
540 : scannerBuilder->BatchReadahead(atoi(pszBatchReadAhead)));
541 : }
542 : #endif
543 :
544 941 : cp::Expression expression;
545 1526 : if (m_poFilterGeom && !m_poFilterGeom->IsEmpty() &&
546 585 : CPLTestBool(CPLGetConfigOption(
547 : "OGR_PARQUET_OPTIMIZED_SPATIAL_FILTER", "YES")))
548 : {
549 : const auto oIter =
550 585 : m_oMapGeomFieldIndexToGeomColBBOX.find(m_iGeomFieldFilter);
551 585 : if (oIter != m_oMapGeomFieldIndexToGeomColBBOX.end())
552 : {
553 : // This actually requires Arrow >= 15 (https://github.com/apache/arrow/issues/39064)
554 : // to be more efficient.
555 : #ifdef SUPPORTS_INDICES_IN_FIELD_REF
556 380 : const auto &oBBOXDef = oIter->second;
557 3420 : expression = cp::and_(
558 : {cp::less_equal(
559 760 : cp::field_ref(arrow::FieldRef(
560 380 : oBBOXDef.iArrowCol, oBBOXDef.iArrowSubfieldXMin)),
561 760 : cp::literal(m_sFilterEnvelope.MaxX)),
562 : cp::less_equal(
563 760 : cp::field_ref(arrow::FieldRef(
564 380 : oBBOXDef.iArrowCol, oBBOXDef.iArrowSubfieldYMin)),
565 760 : cp::literal(m_sFilterEnvelope.MaxY)),
566 : cp::greater_equal(
567 760 : cp::field_ref(arrow::FieldRef(
568 380 : oBBOXDef.iArrowCol, oBBOXDef.iArrowSubfieldXMax)),
569 760 : cp::literal(m_sFilterEnvelope.MinX)),
570 : cp::greater_equal(
571 760 : cp::field_ref(arrow::FieldRef(
572 380 : oBBOXDef.iArrowCol, oBBOXDef.iArrowSubfieldYMax)),
573 2660 : cp::literal(m_sFilterEnvelope.MinY))});
574 : #else
575 : const auto oIter2 = m_oMapGeometryColumns.find(
576 : m_poFeatureDefn->GetGeomFieldDefn(m_iGeomFieldFilter)
577 : ->GetNameRef());
578 : std::string osBBOXColumn;
579 : std::string osXMin, osYMin, osXMax, osYMax;
580 : if (ParseGeometryColumnCovering(oIter2->second, osBBOXColumn,
581 : osXMin, osYMin, osXMax, osYMax))
582 : {
583 : expression = cp::and_(
584 : {cp::less_equal(cp::field_ref(arrow::FieldRef(
585 : osBBOXColumn, osXMin)),
586 : cp::literal(m_sFilterEnvelope.MaxX)),
587 : cp::less_equal(cp::field_ref(arrow::FieldRef(
588 : osBBOXColumn, osYMin)),
589 : cp::literal(m_sFilterEnvelope.MaxY)),
590 : cp::greater_equal(cp::field_ref(arrow::FieldRef(
591 : osBBOXColumn, osXMax)),
592 : cp::literal(m_sFilterEnvelope.MinX)),
593 : cp::greater_equal(
594 : cp::field_ref(
595 : arrow::FieldRef(osBBOXColumn, osYMax)),
596 : cp::literal(m_sFilterEnvelope.MinY))});
597 : }
598 : #endif
599 : }
600 615 : else if (m_iGeomFieldFilter >= 0 &&
601 205 : m_iGeomFieldFilter <
602 410 : static_cast<int>(m_aeGeomEncoding.size()) &&
603 205 : m_aeGeomEncoding[m_iGeomFieldFilter] ==
604 : OGRArrowGeomEncoding::GEOARROW_STRUCT_POINT)
605 : {
606 : const int iCol =
607 28 : m_anMapGeomFieldIndexToArrowColumn[m_iGeomFieldFilter];
608 28 : const auto &field = m_poSchema->fields()[iCol];
609 56 : auto type = field->type();
610 56 : std::vector<arrow::FieldRef> fieldRefs;
611 : #ifdef SUPPORTS_INDICES_IN_FIELD_REF
612 28 : fieldRefs.emplace_back(iCol);
613 : #else
614 : fieldRefs.emplace_back(field->name());
615 : #endif
616 28 : if (type->id() == arrow::Type::STRUCT)
617 : {
618 : const auto fieldStruct =
619 56 : std::static_pointer_cast<arrow::StructType>(type);
620 84 : const auto fieldX = fieldStruct->GetFieldByName("x");
621 84 : const auto fieldY = fieldStruct->GetFieldByName("y");
622 28 : if (fieldX && fieldY)
623 : {
624 56 : auto fieldRefX(fieldRefs);
625 28 : fieldRefX.emplace_back("x");
626 28 : auto fieldRefY(std::move(fieldRefs));
627 28 : fieldRefY.emplace_back("y");
628 252 : expression = cp::and_(
629 : {cp::less_equal(
630 56 : cp::field_ref(arrow::FieldRef(fieldRefX)),
631 56 : cp::literal(m_sFilterEnvelope.MaxX)),
632 : cp::less_equal(
633 56 : cp::field_ref(arrow::FieldRef(fieldRefY)),
634 56 : cp::literal(m_sFilterEnvelope.MaxY)),
635 : cp::greater_equal(
636 56 : cp::field_ref(arrow::FieldRef(fieldRefX)),
637 56 : cp::literal(m_sFilterEnvelope.MinX)),
638 : cp::greater_equal(
639 56 : cp::field_ref(arrow::FieldRef(fieldRefY)),
640 196 : cp::literal(m_sFilterEnvelope.MinY))});
641 : }
642 : }
643 : }
644 531 : else if (m_iGeomFieldFilter >= 0 &&
645 177 : m_iGeomFieldFilter <
646 354 : static_cast<int>(m_aeGeomEncoding.size()) &&
647 177 : m_aeGeomEncoding[m_iGeomFieldFilter] ==
648 : OGRArrowGeomEncoding::WKB)
649 : {
650 : const int iCol =
651 19 : m_anMapGeomFieldIndexToArrowColumn[m_iGeomFieldFilter];
652 19 : const auto &field = m_poSchema->fields()[iCol];
653 38 : if (field->type()->id() == arrow::Type::BINARY &&
654 19 : RegisterOGRWKBIntersectsIfNeeded())
655 : {
656 : #ifdef SUPPORTS_INDICES_IN_FIELD_REF
657 38 : auto oFieldRef = arrow::FieldRef(iCol);
658 : #else
659 : auto oFieldRef = arrow::FieldRef(field->name());
660 : #endif
661 38 : std::vector<GByte> abyFilterGeomWkb;
662 19 : abyFilterGeomWkb.resize(m_poFilterGeom->WkbSize());
663 19 : m_poFilterGeom->exportToWkb(wkbNDR, abyFilterGeomWkb.data(),
664 : wkbVariantIso);
665 : // Silence 'Using uninitialized value oFieldRef. Field oFieldRef.impl_._M_u is uninitialized when calling FieldRef.'
666 : // coverity[uninit_use_in_call]
667 57 : expression = cp::call("OGRWKBIntersects",
668 19 : {cp::field_ref(std::move(oFieldRef))},
669 76 : WKBGeometryOptions(abyFilterGeomWkb));
670 :
671 19 : if (expression.is_valid())
672 : {
673 19 : m_bBaseArrowIgnoreSpatialFilterRect = true;
674 19 : m_bBaseArrowIgnoreSpatialFilter = true;
675 19 : m_bSkipFilterGeometry = true;
676 : }
677 : }
678 : }
679 :
680 585 : if (expression.is_valid() && !m_bSkipFilterGeometry)
681 : {
682 408 : m_bBaseArrowIgnoreSpatialFilterRect = true;
683 :
684 : const bool bIsPoint =
685 408 : wkbFlatten(
686 : m_poFeatureDefn->GetGeomFieldDefn(m_iGeomFieldFilter)
687 408 : ->GetType()) == wkbPoint;
688 408 : m_bBaseArrowIgnoreSpatialFilter =
689 408 : m_bFilterIsEnvelope && bIsPoint;
690 :
691 408 : m_bSkipFilterGeometry =
692 629 : m_bFilterIsEnvelope &&
693 221 : (bIsPoint ||
694 221 : m_poFeatureDefn->GetGeomFieldDefn(m_iGeomFieldFilter)
695 221 : ->IsIgnored());
696 : }
697 : }
698 :
699 1041 : if (m_poAttrQuery &&
700 100 : CPLTestBool(CPLGetConfigOption(
701 : "OGR_PARQUET_OPTIMIZED_ATTRIBUTE_FILTER", "YES")))
702 : {
703 : const swq_expr_node *poNode =
704 100 : static_cast<swq_expr_node *>(m_poAttrQuery->GetSWQExpr());
705 100 : bool bFullyTranslated = true;
706 200 : auto expressionFilter = BuildArrowFilter(poNode, bFullyTranslated);
707 100 : if (expressionFilter.is_valid())
708 : {
709 95 : if (bFullyTranslated)
710 : {
711 93 : CPLDebugOnly("PARQUET",
712 : "Attribute filter fully translated to Arrow");
713 93 : m_asAttributeFilterConstraints.clear();
714 93 : m_bBaseArrowIgnoreAttributeFilter = true;
715 : }
716 :
717 95 : if (expression.is_valid())
718 : expression =
719 7 : cp::and_(expression, std::move(expressionFilter));
720 : else
721 88 : expression = std::move(expressionFilter);
722 : }
723 : }
724 :
725 941 : if (expression.is_valid())
726 : {
727 1030 : PARQUET_ASSIGN_OR_THROW(expression,
728 : expression.Bind(*m_poDataset->schema()));
729 :
730 : #ifdef DEBUG
731 1030 : arrow::dataset::FragmentIterator fragment_it;
732 1545 : PARQUET_ASSIGN_OR_THROW(fragment_it,
733 : m_poDataset->GetFragments(expression));
734 1026 : for (const auto &maybe_fragment : fragment_it)
735 : {
736 511 : auto fragment = maybe_fragment.ValueOrDie();
737 511 : CPLDebug("PARQUET", "Scanner planning to read '%s'",
738 1022 : fragment->ToString().c_str());
739 : }
740 : #endif
741 :
742 1030 : PARQUET_THROW_NOT_OK(scannerBuilder->Filter(expression));
743 : }
744 :
745 941 : if (m_bIgnoredFields)
746 : {
747 : #ifdef DEBUG
748 362 : std::string osFields;
749 1055 : for (const std::string &osField : m_aosProjectedFields)
750 : {
751 874 : if (!osFields.empty())
752 694 : osFields += ',';
753 874 : osFields += osField;
754 : }
755 181 : CPLDebug("PARQUET", "Projected fields: %s", osFields.c_str());
756 : #endif
757 362 : PARQUET_THROW_NOT_OK(scannerBuilder->Project(m_aosProjectedFields));
758 : }
759 :
760 941 : PARQUET_ASSIGN_OR_THROW(m_poScanner, scannerBuilder->Finish());
761 : }
762 0 : catch (const std::exception &e)
763 : {
764 0 : CPLError(CE_Failure, CPLE_AppDefined, "Arrow/Parquet exception: %s",
765 0 : e.what());
766 : }
767 941 : }
768 :
769 : /************************************************************************/
770 : /* BuildArrowFilter() */
771 : /************************************************************************/
772 :
773 : cp::Expression
774 357 : OGRParquetDatasetLayer::BuildArrowFilter(const swq_expr_node *poNode,
775 : bool &bFullyTranslated)
776 : {
777 357 : if (poNode->eNodeType == SNT_OPERATION && poNode->nOperation == SWQ_AND &&
778 10 : poNode->nSubExprCount == 2)
779 : {
780 10 : auto sLeft = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
781 : auto sRight =
782 10 : BuildArrowFilter(poNode->papoSubExpr[1], bFullyTranslated);
783 10 : if (sLeft.is_valid() && sRight.is_valid())
784 7 : return cp::and_(std::move(sLeft), std::move(sRight));
785 3 : else if (sLeft.is_valid())
786 1 : return sLeft;
787 2 : else if (sRight.is_valid())
788 2 : return sRight;
789 : }
790 :
791 347 : else if (poNode->eNodeType == SNT_OPERATION &&
792 131 : poNode->nOperation == SWQ_OR && poNode->nSubExprCount == 2)
793 : {
794 5 : auto sLeft = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
795 : auto sRight =
796 5 : BuildArrowFilter(poNode->papoSubExpr[1], bFullyTranslated);
797 5 : if (sLeft.is_valid() && sRight.is_valid())
798 5 : return cp::or_(std::move(sLeft), std::move(sRight));
799 : }
800 :
801 342 : else if (poNode->eNodeType == SNT_OPERATION &&
802 126 : poNode->nOperation == SWQ_NOT && poNode->nSubExprCount == 1)
803 : {
804 11 : auto expr = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
805 11 : if (expr.is_valid())
806 11 : return cp::not_(std::move(expr));
807 : }
808 :
809 331 : else if (poNode->eNodeType == SNT_COLUMN)
810 : {
811 202 : if (poNode->field_index >= 0 &&
812 101 : poNode->field_index < m_poFeatureDefn->GetFieldCount())
813 : {
814 184 : std::vector<arrow::FieldRef> fieldRefs;
815 92 : std::shared_ptr<arrow::Field> field;
816 186 : for (int idx : m_anMapFieldIndexToArrowColumn[poNode->field_index])
817 : {
818 94 : if (!field)
819 : {
820 92 : field = m_poSchema->fields()[idx];
821 : }
822 : else
823 : {
824 2 : CPLAssert(field->type()->id() == arrow::Type::STRUCT);
825 : const auto fieldStruct =
826 : std::static_pointer_cast<arrow::StructType>(
827 4 : field->type());
828 2 : field = fieldStruct->fields()[idx];
829 : }
830 94 : fieldRefs.emplace_back(field->name());
831 : }
832 :
833 276 : auto expr = cp::field_ref(arrow::FieldRef(std::move(fieldRefs)));
834 :
835 : // Comparing a boolean column to 0 or 1 fails without explicit cast
836 184 : if (m_poFeatureDefn->GetFieldDefn(poNode->field_index)
837 92 : ->GetSubType() == OFSTBoolean)
838 : {
839 48 : expr = cp::call("cast", {expr},
840 64 : cp::CastOptions::Safe(arrow::uint8()));
841 : }
842 92 : return expr;
843 : }
844 18 : else if (poNode->field_index ==
845 12 : m_poFeatureDefn->GetFieldCount() + SPF_FID &&
846 3 : m_iFIDArrowColumn >= 0)
847 : {
848 : #ifdef SUPPORTS_INDICES_IN_FIELD_REF
849 2 : return cp::field_ref(arrow::FieldRef(m_iFIDArrowColumn));
850 : #else
851 : return cp::field_ref(arrow::FieldRef(
852 : m_poSchema->fields()[m_iFIDArrowColumn]->name()));
853 : #endif
854 : }
855 : }
856 :
857 230 : else if (poNode->eNodeType == SNT_CONSTANT)
858 : {
859 115 : switch (poNode->field_type)
860 : {
861 77 : case SWQ_INTEGER:
862 154 : return cp::literal(static_cast<int32_t>(poNode->int_value));
863 :
864 13 : case SWQ_INTEGER64:
865 26 : return cp::literal(static_cast<int64_t>(poNode->int_value));
866 :
867 11 : case SWQ_FLOAT:
868 11 : return cp::literal(poNode->float_value);
869 :
870 10 : case SWQ_STRING:
871 10 : return cp::literal(poNode->string_value);
872 :
873 4 : case SWQ_TIMESTAMP:
874 : {
875 : OGRField sField;
876 4 : if (OGRParseDate(poNode->string_value, &sField, 0))
877 : {
878 : struct tm brokenDown;
879 4 : brokenDown.tm_year = sField.Date.Year - 1900;
880 4 : brokenDown.tm_mon = sField.Date.Month - 1;
881 4 : brokenDown.tm_mday = sField.Date.Day;
882 4 : brokenDown.tm_hour = sField.Date.Hour;
883 4 : brokenDown.tm_min = sField.Date.Minute;
884 4 : brokenDown.tm_sec = static_cast<int>(sField.Date.Second);
885 : int64_t nVal =
886 4 : CPLYMDHMSToUnixTime(&brokenDown) * 1000 +
887 4 : (static_cast<int>(sField.Date.Second * 1000 + 0.5) %
888 4 : 1000);
889 4 : if (sField.Date.TZFlag > OGR_TZFLAG_MIXED_TZ)
890 : {
891 : // Convert for sField.Date.TZFlag to UTC
892 2 : const int TZOffset =
893 2 : (sField.Date.TZFlag - OGR_TZFLAG_UTC) * 15;
894 2 : const int TZOffsetMS = TZOffset * 60 * 1000;
895 2 : nVal -= TZOffsetMS;
896 4 : return cp::literal(arrow::TimestampScalar(
897 2 : nVal, arrow::TimeUnit::MILLI, "UTC"));
898 : }
899 : else
900 : {
901 4 : return cp::literal(arrow::TimestampScalar(
902 2 : nVal, arrow::TimeUnit::MILLI));
903 : }
904 : }
905 0 : break;
906 : }
907 :
908 0 : default:
909 0 : break;
910 : }
911 : }
912 :
913 220 : else if (poNode->eNodeType == SNT_OPERATION && poNode->nSubExprCount == 2 &&
914 105 : IsComparisonOp(poNode->nOperation))
915 : {
916 101 : auto sLeft = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
917 : auto sRight =
918 101 : BuildArrowFilter(poNode->papoSubExpr[1], bFullyTranslated);
919 101 : if (sLeft.is_valid() && sRight.is_valid())
920 : {
921 99 : if (poNode->nOperation == SWQ_EQ)
922 67 : return cp::equal(std::move(sLeft), std::move(sRight));
923 32 : if (poNode->nOperation == SWQ_LT)
924 7 : return cp::less(std::move(sLeft), std::move(sRight));
925 25 : if (poNode->nOperation == SWQ_LE)
926 5 : return cp::less_equal(std::move(sLeft), std::move(sRight));
927 20 : if (poNode->nOperation == SWQ_GT)
928 5 : return cp::greater(std::move(sLeft), std::move(sRight));
929 15 : if (poNode->nOperation == SWQ_GE)
930 5 : return cp::greater_equal(std::move(sLeft), std::move(sRight));
931 10 : if (poNode->nOperation == SWQ_NE)
932 10 : return cp::not_equal(std::move(sLeft), std::move(sRight));
933 : }
934 : }
935 :
936 14 : else if (poNode->eNodeType == SNT_OPERATION && poNode->nSubExprCount == 2 &&
937 4 : (poNode->nOperation == SWQ_LIKE ||
938 1 : poNode->nOperation == SWQ_ILIKE) &&
939 4 : poNode->papoSubExpr[1]->eNodeType == SNT_CONSTANT &&
940 4 : poNode->papoSubExpr[1]->field_type == SWQ_STRING)
941 : {
942 4 : auto sLeft = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
943 4 : if (sLeft.is_valid())
944 : {
945 4 : if (cp::GetFunctionRegistry()
946 8 : ->GetFunction("match_like")
947 4 : .ValueOr(nullptr))
948 : {
949 : // match_like is only available is Arrow built against RE2.
950 : return cp::call(
951 4 : "match_like", {std::move(sLeft)},
952 8 : cp::MatchSubstringOptions(
953 4 : poNode->papoSubExpr[1]->string_value,
954 12 : /* ignore_case=*/poNode->nOperation == SWQ_ILIKE));
955 : }
956 0 : }
957 : }
958 :
959 10 : else if (poNode->eNodeType == SNT_OPERATION &&
960 10 : poNode->nOperation == SWQ_ISNULL && poNode->nSubExprCount == 1)
961 : {
962 10 : auto expr = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
963 10 : if (expr.is_valid())
964 4 : return cp::is_null(std::move(expr));
965 : }
966 :
967 25 : bFullyTranslated = false;
968 25 : return {};
969 : }
970 :
971 : /************************************************************************/
972 : /* ReadNextBatch() */
973 : /************************************************************************/
974 :
975 1558 : bool OGRParquetDatasetLayer::ReadNextBatch()
976 : {
977 1558 : if (m_bRebuildScanner)
978 908 : BuildScanner();
979 :
980 1558 : m_nIdxInBatch = 0;
981 :
982 1558 : if (m_poRecordBatchReader == nullptr)
983 : {
984 971 : if (!m_poScanner)
985 0 : return false;
986 971 : auto result = m_poScanner->ToRecordBatchReader();
987 971 : if (!result.ok())
988 : {
989 0 : CPLError(CE_Failure, CPLE_AppDefined,
990 : "ToRecordBatchReader() failed: %s",
991 0 : result.status().message().c_str());
992 0 : return false;
993 : }
994 971 : m_poRecordBatchReader = *result;
995 971 : if (m_poRecordBatchReader == nullptr)
996 0 : return false;
997 : }
998 :
999 3116 : std::shared_ptr<arrow::RecordBatch> poNextBatch;
1000 287 : do
1001 : {
1002 1845 : ++m_iRecordBatch;
1003 :
1004 1845 : poNextBatch.reset();
1005 1845 : auto status = m_poRecordBatchReader->ReadNext(&poNextBatch);
1006 1845 : if (!status.ok())
1007 : {
1008 0 : CPLError(CE_Failure, CPLE_AppDefined, "ReadNext() failed: %s",
1009 0 : status.message().c_str());
1010 0 : poNextBatch.reset();
1011 : }
1012 1845 : if (poNextBatch == nullptr)
1013 : {
1014 751 : m_poBatch.reset();
1015 751 : return false;
1016 : }
1017 1094 : } while (poNextBatch->num_rows() == 0);
1018 :
1019 : // CPLDebug("PARQUET", "Current batch has %d rows", int(poNextBatch->num_rows()));
1020 :
1021 807 : SetBatch(poNextBatch);
1022 :
1023 807 : return true;
1024 : }
1025 :
1026 : /************************************************************************/
1027 : /* GetNextFeature() */
1028 : /************************************************************************/
1029 :
1030 2783 : OGRFeature *OGRParquetDatasetLayer::GetNextFeature()
1031 : {
1032 : while (true)
1033 : {
1034 2783 : OGRFeature *poFeature = GetNextRawFeature();
1035 2783 : if (poFeature == nullptr)
1036 754 : return nullptr;
1037 :
1038 600 : if ((m_poFilterGeom == nullptr || m_bSkipFilterGeometry ||
1039 4575 : FilterGeometry(poFeature->GetGeometryRef())) &&
1040 2027 : (m_poAttrQuery == nullptr || m_bBaseArrowIgnoreAttributeFilter ||
1041 23 : m_poAttrQuery->Evaluate(poFeature)))
1042 : {
1043 2024 : return poFeature;
1044 : }
1045 : else
1046 5 : delete poFeature;
1047 5 : }
1048 : }
1049 :
1050 : /************************************************************************/
1051 : /* GetFeatureCount() */
1052 : /************************************************************************/
1053 :
1054 693 : GIntBig OGRParquetDatasetLayer::GetFeatureCount(int bForce)
1055 : {
1056 693 : if (m_poAttrQuery == nullptr && m_poFilterGeom == nullptr)
1057 : {
1058 39 : if (m_bRebuildScanner)
1059 33 : BuildScanner();
1060 39 : if (!m_poScanner)
1061 39 : return -1;
1062 39 : auto status = m_poScanner->CountRows();
1063 39 : if (status.ok())
1064 39 : return *status;
1065 : }
1066 654 : return OGRLayer::GetFeatureCount(bForce);
1067 : }
1068 :
1069 : /************************************************************************/
1070 : /* FastGetExtent() */
1071 : /************************************************************************/
1072 :
1073 609 : bool OGRParquetDatasetLayer::FastGetExtent(int iGeomField,
1074 : OGREnvelope *psExtent) const
1075 : {
1076 609 : const auto oIter = m_oMapExtents.find(iGeomField);
1077 609 : if (oIter != m_oMapExtents.end())
1078 : {
1079 24 : *psExtent = oIter->second;
1080 24 : return true;
1081 : }
1082 :
1083 585 : return false;
1084 : }
1085 :
1086 : /************************************************************************/
1087 : /* IGetExtent() */
1088 : /************************************************************************/
1089 :
1090 14 : OGRErr OGRParquetDatasetLayer::IGetExtent(int iGeomField, OGREnvelope *psExtent,
1091 : bool bForce)
1092 : {
1093 14 : if (FastGetExtent(iGeomField, psExtent))
1094 : {
1095 6 : return OGRERR_NONE;
1096 : }
1097 :
1098 : // bbox in general m_oMapGeometryColumns can not be trusted (at least at
1099 : // time of writing), so we have to iterate over each fragment.
1100 : const char *pszGeomFieldName =
1101 8 : m_poFeatureDefn->GetGeomFieldDefn(iGeomField)->GetNameRef();
1102 8 : auto oIter = m_oMapGeometryColumns.find(pszGeomFieldName);
1103 8 : if (oIter != m_oMapGeometryColumns.end())
1104 : {
1105 8 : auto statusFragments = m_poDataset->GetFragments();
1106 8 : if (statusFragments.ok())
1107 : {
1108 8 : *psExtent = OGREnvelope();
1109 8 : int nFragmentCount = 0;
1110 8 : int nBBoxFragmentCount = 0;
1111 18 : for (const auto &oFragmentStatus : *statusFragments)
1112 : {
1113 10 : if (oFragmentStatus.ok())
1114 : {
1115 : auto statusSchema =
1116 10 : (*oFragmentStatus)->ReadPhysicalSchema();
1117 10 : if (statusSchema.ok())
1118 : {
1119 10 : nFragmentCount++;
1120 10 : const auto &kv_metadata = (*statusSchema)->metadata();
1121 10 : if (kv_metadata && kv_metadata->Contains("geo"))
1122 : {
1123 20 : auto geo = kv_metadata->Get("geo");
1124 20 : CPLJSONDocument oDoc;
1125 10 : if (geo.ok() && oDoc.LoadMemory(*geo))
1126 : {
1127 20 : auto oRoot = oDoc.GetRoot();
1128 30 : auto oColumns = oRoot.GetObj("columns");
1129 30 : auto oCol = oColumns.GetObj(pszGeomFieldName);
1130 10 : OGREnvelope3D sFragmentExtent;
1131 20 : if (oCol.IsValid() &&
1132 10 : GetExtentFromMetadata(
1133 : oCol, &sFragmentExtent) == OGRERR_NONE)
1134 : {
1135 8 : nBBoxFragmentCount++;
1136 8 : psExtent->Merge(sFragmentExtent);
1137 : }
1138 : }
1139 : }
1140 10 : if (nFragmentCount != nBBoxFragmentCount)
1141 2 : break;
1142 : }
1143 : }
1144 : }
1145 8 : if (nFragmentCount == nBBoxFragmentCount)
1146 : {
1147 6 : m_oMapExtents[iGeomField] = *psExtent;
1148 6 : return OGRERR_NONE;
1149 : }
1150 : }
1151 : }
1152 :
1153 2 : return OGRParquetLayerBase::IGetExtent(iGeomField, psExtent, bForce);
1154 : }
1155 :
1156 : /************************************************************************/
1157 : /* ISetSpatialFilter() */
1158 : /************************************************************************/
1159 :
1160 674 : OGRErr OGRParquetDatasetLayer::ISetSpatialFilter(int iGeomField,
1161 : const OGRGeometry *poGeomIn)
1162 :
1163 : {
1164 : const OGRErr eErr =
1165 674 : OGRParquetLayerBase::ISetSpatialFilter(iGeomField, poGeomIn);
1166 674 : m_bRebuildScanner = true;
1167 :
1168 : // Full invalidation
1169 674 : InvalidateCachedBatches();
1170 674 : return eErr;
1171 : }
1172 :
1173 : /************************************************************************/
1174 : /* SetIgnoredFields() */
1175 : /************************************************************************/
1176 :
1177 131 : OGRErr OGRParquetDatasetLayer::SetIgnoredFields(CSLConstList papszFields)
1178 : {
1179 131 : m_bRebuildScanner = true;
1180 131 : m_aosProjectedFields.clear();
1181 131 : m_bIgnoredFields = false;
1182 131 : m_anMapFieldIndexToArrayIndex.clear();
1183 131 : m_anMapGeomFieldIndexToArrayIndex.clear();
1184 131 : m_nRequestedFIDColumn = -1;
1185 131 : OGRErr eErr = OGRParquetLayerBase::SetIgnoredFields(papszFields);
1186 131 : if (eErr == OGRERR_NONE)
1187 : {
1188 131 : m_bIgnoredFields = papszFields != nullptr && papszFields[0] != nullptr;
1189 131 : if (m_bIgnoredFields)
1190 : {
1191 96 : if (m_iFIDArrowColumn >= 0)
1192 : {
1193 1 : m_nRequestedFIDColumn =
1194 1 : static_cast<int>(m_aosProjectedFields.size());
1195 1 : m_aosProjectedFields.emplace_back(GetFIDColumn());
1196 : }
1197 :
1198 96 : const auto &fields = m_poSchema->fields();
1199 884 : for (int i = 0; i < m_poFeatureDefn->GetFieldCount(); ++i)
1200 : {
1201 : const auto &field =
1202 788 : fields[m_anMapFieldIndexToArrowColumn[i][0]];
1203 788 : const auto eArrowType = field->type()->id();
1204 788 : if (eArrowType == arrow::Type::STRUCT)
1205 : {
1206 : // For a struct, for the sake of simplicity in
1207 : // GetNextRawFeature(), as soon as one of the member if
1208 : // requested, request the struct field, so that the Arrow
1209 : // type doesn't change
1210 9 : bool bFoundNotIgnored = false;
1211 46 : for (int j = i; j < m_poFeatureDefn->GetFieldCount() &&
1212 46 : m_anMapFieldIndexToArrowColumn[i][0] ==
1213 23 : m_anMapFieldIndexToArrowColumn[j][0];
1214 : ++j)
1215 : {
1216 21 : if (!m_poFeatureDefn->GetFieldDefn(j)->IsIgnored())
1217 : {
1218 7 : bFoundNotIgnored = true;
1219 7 : break;
1220 : }
1221 : }
1222 9 : if (bFoundNotIgnored)
1223 : {
1224 : int j;
1225 98 : for (j = i; j < m_poFeatureDefn->GetFieldCount() &&
1226 98 : m_anMapFieldIndexToArrowColumn[i][0] ==
1227 49 : m_anMapFieldIndexToArrowColumn[j][0];
1228 : ++j)
1229 : {
1230 42 : if (!m_poFeatureDefn->GetFieldDefn(j)->IsIgnored())
1231 : {
1232 80 : m_anMapFieldIndexToArrayIndex.push_back(
1233 40 : static_cast<int>(
1234 40 : m_aosProjectedFields.size()));
1235 : }
1236 : else
1237 : {
1238 2 : m_anMapFieldIndexToArrayIndex.push_back(-1);
1239 : }
1240 : }
1241 7 : i = j - 1;
1242 :
1243 7 : m_aosProjectedFields.emplace_back(field->name());
1244 : }
1245 : else
1246 : {
1247 : int j;
1248 28 : for (j = i; j < m_poFeatureDefn->GetFieldCount() &&
1249 28 : m_anMapFieldIndexToArrowColumn[i][0] ==
1250 14 : m_anMapFieldIndexToArrowColumn[j][0];
1251 : ++j)
1252 : {
1253 12 : m_anMapFieldIndexToArrayIndex.push_back(-1);
1254 : }
1255 2 : i = j - 1;
1256 : }
1257 : }
1258 779 : else if (!m_poFeatureDefn->GetFieldDefn(i)->IsIgnored())
1259 : {
1260 1382 : m_anMapFieldIndexToArrayIndex.push_back(
1261 691 : static_cast<int>(m_aosProjectedFields.size()));
1262 691 : m_aosProjectedFields.emplace_back(field->name());
1263 : }
1264 : else
1265 : {
1266 88 : m_anMapFieldIndexToArrayIndex.push_back(-1);
1267 : }
1268 : }
1269 :
1270 191 : for (int i = 0; i < m_poFeatureDefn->GetGeomFieldCount(); ++i)
1271 : {
1272 : const auto &field =
1273 95 : fields[m_anMapGeomFieldIndexToArrowColumn[i]];
1274 95 : if (!m_poFeatureDefn->GetGeomFieldDefn(i)->IsIgnored())
1275 : {
1276 180 : m_anMapGeomFieldIndexToArrayIndex.push_back(
1277 90 : static_cast<int>(m_aosProjectedFields.size()));
1278 90 : m_aosProjectedFields.emplace_back(field->name());
1279 : }
1280 : else
1281 : {
1282 5 : m_anMapGeomFieldIndexToArrayIndex.push_back(-1);
1283 : }
1284 : }
1285 : }
1286 : }
1287 :
1288 131 : m_nExpectedBatchColumns =
1289 131 : m_bIgnoredFields ? static_cast<int>(m_aosProjectedFields.size()) : -1;
1290 :
1291 : // Full invalidation
1292 131 : InvalidateCachedBatches();
1293 :
1294 131 : return eErr;
1295 : }
1296 :
1297 : /************************************************************************/
1298 : /* TestCapability() */
1299 : /************************************************************************/
1300 :
1301 270 : int OGRParquetDatasetLayer::TestCapability(const char *pszCap) const
1302 : {
1303 270 : if (EQUAL(pszCap, OLCIgnoreFields))
1304 5 : return true;
1305 :
1306 265 : if (EQUAL(pszCap, OLCFastSpatialFilter))
1307 : {
1308 255 : if (m_iGeomFieldFilter >= 0 &&
1309 170 : m_iGeomFieldFilter < static_cast<int>(m_aeGeomEncoding.size()) &&
1310 85 : m_aeGeomEncoding[m_iGeomFieldFilter] ==
1311 : OGRArrowGeomEncoding::GEOARROW_STRUCT_POINT)
1312 : {
1313 12 : return true;
1314 : }
1315 : // fallback to base method
1316 : }
1317 :
1318 253 : return OGRParquetLayerBase::TestCapability(pszCap);
1319 : }
1320 :
1321 : /***********************************************************************/
1322 : /* SetAttributeFilter() */
1323 : /***********************************************************************/
1324 :
1325 196 : OGRErr OGRParquetDatasetLayer::SetAttributeFilter(const char *pszFilter)
1326 : {
1327 196 : m_bRebuildScanner = true;
1328 196 : return OGRParquetLayerBase::SetAttributeFilter(pszFilter);
1329 : }
|
