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 273 : OGRParquetDatasetLayer::OGRParquetDatasetLayer(
42 : OGRParquetDataset *poDS, const char *pszLayerName, bool bIsVSI,
43 : const std::shared_ptr<arrow::dataset::Dataset> &dataset,
44 273 : CSLConstList papszOpenOptions)
45 : : OGRParquetLayerBase(poDS, pszLayerName, papszOpenOptions),
46 273 : m_bIsVSI(bIsVSI), m_poDataset(dataset)
47 : {
48 273 : m_poSchema = m_poDataset->schema();
49 273 : EstablishFeatureDefn();
50 273 : CPLAssert(static_cast<int>(m_aeGeomEncoding.size()) ==
51 : m_poFeatureDefn->GetGeomFieldCount());
52 273 : }
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 254 : void OGRParquetDatasetLayer::ProcessGeometryColumnCovering(
63 : const std::shared_ptr<arrow::Field> &field,
64 : const CPLJSONObject &oJSONGeometryColumn)
65 : {
66 508 : std::string osBBOXColumn;
67 508 : std::string osXMin, osYMin, osXMax, osYMax;
68 254 : if (ParseGeometryColumnCovering(oJSONGeometryColumn, osBBOXColumn, osXMin,
69 : osYMin, osXMax, osYMax))
70 : {
71 90 : OGRArrowLayer::GeomColBBOX sDesc;
72 90 : sDesc.iArrowCol = m_poSchema->GetFieldIndex(osBBOXColumn);
73 180 : const auto fieldBBOX = m_poSchema->GetFieldByName(osBBOXColumn);
74 180 : if (sDesc.iArrowCol >= 0 && fieldBBOX &&
75 90 : fieldBBOX->type()->id() == arrow::Type::STRUCT)
76 : {
77 : const auto fieldBBOXStruct =
78 180 : std::static_pointer_cast<arrow::StructType>(fieldBBOX->type());
79 180 : const auto fieldXMin = fieldBBOXStruct->GetFieldByName(osXMin);
80 180 : const auto fieldYMin = fieldBBOXStruct->GetFieldByName(osYMin);
81 180 : const auto fieldXMax = fieldBBOXStruct->GetFieldByName(osXMax);
82 180 : const auto fieldYMax = fieldBBOXStruct->GetFieldByName(osYMax);
83 90 : const int nXMinIdx = fieldBBOXStruct->GetFieldIndex(osXMin);
84 90 : const int nYMinIdx = fieldBBOXStruct->GetFieldIndex(osYMin);
85 90 : const int nXMaxIdx = fieldBBOXStruct->GetFieldIndex(osXMax);
86 90 : const int nYMaxIdx = fieldBBOXStruct->GetFieldIndex(osYMax);
87 90 : if (nXMinIdx >= 0 && nYMinIdx >= 0 && nXMaxIdx >= 0 &&
88 180 : nYMaxIdx >= 0 && fieldXMin && fieldYMin && fieldXMax &&
89 180 : fieldYMax &&
90 90 : (fieldXMin->type()->id() == arrow::Type::FLOAT ||
91 0 : fieldXMin->type()->id() == arrow::Type::DOUBLE) &&
92 90 : fieldXMin->type()->id() == fieldYMin->type()->id() &&
93 270 : fieldXMin->type()->id() == fieldXMax->type()->id() &&
94 90 : fieldXMin->type()->id() == fieldYMax->type()->id())
95 : {
96 90 : CPLDebug("PARQUET",
97 : "Bounding box column '%s' detected for "
98 : "geometry column '%s'",
99 90 : osBBOXColumn.c_str(), field->name().c_str());
100 90 : sDesc.iArrowSubfieldXMin = nXMinIdx;
101 90 : sDesc.iArrowSubfieldYMin = nYMinIdx;
102 90 : sDesc.iArrowSubfieldXMax = nXMaxIdx;
103 90 : sDesc.iArrowSubfieldYMax = nYMaxIdx;
104 90 : sDesc.bIsFloat =
105 90 : (fieldXMin->type()->id() == arrow::Type::FLOAT);
106 :
107 : m_oMapGeomFieldIndexToGeomColBBOX
108 90 : [m_poFeatureDefn->GetGeomFieldCount() - 1] =
109 90 : std::move(sDesc);
110 : }
111 : }
112 : }
113 254 : }
114 :
115 : /************************************************************************/
116 : /* EstablishFeatureDefn() */
117 : /************************************************************************/
118 :
119 273 : void OGRParquetDatasetLayer::EstablishFeatureDefn()
120 : {
121 273 : const auto &kv_metadata = m_poSchema->metadata();
122 :
123 273 : LoadGeoMetadata(kv_metadata);
124 : const auto oMapFieldNameToGDALSchemaFieldDefn =
125 546 : LoadGDALSchema(kv_metadata.get());
126 :
127 273 : LoadGDALMetadata(kv_metadata.get());
128 :
129 : const bool bUseBBOX =
130 273 : 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 546 : std::set<std::string> oSetBBOXColumns;
135 273 : if (bUseBBOX)
136 : {
137 526 : for (const auto &iter : m_oMapGeometryColumns)
138 : {
139 508 : std::string osBBOXColumn;
140 508 : std::string osXMin, osYMin, osXMax, osYMax;
141 254 : if (ParseGeometryColumnCovering(iter.second, osBBOXColumn, osXMin,
142 : osYMin, osXMax, osYMax))
143 : {
144 89 : oSetBBOXColumns.insert(std::move(osBBOXColumn));
145 : }
146 : }
147 : }
148 :
149 273 : 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 801 : if (m_oMapGeometryColumns.find("geometry") != m_oMapGeometryColumns.end() &&
154 508 : bUseBBOX &&
155 1239 : !m_oMapGeometryColumns["geometry"].GetObj("covering").IsValid() &&
156 438 : 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 7178 : for (int i = 0; i < m_poSchema->num_fields(); ++i)
224 : {
225 6905 : const auto &field = fields[i];
226 :
227 6905 : if (!m_osFIDColumn.empty() && field->name() == m_osFIDColumn)
228 : {
229 2 : m_iFIDArrowColumn = i;
230 2 : continue;
231 : }
232 :
233 6903 : if (oSetBBOXColumns.find(field->name()) != oSetBBOXColumns.end())
234 : {
235 90 : m_oSetBBoxArrowColumns.insert(i);
236 90 : continue;
237 : }
238 :
239 6813 : const bool bGeometryField = DealWithGeometryColumn(
240 79 : i, field, []() { return wkbUnknown; }, nullptr, nullptr, -1);
241 6813 : if (bGeometryField)
242 : {
243 255 : const auto oIter = m_oMapGeometryColumns.find(field->name());
244 255 : if (bUseBBOX && oIter != m_oMapGeometryColumns.end())
245 : {
246 254 : ProcessGeometryColumnCovering(field, oIter->second);
247 : }
248 : }
249 : else
250 : {
251 6558 : CreateFieldFromSchema(field, {i},
252 : oMapFieldNameToGDALSchemaFieldDefn);
253 : }
254 : }
255 :
256 273 : CPLAssert(static_cast<int>(m_anMapFieldIndexToArrowColumn.size()) ==
257 : m_poFeatureDefn->GetFieldCount());
258 273 : CPLAssert(static_cast<int>(m_anMapGeomFieldIndexToArrowColumn.size()) ==
259 : m_poFeatureDefn->GetGeomFieldCount());
260 273 : }
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 19 : explicit OptionsWrapper(OptionsType optionsIn)
357 19 : : options(std::move(optionsIn))
358 : {
359 19 : }
360 :
361 : static arrow::Result<std::unique_ptr<cp::KernelState>>
362 19 : Init(cp::KernelContext *, const cp::KernelInitArgs &args)
363 : {
364 19 : auto options = cpl::down_cast<const OptionsType *>(args.options);
365 19 : CPLAssert(options);
366 19 : 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 694 : void OGRParquetDatasetLayer::BuildScanner()
487 : {
488 694 : m_bRebuildScanner = false;
489 694 : m_bSkipFilterGeometry = false;
490 694 : m_bBaseArrowIgnoreSpatialFilterRect = false;
491 694 : m_bBaseArrowIgnoreSpatialFilter = false;
492 694 : m_bBaseArrowIgnoreAttributeFilter = false;
493 :
494 : try
495 : {
496 694 : std::shared_ptr<arrow::dataset::ScannerBuilder> scannerBuilder;
497 1388 : PARQUET_ASSIGN_OR_THROW(scannerBuilder, m_poDataset->NewScan());
498 694 : 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 694 : if (m_bIsVSI)
507 : {
508 254 : const int nFragmentReadAhead = atoi(
509 : CPLGetConfigOption("OGR_PARQUET_FRAGMENT_READ_AHEAD", "2"));
510 508 : PARQUET_THROW_NOT_OK(
511 : scannerBuilder->FragmentReadahead(nFragmentReadAhead));
512 : }
513 :
514 : const char *pszBatchSize =
515 694 : CPLGetConfigOption("OGR_PARQUET_BATCH_SIZE", nullptr);
516 694 : if (pszBatchSize)
517 : {
518 0 : PARQUET_THROW_NOT_OK(
519 : scannerBuilder->BatchSize(CPLAtoGIntBig(pszBatchSize)));
520 : }
521 :
522 694 : const int nNumCPUs = GetNumCPUs();
523 : const char *pszUseThreads =
524 694 : CPLGetConfigOption("OGR_PARQUET_USE_THREADS", nullptr);
525 694 : if (!pszUseThreads && nNumCPUs > 1)
526 : {
527 694 : pszUseThreads = "YES";
528 : }
529 694 : if (pszUseThreads && CPLTestBool(pszUseThreads))
530 : {
531 1388 : PARQUET_THROW_NOT_OK(scannerBuilder->UseThreads(true));
532 : }
533 :
534 : #if PARQUET_VERSION_MAJOR >= 10
535 : const char *pszBatchReadAhead =
536 694 : CPLGetConfigOption("OGR_PARQUET_BATCH_READ_AHEAD", nullptr);
537 694 : if (pszBatchReadAhead)
538 : {
539 0 : PARQUET_THROW_NOT_OK(
540 : scannerBuilder->BatchReadahead(atoi(pszBatchReadAhead)));
541 : }
542 : #endif
543 :
544 694 : cp::Expression expression;
545 1093 : if (m_poFilterGeom && !m_poFilterGeom->IsEmpty() &&
546 399 : CPLTestBool(CPLGetConfigOption(
547 : "OGR_PARQUET_OPTIMIZED_SPATIAL_FILTER", "YES")))
548 : {
549 : const auto oIter =
550 399 : m_oMapGeomFieldIndexToGeomColBBOX.find(m_iGeomFieldFilter);
551 399 : 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 194 : const auto &oBBOXDef = oIter->second;
557 1746 : expression = cp::and_(
558 : {cp::less_equal(
559 388 : cp::field_ref(arrow::FieldRef(
560 194 : oBBOXDef.iArrowCol, oBBOXDef.iArrowSubfieldXMin)),
561 388 : cp::literal(m_sFilterEnvelope.MaxX)),
562 : cp::less_equal(
563 388 : cp::field_ref(arrow::FieldRef(
564 194 : oBBOXDef.iArrowCol, oBBOXDef.iArrowSubfieldYMin)),
565 388 : cp::literal(m_sFilterEnvelope.MaxY)),
566 : cp::greater_equal(
567 388 : cp::field_ref(arrow::FieldRef(
568 194 : oBBOXDef.iArrowCol, oBBOXDef.iArrowSubfieldXMax)),
569 388 : cp::literal(m_sFilterEnvelope.MinX)),
570 : cp::greater_equal(
571 388 : cp::field_ref(arrow::FieldRef(
572 194 : oBBOXDef.iArrowCol, oBBOXDef.iArrowSubfieldYMax)),
573 1358 : 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 399 : if (expression.is_valid() && !m_bSkipFilterGeometry)
681 : {
682 222 : m_bBaseArrowIgnoreSpatialFilterRect = true;
683 :
684 : const bool bIsPoint =
685 222 : wkbFlatten(
686 : m_poFeatureDefn->GetGeomFieldDefn(m_iGeomFieldFilter)
687 222 : ->GetType()) == wkbPoint;
688 222 : m_bBaseArrowIgnoreSpatialFilter =
689 222 : m_bFilterIsEnvelope && bIsPoint;
690 :
691 222 : m_bSkipFilterGeometry =
692 333 : m_bFilterIsEnvelope &&
693 111 : (bIsPoint ||
694 111 : m_poFeatureDefn->GetGeomFieldDefn(m_iGeomFieldFilter)
695 111 : ->IsIgnored());
696 : }
697 : }
698 :
699 789 : if (m_poAttrQuery &&
700 95 : CPLTestBool(CPLGetConfigOption(
701 : "OGR_PARQUET_OPTIMIZED_ATTRIBUTE_FILTER", "YES")))
702 : {
703 : const swq_expr_node *poNode =
704 95 : static_cast<swq_expr_node *>(m_poAttrQuery->GetSWQExpr());
705 95 : bool bFullyTranslated = true;
706 190 : auto expressionFilter = BuildArrowFilter(poNode, bFullyTranslated);
707 95 : if (expressionFilter.is_valid())
708 : {
709 90 : if (bFullyTranslated)
710 : {
711 88 : CPLDebugOnly("PARQUET",
712 : "Attribute filter fully translated to Arrow");
713 88 : m_asAttributeFilterConstraints.clear();
714 88 : m_bBaseArrowIgnoreAttributeFilter = true;
715 : }
716 :
717 90 : if (expression.is_valid())
718 : expression =
719 7 : cp::and_(expression, std::move(expressionFilter));
720 : else
721 83 : expression = std::move(expressionFilter);
722 : }
723 : }
724 :
725 694 : if (expression.is_valid())
726 : {
727 648 : PARQUET_THROW_NOT_OK(scannerBuilder->Filter(expression));
728 : }
729 :
730 694 : if (m_bIgnoredFields)
731 : {
732 : #ifdef DEBUG
733 250 : std::string osFields;
734 943 : for (const std::string &osField : m_aosProjectedFields)
735 : {
736 818 : if (!osFields.empty())
737 694 : osFields += ',';
738 818 : osFields += osField;
739 : }
740 125 : CPLDebug("PARQUET", "Projected fields: %s", osFields.c_str());
741 : #endif
742 250 : PARQUET_THROW_NOT_OK(scannerBuilder->Project(m_aosProjectedFields));
743 : }
744 :
745 694 : PARQUET_ASSIGN_OR_THROW(m_poScanner, scannerBuilder->Finish());
746 : }
747 0 : catch (const std::exception &e)
748 : {
749 0 : CPLError(CE_Failure, CPLE_AppDefined, "Arrow/Parquet exception: %s",
750 0 : e.what());
751 : }
752 694 : }
753 :
754 : /************************************************************************/
755 : /* BuildArrowFilter() */
756 : /************************************************************************/
757 :
758 : cp::Expression
759 338 : OGRParquetDatasetLayer::BuildArrowFilter(const swq_expr_node *poNode,
760 : bool &bFullyTranslated)
761 : {
762 338 : if (poNode->eNodeType == SNT_OPERATION && poNode->nOperation == SWQ_AND &&
763 9 : poNode->nSubExprCount == 2)
764 : {
765 9 : auto sLeft = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
766 : auto sRight =
767 9 : BuildArrowFilter(poNode->papoSubExpr[1], bFullyTranslated);
768 9 : if (sLeft.is_valid() && sRight.is_valid())
769 6 : return cp::and_(std::move(sLeft), std::move(sRight));
770 3 : else if (sLeft.is_valid())
771 1 : return sLeft;
772 2 : else if (sRight.is_valid())
773 2 : return sRight;
774 : }
775 :
776 329 : else if (poNode->eNodeType == SNT_OPERATION &&
777 125 : poNode->nOperation == SWQ_OR && poNode->nSubExprCount == 2)
778 : {
779 5 : auto sLeft = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
780 : auto sRight =
781 5 : BuildArrowFilter(poNode->papoSubExpr[1], bFullyTranslated);
782 5 : if (sLeft.is_valid() && sRight.is_valid())
783 5 : return cp::or_(std::move(sLeft), std::move(sRight));
784 : }
785 :
786 324 : else if (poNode->eNodeType == SNT_OPERATION &&
787 120 : poNode->nOperation == SWQ_NOT && poNode->nSubExprCount == 1)
788 : {
789 11 : auto expr = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
790 11 : if (expr.is_valid())
791 11 : return cp::not_(std::move(expr));
792 : }
793 :
794 313 : else if (poNode->eNodeType == SNT_COLUMN)
795 : {
796 190 : if (poNode->field_index >= 0 &&
797 95 : poNode->field_index < m_poFeatureDefn->GetFieldCount())
798 : {
799 172 : std::vector<arrow::FieldRef> fieldRefs;
800 : #ifdef SUPPORTS_INDICES_IN_FIELD_REF
801 174 : for (int idx : m_anMapFieldIndexToArrowColumn[poNode->field_index])
802 88 : fieldRefs.emplace_back(idx);
803 : #else
804 : std::shared_ptr<arrow::Field> field;
805 : for (int idx : m_anMapFieldIndexToArrowColumn[poNode->field_index])
806 : {
807 : if (!field)
808 : {
809 : field = m_poSchema->fields()[idx];
810 : }
811 : else
812 : {
813 : CPLAssert(field->type()->id() == arrow::Type::STRUCT);
814 : const auto fieldStruct =
815 : std::static_pointer_cast<arrow::StructType>(
816 : field->type());
817 : field = fieldStruct->fields()[idx];
818 : }
819 : fieldRefs.emplace_back(field->name());
820 : }
821 : #endif
822 258 : auto expr = cp::field_ref(arrow::FieldRef(std::move(fieldRefs)));
823 :
824 : // Comparing a boolean column to 0 or 1 fails without explicit cast
825 172 : if (m_poFeatureDefn->GetFieldDefn(poNode->field_index)
826 86 : ->GetSubType() == OFSTBoolean)
827 : {
828 48 : expr = cp::call("cast", {expr},
829 64 : cp::CastOptions::Safe(arrow::uint8()));
830 : }
831 86 : return expr;
832 : }
833 18 : else if (poNode->field_index ==
834 12 : m_poFeatureDefn->GetFieldCount() + SPF_FID &&
835 3 : m_iFIDArrowColumn >= 0)
836 : {
837 : #ifdef SUPPORTS_INDICES_IN_FIELD_REF
838 2 : return cp::field_ref(arrow::FieldRef(m_iFIDArrowColumn));
839 : #else
840 : return cp::field_ref(arrow::FieldRef(
841 : m_poSchema->fields()[m_iFIDArrowColumn]->name()));
842 : #endif
843 : }
844 : }
845 :
846 218 : else if (poNode->eNodeType == SNT_CONSTANT)
847 : {
848 109 : switch (poNode->field_type)
849 : {
850 88 : case SWQ_INTEGER:
851 : case SWQ_INTEGER64:
852 176 : return cp::literal(static_cast<int64_t>(poNode->int_value));
853 :
854 11 : case SWQ_FLOAT:
855 11 : return cp::literal(poNode->float_value);
856 :
857 6 : case SWQ_STRING:
858 6 : return cp::literal(poNode->string_value);
859 :
860 4 : case SWQ_TIMESTAMP:
861 : {
862 : OGRField sField;
863 4 : if (OGRParseDate(poNode->string_value, &sField, 0))
864 : {
865 : struct tm brokenDown;
866 4 : brokenDown.tm_year = sField.Date.Year - 1900;
867 4 : brokenDown.tm_mon = sField.Date.Month - 1;
868 4 : brokenDown.tm_mday = sField.Date.Day;
869 4 : brokenDown.tm_hour = sField.Date.Hour;
870 4 : brokenDown.tm_min = sField.Date.Minute;
871 4 : brokenDown.tm_sec = static_cast<int>(sField.Date.Second);
872 : int64_t nVal =
873 4 : CPLYMDHMSToUnixTime(&brokenDown) * 1000 +
874 4 : (static_cast<int>(sField.Date.Second * 1000 + 0.5) %
875 4 : 1000);
876 4 : if (sField.Date.TZFlag > OGR_TZFLAG_MIXED_TZ)
877 : {
878 : // Convert for sField.Date.TZFlag to UTC
879 2 : const int TZOffset =
880 2 : (sField.Date.TZFlag - OGR_TZFLAG_UTC) * 15;
881 2 : const int TZOffsetMS = TZOffset * 60 * 1000;
882 2 : nVal -= TZOffsetMS;
883 4 : return cp::literal(arrow::TimestampScalar(
884 2 : nVal, arrow::TimeUnit::MILLI, "UTC"));
885 : }
886 : else
887 : {
888 4 : return cp::literal(arrow::TimestampScalar(
889 2 : nVal, arrow::TimeUnit::MILLI));
890 : }
891 : }
892 0 : break;
893 : }
894 :
895 0 : default:
896 0 : break;
897 : }
898 : }
899 :
900 208 : else if (poNode->eNodeType == SNT_OPERATION && poNode->nSubExprCount == 2 &&
901 99 : IsComparisonOp(poNode->nOperation))
902 : {
903 95 : auto sLeft = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
904 : auto sRight =
905 95 : BuildArrowFilter(poNode->papoSubExpr[1], bFullyTranslated);
906 95 : if (sLeft.is_valid() && sRight.is_valid())
907 : {
908 93 : if (poNode->nOperation == SWQ_EQ)
909 61 : return cp::equal(std::move(sLeft), std::move(sRight));
910 32 : if (poNode->nOperation == SWQ_LT)
911 7 : return cp::less(std::move(sLeft), std::move(sRight));
912 25 : if (poNode->nOperation == SWQ_LE)
913 5 : return cp::less_equal(std::move(sLeft), std::move(sRight));
914 20 : if (poNode->nOperation == SWQ_GT)
915 5 : return cp::greater(std::move(sLeft), std::move(sRight));
916 15 : if (poNode->nOperation == SWQ_GE)
917 5 : return cp::greater_equal(std::move(sLeft), std::move(sRight));
918 10 : if (poNode->nOperation == SWQ_NE)
919 10 : return cp::not_equal(std::move(sLeft), std::move(sRight));
920 : }
921 : }
922 :
923 14 : else if (poNode->eNodeType == SNT_OPERATION && poNode->nSubExprCount == 2 &&
924 4 : (poNode->nOperation == SWQ_LIKE ||
925 1 : poNode->nOperation == SWQ_ILIKE) &&
926 4 : poNode->papoSubExpr[1]->eNodeType == SNT_CONSTANT &&
927 4 : poNode->papoSubExpr[1]->field_type == SWQ_STRING)
928 : {
929 4 : auto sLeft = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
930 4 : if (sLeft.is_valid())
931 : {
932 4 : if (cp::GetFunctionRegistry()
933 8 : ->GetFunction("match_like")
934 4 : .ValueOr(nullptr))
935 : {
936 : // match_like is only available is Arrow built against RE2.
937 : return cp::call(
938 4 : "match_like", {std::move(sLeft)},
939 8 : cp::MatchSubstringOptions(
940 4 : poNode->papoSubExpr[1]->string_value,
941 12 : /* ignore_case=*/poNode->nOperation == SWQ_ILIKE));
942 : }
943 0 : }
944 : }
945 :
946 10 : else if (poNode->eNodeType == SNT_OPERATION &&
947 10 : poNode->nOperation == SWQ_ISNULL && poNode->nSubExprCount == 1)
948 : {
949 10 : auto expr = BuildArrowFilter(poNode->papoSubExpr[0], bFullyTranslated);
950 10 : if (expr.is_valid())
951 4 : return cp::is_null(std::move(expr));
952 : }
953 :
954 25 : bFullyTranslated = false;
955 25 : return {};
956 : }
957 :
958 : /************************************************************************/
959 : /* ReadNextBatch() */
960 : /************************************************************************/
961 :
962 1224 : bool OGRParquetDatasetLayer::ReadNextBatch()
963 : {
964 1224 : if (m_bRebuildScanner)
965 661 : BuildScanner();
966 :
967 1224 : m_nIdxInBatch = 0;
968 :
969 1224 : if (m_poRecordBatchReader == nullptr)
970 : {
971 721 : if (!m_poScanner)
972 0 : return false;
973 721 : auto result = m_poScanner->ToRecordBatchReader();
974 721 : if (!result.ok())
975 : {
976 0 : CPLError(CE_Failure, CPLE_AppDefined,
977 : "ToRecordBatchReader() failed: %s",
978 0 : result.status().message().c_str());
979 0 : return false;
980 : }
981 721 : m_poRecordBatchReader = *result;
982 721 : if (m_poRecordBatchReader == nullptr)
983 0 : return false;
984 : }
985 :
986 2448 : std::shared_ptr<arrow::RecordBatch> poNextBatch;
987 175 : do
988 : {
989 1399 : ++m_iRecordBatch;
990 :
991 1399 : poNextBatch.reset();
992 1399 : auto status = m_poRecordBatchReader->ReadNext(&poNextBatch);
993 1399 : if (!status.ok())
994 : {
995 0 : CPLError(CE_Failure, CPLE_AppDefined, "ReadNext() failed: %s",
996 0 : status.message().c_str());
997 0 : poNextBatch.reset();
998 : }
999 1399 : if (poNextBatch == nullptr)
1000 : {
1001 557 : m_poBatch.reset();
1002 557 : return false;
1003 : }
1004 842 : } while (poNextBatch->num_rows() == 0);
1005 :
1006 : // CPLDebug("PARQUET", "Current batch has %d rows", int(poNextBatch->num_rows()));
1007 :
1008 667 : SetBatch(poNextBatch);
1009 :
1010 667 : return true;
1011 : }
1012 :
1013 : /************************************************************************/
1014 : /* GetNextFeature() */
1015 : /************************************************************************/
1016 :
1017 2120 : OGRFeature *OGRParquetDatasetLayer::GetNextFeature()
1018 : {
1019 : while (true)
1020 : {
1021 2120 : OGRFeature *poFeature = GetNextRawFeature();
1022 2120 : if (poFeature == nullptr)
1023 554 : return nullptr;
1024 :
1025 420 : if ((m_poFilterGeom == nullptr || m_bSkipFilterGeometry ||
1026 3477 : FilterGeometry(poFeature->GetGeometryRef())) &&
1027 1564 : (m_poAttrQuery == nullptr || m_bBaseArrowIgnoreAttributeFilter ||
1028 23 : m_poAttrQuery->Evaluate(poFeature)))
1029 : {
1030 1561 : return poFeature;
1031 : }
1032 : else
1033 5 : delete poFeature;
1034 5 : }
1035 : }
1036 :
1037 : /************************************************************************/
1038 : /* GetFeatureCount() */
1039 : /************************************************************************/
1040 :
1041 502 : GIntBig OGRParquetDatasetLayer::GetFeatureCount(int bForce)
1042 : {
1043 502 : if (m_poAttrQuery == nullptr && m_poFilterGeom == nullptr)
1044 : {
1045 39 : if (m_bRebuildScanner)
1046 33 : BuildScanner();
1047 39 : if (!m_poScanner)
1048 39 : return -1;
1049 39 : auto status = m_poScanner->CountRows();
1050 39 : if (status.ok())
1051 39 : return *status;
1052 : }
1053 463 : return OGRLayer::GetFeatureCount(bForce);
1054 : }
1055 :
1056 : /************************************************************************/
1057 : /* FastGetExtent() */
1058 : /************************************************************************/
1059 :
1060 423 : bool OGRParquetDatasetLayer::FastGetExtent(int iGeomField,
1061 : OGREnvelope *psExtent) const
1062 : {
1063 423 : const auto oIter = m_oMapExtents.find(iGeomField);
1064 423 : if (oIter != m_oMapExtents.end())
1065 : {
1066 24 : *psExtent = oIter->second;
1067 24 : return true;
1068 : }
1069 :
1070 399 : return false;
1071 : }
1072 :
1073 : /************************************************************************/
1074 : /* IGetExtent() */
1075 : /************************************************************************/
1076 :
1077 14 : OGRErr OGRParquetDatasetLayer::IGetExtent(int iGeomField, OGREnvelope *psExtent,
1078 : bool bForce)
1079 : {
1080 14 : if (FastGetExtent(iGeomField, psExtent))
1081 : {
1082 6 : return OGRERR_NONE;
1083 : }
1084 :
1085 : // bbox in general m_oMapGeometryColumns can not be trusted (at least at
1086 : // time of writing), so we have to iterate over each fragment.
1087 : const char *pszGeomFieldName =
1088 8 : m_poFeatureDefn->GetGeomFieldDefn(iGeomField)->GetNameRef();
1089 8 : auto oIter = m_oMapGeometryColumns.find(pszGeomFieldName);
1090 8 : if (oIter != m_oMapGeometryColumns.end())
1091 : {
1092 8 : auto statusFragments = m_poDataset->GetFragments();
1093 8 : if (statusFragments.ok())
1094 : {
1095 8 : *psExtent = OGREnvelope();
1096 8 : int nFragmentCount = 0;
1097 8 : int nBBoxFragmentCount = 0;
1098 18 : for (const auto &oFragmentStatus : *statusFragments)
1099 : {
1100 10 : if (oFragmentStatus.ok())
1101 : {
1102 : auto statusSchema =
1103 10 : (*oFragmentStatus)->ReadPhysicalSchema();
1104 10 : if (statusSchema.ok())
1105 : {
1106 10 : nFragmentCount++;
1107 10 : const auto &kv_metadata = (*statusSchema)->metadata();
1108 10 : if (kv_metadata && kv_metadata->Contains("geo"))
1109 : {
1110 20 : auto geo = kv_metadata->Get("geo");
1111 20 : CPLJSONDocument oDoc;
1112 10 : if (geo.ok() && oDoc.LoadMemory(*geo))
1113 : {
1114 20 : auto oRoot = oDoc.GetRoot();
1115 30 : auto oColumns = oRoot.GetObj("columns");
1116 30 : auto oCol = oColumns.GetObj(pszGeomFieldName);
1117 10 : OGREnvelope3D sFragmentExtent;
1118 20 : if (oCol.IsValid() &&
1119 10 : GetExtentFromMetadata(
1120 : oCol, &sFragmentExtent) == OGRERR_NONE)
1121 : {
1122 8 : nBBoxFragmentCount++;
1123 8 : psExtent->Merge(sFragmentExtent);
1124 : }
1125 : }
1126 : }
1127 10 : if (nFragmentCount != nBBoxFragmentCount)
1128 2 : break;
1129 : }
1130 : }
1131 : }
1132 8 : if (nFragmentCount == nBBoxFragmentCount)
1133 : {
1134 6 : m_oMapExtents[iGeomField] = *psExtent;
1135 6 : return OGRERR_NONE;
1136 : }
1137 : }
1138 : }
1139 :
1140 2 : return OGRParquetLayerBase::IGetExtent(iGeomField, psExtent, bForce);
1141 : }
1142 :
1143 : /************************************************************************/
1144 : /* ISetSpatialFilter() */
1145 : /************************************************************************/
1146 :
1147 488 : OGRErr OGRParquetDatasetLayer::ISetSpatialFilter(int iGeomField,
1148 : const OGRGeometry *poGeomIn)
1149 :
1150 : {
1151 : const OGRErr eErr =
1152 488 : OGRParquetLayerBase::ISetSpatialFilter(iGeomField, poGeomIn);
1153 488 : m_bRebuildScanner = true;
1154 :
1155 : // Full invalidation
1156 488 : InvalidateCachedBatches();
1157 488 : return eErr;
1158 : }
1159 :
1160 : /************************************************************************/
1161 : /* SetIgnoredFields() */
1162 : /************************************************************************/
1163 :
1164 103 : OGRErr OGRParquetDatasetLayer::SetIgnoredFields(CSLConstList papszFields)
1165 : {
1166 103 : m_bRebuildScanner = true;
1167 103 : m_aosProjectedFields.clear();
1168 103 : m_bIgnoredFields = false;
1169 103 : m_anMapFieldIndexToArrayIndex.clear();
1170 103 : m_anMapGeomFieldIndexToArrayIndex.clear();
1171 103 : m_nRequestedFIDColumn = -1;
1172 103 : OGRErr eErr = OGRParquetLayerBase::SetIgnoredFields(papszFields);
1173 103 : if (eErr == OGRERR_NONE)
1174 : {
1175 103 : m_bIgnoredFields = papszFields != nullptr && papszFields[0] != nullptr;
1176 103 : if (m_bIgnoredFields)
1177 : {
1178 68 : if (m_iFIDArrowColumn >= 0)
1179 : {
1180 1 : m_nRequestedFIDColumn =
1181 1 : static_cast<int>(m_aosProjectedFields.size());
1182 1 : m_aosProjectedFields.emplace_back(GetFIDColumn());
1183 : }
1184 :
1185 68 : const auto &fields = m_poSchema->fields();
1186 828 : for (int i = 0; i < m_poFeatureDefn->GetFieldCount(); ++i)
1187 : {
1188 : const auto &field =
1189 760 : fields[m_anMapFieldIndexToArrowColumn[i][0]];
1190 760 : const auto eArrowType = field->type()->id();
1191 760 : if (eArrowType == arrow::Type::STRUCT)
1192 : {
1193 : // For a struct, for the sake of simplicity in
1194 : // GetNextRawFeature(), as soon as one of the member if
1195 : // requested, request the struct field, so that the Arrow
1196 : // type doesn't change
1197 9 : bool bFoundNotIgnored = false;
1198 46 : for (int j = i; j < m_poFeatureDefn->GetFieldCount() &&
1199 46 : m_anMapFieldIndexToArrowColumn[i][0] ==
1200 23 : m_anMapFieldIndexToArrowColumn[j][0];
1201 : ++j)
1202 : {
1203 21 : if (!m_poFeatureDefn->GetFieldDefn(j)->IsIgnored())
1204 : {
1205 7 : bFoundNotIgnored = true;
1206 7 : break;
1207 : }
1208 : }
1209 9 : if (bFoundNotIgnored)
1210 : {
1211 : int j;
1212 98 : for (j = i; j < m_poFeatureDefn->GetFieldCount() &&
1213 98 : m_anMapFieldIndexToArrowColumn[i][0] ==
1214 49 : m_anMapFieldIndexToArrowColumn[j][0];
1215 : ++j)
1216 : {
1217 42 : if (!m_poFeatureDefn->GetFieldDefn(j)->IsIgnored())
1218 : {
1219 80 : m_anMapFieldIndexToArrayIndex.push_back(
1220 40 : static_cast<int>(
1221 40 : m_aosProjectedFields.size()));
1222 : }
1223 : else
1224 : {
1225 2 : m_anMapFieldIndexToArrayIndex.push_back(-1);
1226 : }
1227 : }
1228 7 : i = j - 1;
1229 :
1230 7 : m_aosProjectedFields.emplace_back(field->name());
1231 : }
1232 : else
1233 : {
1234 : int j;
1235 28 : for (j = i; j < m_poFeatureDefn->GetFieldCount() &&
1236 28 : m_anMapFieldIndexToArrowColumn[i][0] ==
1237 14 : m_anMapFieldIndexToArrowColumn[j][0];
1238 : ++j)
1239 : {
1240 12 : m_anMapFieldIndexToArrayIndex.push_back(-1);
1241 : }
1242 2 : i = j - 1;
1243 : }
1244 : }
1245 751 : else if (!m_poFeatureDefn->GetFieldDefn(i)->IsIgnored())
1246 : {
1247 1382 : m_anMapFieldIndexToArrayIndex.push_back(
1248 691 : static_cast<int>(m_aosProjectedFields.size()));
1249 691 : m_aosProjectedFields.emplace_back(field->name());
1250 : }
1251 : else
1252 : {
1253 60 : m_anMapFieldIndexToArrayIndex.push_back(-1);
1254 : }
1255 : }
1256 :
1257 135 : for (int i = 0; i < m_poFeatureDefn->GetGeomFieldCount(); ++i)
1258 : {
1259 : const auto &field =
1260 67 : fields[m_anMapGeomFieldIndexToArrowColumn[i]];
1261 67 : if (!m_poFeatureDefn->GetGeomFieldDefn(i)->IsIgnored())
1262 : {
1263 124 : m_anMapGeomFieldIndexToArrayIndex.push_back(
1264 62 : static_cast<int>(m_aosProjectedFields.size()));
1265 62 : m_aosProjectedFields.emplace_back(field->name());
1266 : }
1267 : else
1268 : {
1269 5 : m_anMapGeomFieldIndexToArrayIndex.push_back(-1);
1270 : }
1271 : }
1272 : }
1273 : }
1274 :
1275 103 : m_nExpectedBatchColumns =
1276 103 : m_bIgnoredFields ? static_cast<int>(m_aosProjectedFields.size()) : -1;
1277 :
1278 : // Full invalidation
1279 103 : InvalidateCachedBatches();
1280 :
1281 103 : return eErr;
1282 : }
1283 :
1284 : /************************************************************************/
1285 : /* TestCapability() */
1286 : /************************************************************************/
1287 :
1288 234 : int OGRParquetDatasetLayer::TestCapability(const char *pszCap)
1289 : {
1290 234 : if (EQUAL(pszCap, OLCIgnoreFields))
1291 5 : return true;
1292 :
1293 229 : if (EQUAL(pszCap, OLCFastSpatialFilter))
1294 : {
1295 171 : if (m_iGeomFieldFilter >= 0 &&
1296 114 : m_iGeomFieldFilter < static_cast<int>(m_aeGeomEncoding.size()) &&
1297 57 : m_aeGeomEncoding[m_iGeomFieldFilter] ==
1298 : OGRArrowGeomEncoding::GEOARROW_STRUCT_POINT)
1299 : {
1300 8 : return true;
1301 : }
1302 : // fallback to base method
1303 : }
1304 :
1305 221 : return OGRParquetLayerBase::TestCapability(pszCap);
1306 : }
1307 :
1308 : /***********************************************************************/
1309 : /* SetAttributeFilter() */
1310 : /***********************************************************************/
1311 :
1312 191 : OGRErr OGRParquetDatasetLayer::SetAttributeFilter(const char *pszFilter)
1313 : {
1314 191 : m_bRebuildScanner = true;
1315 191 : return OGRParquetLayerBase::SetAttributeFilter(pszFilter);
1316 : }
|