Line data Source code
1 : /******************************************************************************
2 : *
3 : * Project: GDAL
4 : * Purpose: "gdal vector sort" subcommand
5 : * Author: Daniel Baston
6 : *
7 : ******************************************************************************
8 : * Copyright (c) 2025-2026, ISciences LLC
9 : *
10 : * SPDX-License-Identifier: MIT
11 : ****************************************************************************/
12 :
13 : #include "gdalalg_vector_sort.h"
14 :
15 : #include "cpl_error.h"
16 : #include "gdal_alg.h"
17 : #include "gdal_priv.h"
18 : #include "gdalalg_vector_geom.h"
19 : #include "ogr_geometry.h"
20 :
21 : #include "ogr_geos.h"
22 :
23 : #include <algorithm>
24 : #include <cinttypes>
25 : #include <limits>
26 :
27 : #ifndef _
28 : #define _(x) (x)
29 : #endif
30 :
31 : //! @cond Doxygen_Suppress
32 :
33 59 : GDALVectorSortAlgorithm::GDALVectorSortAlgorithm(bool standaloneStep)
34 : : GDALVectorPipelineStepAlgorithm(NAME, DESCRIPTION, HELP_URL,
35 59 : standaloneStep)
36 : {
37 : AddArg("geometry-field", 0, _("Name of geometry field to use in sort"),
38 59 : &m_geomField);
39 118 : AddArg("method", 0, _("Geometry sorting algorithm"), &m_sortMethod)
40 59 : .SetChoices("hilbert", "strtree")
41 59 : .SetDefault(m_sortMethod);
42 : AddArg("use-tempfile", 0,
43 : _("Write features to a temporary file to avoid reading the entire "
44 : "input dataset into memory"),
45 59 : &m_useTempfile);
46 59 : }
47 :
48 : namespace
49 : {
50 :
51 : /// Simple container allowing to store features and retrieve each one _a single
52 : /// time_ using a random access pattern. The number of stored features does not
53 : /// need to be known at construction time.
54 : class GDALFeatureStore
55 : {
56 : public:
57 14 : virtual ~GDALFeatureStore() = default;
58 :
59 : virtual std::unique_ptr<OGRFeature> Load(std::size_t i) = 0;
60 :
61 : virtual bool Store(std::unique_ptr<OGRFeature> f) = 0;
62 :
63 : virtual std::size_t Size() const = 0;
64 : };
65 :
66 : /// FeatureStore backed by a temporary file on disk.
67 : class GDALFileFeatureStore : public GDALFeatureStore
68 : {
69 : public:
70 6 : GDALFileFeatureStore()
71 6 : : m_fileName(CPLGenerateTempFilenameSafe(nullptr)),
72 6 : m_file(VSIFOpenL(m_fileName.c_str(), "wb+"))
73 : {
74 6 : if (m_file == nullptr)
75 : {
76 0 : CPLError(CE_Failure, CPLE_AppDefined,
77 : "Failed to create temporary file");
78 : }
79 : else
80 : {
81 : // Unlink immediately so that the file is cleaned up if the process is killed
82 : // (at least on Linux)
83 6 : VSIUnlink(m_fileName.c_str());
84 : }
85 6 : }
86 :
87 12 : ~GDALFileFeatureStore() override
88 6 : {
89 6 : if (m_defn != nullptr)
90 : {
91 6 : const_cast<OGRFeatureDefn *>(m_defn)->Release();
92 : }
93 6 : if (m_file != nullptr)
94 : {
95 6 : VSIFCloseL(m_file);
96 : }
97 :
98 6 : VSIUnlink(m_fileName.c_str());
99 12 : }
100 :
101 400 : std::unique_ptr<OGRFeature> Load(std::size_t i) override
102 : {
103 400 : auto loc = m_locs[i];
104 400 : m_buf.resize(loc.size);
105 400 : if (VSIFSeekL(m_file, loc.offset, SEEK_SET) == -1)
106 : {
107 0 : return nullptr;
108 : }
109 400 : auto nBytesRead = VSIFReadL(m_buf.data(), 1, loc.size, m_file);
110 400 : if (nBytesRead != loc.size)
111 : {
112 0 : return nullptr;
113 : }
114 :
115 800 : auto poFeature = std::make_unique<OGRFeature>(m_defn);
116 400 : if (!poFeature->DeserializeFromBinary(m_buf.data(), m_buf.size()))
117 : {
118 0 : return nullptr;
119 : }
120 :
121 400 : return poFeature;
122 : }
123 :
124 812 : size_t Size() const override
125 : {
126 812 : return m_locs.size();
127 : }
128 :
129 400 : bool Store(std::unique_ptr<OGRFeature> f) override
130 : {
131 400 : if (m_file == nullptr)
132 : {
133 0 : return false;
134 : }
135 :
136 400 : if (m_defn == nullptr)
137 : {
138 6 : m_defn = f->GetDefnRef();
139 6 : const_cast<OGRFeatureDefn *>(m_defn)->Reference();
140 : }
141 :
142 400 : if (!f->SerializeToBinary(m_buf))
143 : {
144 0 : CPLError(CE_Failure, CPLE_AppDefined,
145 : "Failed to serialize feature to buffer");
146 0 : return false;
147 : }
148 :
149 : Loc loc;
150 400 : loc.offset = m_fileSize;
151 400 : loc.size = m_buf.size();
152 400 : m_locs.push_back(loc);
153 :
154 400 : auto nBytesWritten = VSIFWriteL(m_buf.data(), 1, m_buf.size(), m_file);
155 400 : if (nBytesWritten != m_buf.size())
156 : {
157 0 : CPLError(CE_Failure, CPLE_AppDefined,
158 : "Failed to write feature to temporary file");
159 0 : return false;
160 : }
161 :
162 400 : m_fileSize += loc.size;
163 :
164 400 : return true;
165 : }
166 :
167 : private:
168 : struct Loc
169 : {
170 : vsi_l_offset offset;
171 : std::size_t size;
172 : };
173 :
174 : std::string m_fileName{};
175 : const OGRFeatureDefn *m_defn{nullptr};
176 : vsi_l_offset m_fileSize{0};
177 : VSILFILE *m_file{nullptr};
178 : std::vector<Loc> m_locs{};
179 : std::vector<GByte> m_buf{};
180 :
181 : CPL_DISALLOW_COPY_ASSIGN(GDALFileFeatureStore)
182 : };
183 :
184 : /// FeatureStore backed by a std::vector.
185 : class GDALMemFeatureStore : public GDALFeatureStore
186 : {
187 : public:
188 402 : std::unique_ptr<OGRFeature> Load(std::size_t i) override
189 : {
190 402 : return std::unique_ptr<OGRFeature>(m_features[i]->Clone());
191 : }
192 :
193 816 : size_t Size() const override
194 : {
195 816 : return m_features.size();
196 : }
197 :
198 420 : bool Store(std::unique_ptr<OGRFeature> f) override
199 : {
200 420 : m_features.push_back(std::move(f));
201 420 : return true;
202 : }
203 :
204 : private:
205 : std::vector<std::unique_ptr<OGRFeature>> m_features{};
206 : };
207 :
208 : /**
209 : * This base class provides common functionality for layers representing different
210 : * sorting algorithms. An implementation's Process() method should:
211 : * - read the input features and transfer them to the feature store
212 : * - populate the m_sortedIndices vector
213 : */
214 : class GDALVectorSortedLayer : public GDALVectorNonStreamingAlgorithmLayer
215 : {
216 : public:
217 14 : GDALVectorSortedLayer(OGRLayer &srcLayer, int geomFieldIndex,
218 : bool processInMemory)
219 14 : : GDALVectorNonStreamingAlgorithmLayer(srcLayer, geomFieldIndex),
220 14 : m_store(nullptr), m_processInMemory(processInMemory), m_readPos(0)
221 : {
222 14 : }
223 :
224 0 : const OGRFeatureDefn *GetLayerDefn() const override
225 : {
226 0 : return m_srcLayer.GetLayerDefn();
227 : }
228 :
229 0 : int TestCapability(const char *pszCap) const override
230 : {
231 0 : if (EQUAL(pszCap, OLCFastFeatureCount) ||
232 0 : EQUAL(pszCap, OLCFastGetExtent) ||
233 0 : EQUAL(pszCap, OLCFastGetExtent3D) ||
234 0 : EQUAL(pszCap, OLCStringsAsUTF8) || EQUAL(pszCap, OLCIgnoreFields) ||
235 0 : EQUAL(pszCap, OLCCurveGeometries) ||
236 0 : EQUAL(pszCap, OLCMeasuredGeometries) ||
237 0 : EQUAL(pszCap, OLCZGeometries))
238 : {
239 0 : return m_srcLayer.TestCapability(pszCap);
240 : }
241 :
242 0 : return false;
243 : }
244 :
245 12 : GIntBig GetFeatureCount(int bForce) override
246 : {
247 12 : if (!m_poAttrQuery && !m_poFilterGeom)
248 : {
249 12 : return m_srcLayer.GetFeatureCount(bForce);
250 : }
251 :
252 0 : return OGRLayer::GetFeatureCount(bForce);
253 : }
254 :
255 814 : std::unique_ptr<OGRFeature> GetNextProcessedFeature() override
256 : {
257 814 : CPLAssert(m_sortedIndices.size() == m_store->Size());
258 :
259 814 : if (m_readPos < m_store->Size())
260 : {
261 802 : return m_store->Load(m_sortedIndices[m_readPos++]);
262 : }
263 :
264 12 : return nullptr;
265 : }
266 :
267 12 : void ResetReading() override
268 : {
269 12 : m_readPos = 0;
270 12 : }
271 :
272 : protected:
273 14 : void Init()
274 : {
275 14 : if (m_processInMemory)
276 : {
277 8 : m_store = std::make_unique<GDALMemFeatureStore>();
278 : }
279 : else
280 : {
281 6 : m_store = std::make_unique<GDALFileFeatureStore>();
282 : }
283 :
284 14 : m_readPos = 0;
285 14 : }
286 :
287 : std::unique_ptr<GDALFeatureStore> m_store;
288 : std::vector<size_t> m_sortedIndices{};
289 :
290 : private:
291 : bool m_processInMemory;
292 : size_t m_readPos;
293 : };
294 :
295 : class GDALVectorHilbertSortLayer : public GDALVectorSortedLayer
296 : {
297 : public:
298 : using GDALVectorSortedLayer::GDALVectorSortedLayer;
299 :
300 8 : bool Process(GDALProgressFunc pfnProgress, void *pProgressData) override
301 : {
302 8 : Init();
303 :
304 : const GIntBig nLayerFeatures =
305 8 : m_srcLayer.TestCapability(OLCFastFeatureCount)
306 8 : ? m_srcLayer.GetFeatureCount(false)
307 8 : : -1;
308 : const double dfInvLayerFeatures =
309 8 : 1.0 / std::max(1.0, static_cast<double>(nLayerFeatures));
310 8 : const double dfFirstPhaseProgressRatio =
311 : dfInvLayerFeatures * (2.0 / 3.0);
312 :
313 16 : std::vector<OGREnvelope> envelopes;
314 8 : OGREnvelope oLayerExtent;
315 428 : for (auto &poFeature : m_srcLayer)
316 : {
317 : const OGRGeometry *poGeom =
318 420 : poFeature->GetGeomFieldRef(m_geomFieldIndex);
319 :
320 420 : envelopes.emplace_back();
321 :
322 420 : if (poGeom != nullptr && !poGeom->IsEmpty())
323 : {
324 404 : poGeom->getEnvelope(&envelopes.back());
325 404 : oLayerExtent.Merge(envelopes.back());
326 : }
327 :
328 840 : if (!m_store->Store(
329 840 : std::unique_ptr<OGRFeature>(poFeature.release())))
330 : {
331 0 : return false;
332 : }
333 :
334 692 : if (pfnProgress && nLayerFeatures > 0 &&
335 272 : !pfnProgress(static_cast<double>(envelopes.size()) *
336 : dfFirstPhaseProgressRatio,
337 : "", pProgressData))
338 : {
339 0 : CPLError(CE_Failure, CPLE_UserInterrupt, "Interrupted by user");
340 0 : return false;
341 : }
342 : }
343 :
344 : std::vector<std::pair<std::size_t, std::uint32_t>> hilbertCodes(
345 8 : envelopes.size());
346 428 : for (std::size_t i = 0; i < envelopes.size(); i++)
347 : {
348 420 : hilbertCodes[i].first = i;
349 :
350 420 : if (envelopes[i].IsInit())
351 : {
352 : double dfX, dfY;
353 404 : envelopes[i].Center(dfX, dfY);
354 404 : hilbertCodes[i].second =
355 404 : GDALHilbertCode(&oLayerExtent, dfX, dfY);
356 : }
357 : else
358 : {
359 16 : hilbertCodes[i].second =
360 16 : std::numeric_limits<std::uint32_t>::max();
361 : }
362 : }
363 :
364 8 : std::sort(hilbertCodes.begin(), hilbertCodes.end(),
365 2858 : [](const auto &a, const auto &b)
366 2858 : { return a.second < b.second; });
367 :
368 428 : for (const auto &sItem : hilbertCodes)
369 : {
370 420 : m_sortedIndices.push_back(sItem.first);
371 : }
372 :
373 8 : if (pfnProgress)
374 : {
375 4 : pfnProgress(1.0, "", pProgressData);
376 : }
377 :
378 8 : return true;
379 : }
380 :
381 : private:
382 : CPL_DISALLOW_COPY_ASSIGN(GDALVectorHilbertSortLayer)
383 : };
384 :
385 : #ifdef HAVE_GEOS
386 : class GDALVectorSTRTreeSortLayer : public GDALVectorSortedLayer
387 : {
388 : public:
389 : using GDALVectorSortedLayer::GDALVectorSortedLayer;
390 :
391 12 : ~GDALVectorSTRTreeSortLayer() override
392 6 : {
393 6 : if (m_geosContext)
394 : {
395 6 : if (m_poTree)
396 : {
397 6 : GEOSSTRtree_destroy_r(m_geosContext, m_poTree);
398 : }
399 :
400 6 : finishGEOS_r(m_geosContext);
401 : }
402 12 : }
403 :
404 6 : bool Process(GDALProgressFunc pfnProgress, void *pProgressData) override
405 : {
406 6 : Init();
407 :
408 : const GIntBig nLayerFeatures =
409 6 : m_srcLayer.TestCapability(OLCFastFeatureCount)
410 6 : ? m_srcLayer.GetFeatureCount(false)
411 6 : : -1;
412 : const double dfInvLayerFeatures =
413 6 : 1.0 / std::max(1.0, static_cast<double>(nLayerFeatures));
414 6 : const double dfFirstPhaseProgressRatio =
415 : dfInvLayerFeatures * (2.0 / 3.0);
416 :
417 : // TODO: variant of this fn returning unique_ptr
418 6 : m_geosContext = OGRGeometry::createGEOSContext();
419 6 : m_poTree = GEOSSTRtree_create_r(m_geosContext, 10);
420 :
421 6 : OGREnvelope oGeomExtent;
422 12 : std::vector<size_t> nullIndices;
423 6 : std::uintptr_t i = 0;
424 406 : for (auto &poFeature : m_srcLayer)
425 : {
426 :
427 : const OGRGeometry *poGeom =
428 400 : poFeature->GetGeomFieldRef(m_geomFieldIndex);
429 :
430 400 : if (poGeom == nullptr || poGeom->IsEmpty())
431 : {
432 16 : nullIndices.push_back(i);
433 : }
434 : else
435 : {
436 384 : poGeom->getEnvelope(&oGeomExtent);
437 384 : if (!InsertIntoTree(oGeomExtent, i))
438 : {
439 0 : return false;
440 : }
441 : }
442 :
443 800 : if (!m_store->Store(
444 800 : std::unique_ptr<OGRFeature>(poFeature.release())))
445 : {
446 0 : return false;
447 : }
448 :
449 400 : i++;
450 :
451 672 : if (pfnProgress && nLayerFeatures > 0 &&
452 272 : !pfnProgress(static_cast<double>(i) * dfFirstPhaseProgressRatio,
453 : "", pProgressData))
454 : {
455 0 : CPLError(CE_Failure, CPLE_UserInterrupt, "Interrupted by user");
456 0 : return false;
457 : }
458 : }
459 :
460 6 : BuildTree();
461 :
462 6 : m_sortedIndices = ReadTreeIndices();
463 6 : m_sortedIndices.insert(m_sortedIndices.end(), nullIndices.begin(),
464 12 : nullIndices.end());
465 :
466 6 : if (pfnProgress)
467 : {
468 4 : pfnProgress(1.0, "", pProgressData);
469 : }
470 :
471 6 : return true;
472 : }
473 :
474 : private:
475 384 : bool InsertIntoTree(const OGREnvelope &oGeomExtent, std::uintptr_t i)
476 : {
477 384 : GEOSGeometry *poEnv = CreateGEOSEnvelope(oGeomExtent);
478 384 : if (poEnv == nullptr)
479 : {
480 0 : return false;
481 : }
482 384 : GEOSSTRtree_insert_r(m_geosContext, m_poTree, poEnv,
483 : reinterpret_cast<void *>(i));
484 384 : GEOSGeom_destroy_r(m_geosContext, poEnv);
485 :
486 384 : return true;
487 : }
488 :
489 6 : bool BuildTree()
490 : {
491 : #if GEOS_VERSION_MAJOR > 3 || \
492 : (GEOS_VERSION_MAJOR == 3 && GEOS_VERSION_MINOR >= 12)
493 6 : GEOSSTRtree_build_r(m_geosContext, m_poTree);
494 : #else
495 : if (m_store->Size() > 0)
496 : {
497 : // Perform a dummy query to force tree construction.
498 : OGREnvelope oExtent;
499 : oExtent.MinX = oExtent.MaxX = oExtent.MinY = oExtent.MaxY = 0;
500 : GEOSGeometry *poEnv = CreateGEOSEnvelope(oExtent);
501 : if (poEnv == nullptr)
502 : {
503 : return false;
504 : }
505 : GEOSSTRtree_query_r(
506 : m_geosContext, m_poTree, poEnv, [](void *, void *) {}, nullptr);
507 : GEOSGeom_destroy_r(m_geosContext, poEnv);
508 : }
509 : #endif
510 6 : return true;
511 : }
512 :
513 6 : std::vector<size_t> ReadTreeIndices()
514 : {
515 6 : std::vector<size_t> sortedIndices;
516 :
517 6 : GEOSSTRtree_iterate_r(
518 : m_geosContext, m_poTree,
519 384 : [](void *item, void *userData)
520 : {
521 384 : static_cast<std::vector<size_t> *>(userData)->push_back(
522 384 : reinterpret_cast<std::uintptr_t>(item));
523 384 : },
524 : &sortedIndices);
525 :
526 6 : return sortedIndices;
527 : }
528 :
529 : // FIXME: Duplicated from alg/zonal.cpp.
530 : // Put into OGRGeometryFactory?
531 384 : GEOSGeometry *CreateGEOSEnvelope(const OGREnvelope &oEnv) const
532 : {
533 384 : GEOSCoordSequence *seq = GEOSCoordSeq_create_r(m_geosContext, 2, 2);
534 384 : if (seq == nullptr)
535 : {
536 0 : return nullptr;
537 : }
538 384 : GEOSCoordSeq_setXY_r(m_geosContext, seq, 0, oEnv.MinX, oEnv.MinY);
539 384 : GEOSCoordSeq_setXY_r(m_geosContext, seq, 1, oEnv.MaxX, oEnv.MaxY);
540 384 : return GEOSGeom_createLineString_r(m_geosContext, seq);
541 : }
542 :
543 : GEOSContextHandle_t m_geosContext{nullptr};
544 : GEOSSTRtree *m_poTree{nullptr};
545 :
546 : CPL_DISALLOW_COPY_ASSIGN(GDALVectorSTRTreeSortLayer)
547 : };
548 : #endif
549 :
550 : } // namespace
551 :
552 20 : bool GDALVectorSortAlgorithm::RunStep(GDALPipelineStepRunContext &ctxt)
553 : {
554 20 : auto poSrcDS = m_inputDataset[0].GetDatasetRef();
555 40 : auto poDstDS = std::make_unique<GDALVectorNonStreamingAlgorithmDataset>();
556 :
557 40 : GDALVectorAlgorithmLayerProgressHelper progressHelper(ctxt);
558 :
559 40 : for (auto &&poSrcLayer : poSrcDS->GetLayers())
560 : {
561 20 : if (m_inputLayerNames.empty() ||
562 0 : std::find(m_inputLayerNames.begin(), m_inputLayerNames.end(),
563 20 : poSrcLayer->GetDescription()) != m_inputLayerNames.end())
564 : {
565 20 : const auto poSrcLayerDefn = poSrcLayer->GetLayerDefn();
566 20 : if (poSrcLayerDefn->GetGeomFieldCount() > 0)
567 : {
568 16 : progressHelper.AddProcessedLayer(*poSrcLayer);
569 : }
570 : else
571 : {
572 4 : progressHelper.AddPassThroughLayer(*poSrcLayer);
573 : }
574 : }
575 : }
576 :
577 38 : for (auto [poSrcLayer, bProcessed, layerProgressFunc, layerProgressData] :
578 56 : progressHelper)
579 : {
580 20 : if (bProcessed)
581 : {
582 16 : const auto poSrcLayerDefn = poSrcLayer->GetLayerDefn();
583 : const int geomFieldIndex =
584 16 : m_geomField.empty()
585 16 : ? 0
586 3 : : poSrcLayerDefn->GetGeomFieldIndex(m_geomField.c_str());
587 :
588 16 : if (geomFieldIndex == -1)
589 : {
590 2 : ReportError(CE_Failure, CPLE_AppDefined,
591 : "Specified geometry field '%s' does not exist in "
592 : "layer '%s'",
593 2 : m_geomField.c_str(), poSrcLayer->GetDescription());
594 2 : return false;
595 : }
596 :
597 0 : std::unique_ptr<GDALVectorSortedLayer> layer;
598 :
599 14 : if (m_sortMethod == "hilbert")
600 : {
601 8 : layer = std::make_unique<GDALVectorHilbertSortLayer>(
602 16 : *poSrcLayer, geomFieldIndex, !m_useTempfile);
603 : }
604 : else
605 : {
606 : // Already checked for invalid method at arg parsing stage.
607 6 : CPLAssert(m_sortMethod == "strtree");
608 : #ifdef HAVE_GEOS
609 6 : layer = std::make_unique<GDALVectorSTRTreeSortLayer>(
610 12 : *poSrcLayer, geomFieldIndex, !m_useTempfile);
611 : #else
612 : CPLError(CE_Failure, CPLE_AppDefined,
613 : "--method strtree requires a GDAL build against the "
614 : "GEOS library.");
615 : return false;
616 : #endif
617 : }
618 :
619 14 : if (!poDstDS->AddProcessedLayer(std::move(layer), layerProgressFunc,
620 : layerProgressData.get()))
621 : {
622 0 : return false;
623 : }
624 : }
625 : else
626 : {
627 4 : poDstDS->AddPassThroughLayer(*poSrcLayer);
628 : }
629 : }
630 :
631 18 : m_outputDataset.Set(std::move(poDstDS));
632 :
633 18 : return true;
634 : }
635 :
636 : GDALVectorSortAlgorithmStandalone::~GDALVectorSortAlgorithmStandalone() =
637 : default;
638 :
639 : //! @endcond
|