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