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, ISciences LLC
       9             :  *
      10             :  * SPDX-License-Identifier: MIT
      11             :  ****************************************************************************/
      12             : 
      13             : #include "gdalalg_vector_sort.h"
      14             : 
      15             : #include "cpl_enumerate.h"
      16             : #include "cpl_error.h"
      17             : #include "gdal_alg.h"
      18             : #include "gdal_priv.h"
      19             : #include "gdalalg_vector_geom.h"
      20             : #include "ogr_geometry.h"
      21             : 
      22             : #include "ogr_geos.h"
      23             : 
      24             : #include <algorithm>
      25             : #include <cinttypes>
      26             : #include <limits>
      27             : 
      28             : #ifndef _
      29             : #define _(x) (x)
      30             : #endif
      31             : 
      32             : //! @cond Doxygen_Suppress
      33             : 
      34          50 : GDALVectorSortAlgorithm::GDALVectorSortAlgorithm(bool standaloneStep)
      35             :     : GDALVectorPipelineStepAlgorithm(NAME, DESCRIPTION, HELP_URL,
      36          50 :                                       standaloneStep)
      37             : {
      38             :     AddArg("geometry-field", 0, _("Name of geometry field to use in sort"),
      39          50 :            &m_geomField);
      40         100 :     AddArg("method", 0, _("Geometry sorting algorithm"), &m_sortMethod)
      41          50 :         .SetChoices("hilbert", "strtree")
      42          50 :         .SetDefault(m_sortMethod);
      43          50 : }
      44             : 
      45             : namespace
      46             : {
      47             : 
      48           8 : bool CreateDstFeatures(std::vector<std::unique_ptr<OGRFeature>> &srcFeatures,
      49             :                        const std::vector<size_t> &sortedIndices,
      50             :                        OGRLayer &dstLayer, GDALProgressFunc pfnProgress,
      51             :                        void *pProgressData, double dfProgressStart,
      52             :                        double dfProgressRatio)
      53             : {
      54           8 :     uint64_t nCounter = 0;
      55         428 :     for (size_t iSrcFeature : sortedIndices)
      56             :     {
      57         420 :         auto &poSrcFeature = srcFeatures[iSrcFeature];
      58         420 :         poSrcFeature->SetFDefnUnsafe(dstLayer.GetLayerDefn());
      59         420 :         poSrcFeature->SetFID(OGRNullFID);
      60             : 
      61         420 :         if (dstLayer.CreateFeature(std::move(poSrcFeature)) != OGRERR_NONE)
      62             :         {
      63           0 :             return false;
      64             :         }
      65         692 :         if (pfnProgress &&
      66         272 :             !pfnProgress(dfProgressStart +
      67         272 :                              static_cast<double>(++nCounter) * dfProgressRatio,
      68             :                          "", pProgressData))
      69             :         {
      70           0 :             CPLError(CE_Failure, CPLE_UserInterrupt, "Interrupted by user");
      71           0 :             return false;
      72             :         }
      73             :     }
      74             : 
      75           8 :     return true;
      76             : }
      77             : 
      78             : class GDALVectorHilbertSortDataset
      79             :     : public GDALVectorNonStreamingAlgorithmDataset
      80             : {
      81             :   public:
      82             :     using GDALVectorNonStreamingAlgorithmDataset::
      83             :         GDALVectorNonStreamingAlgorithmDataset;
      84             : 
      85           5 :     bool Process(OGRLayer &srcLayer, OGRLayer &dstLayer, int geomFieldIndex,
      86             :                  GDALProgressFunc pfnProgress, void *pProgressData) override
      87             :     {
      88          10 :         std::vector<std::unique_ptr<OGRFeature>> features;
      89             : 
      90             :         const GIntBig nLayerFeatures =
      91           5 :             srcLayer.TestCapability(OLCFastFeatureCount)
      92           5 :                 ? srcLayer.GetFeatureCount(false)
      93           5 :                 : -1;
      94             :         const double dfInvLayerFeatures =
      95           5 :             1.0 / std::max(1.0, static_cast<double>(nLayerFeatures));
      96           5 :         const double dfFirstPhaseProgressRatio =
      97             :             dfInvLayerFeatures * (2.0 / 3.0);
      98         225 :         for (auto &feature : srcLayer)
      99             :         {
     100         220 :             features.emplace_back(feature.release());
     101         356 :             if (pfnProgress && nLayerFeatures > 0 &&
     102         136 :                 !pfnProgress(static_cast<double>(features.size()) *
     103             :                                  dfFirstPhaseProgressRatio,
     104             :                              "", pProgressData))
     105             :             {
     106           0 :                 ReportError(CE_Failure, CPLE_UserInterrupt,
     107             :                             "Interrupted by user");
     108           0 :                 return false;
     109             :             }
     110             :         }
     111           5 :         OGREnvelope oLayerExtent;
     112             : 
     113          10 :         std::vector<OGREnvelope> envelopes(features.size());
     114         225 :         for (const auto &[i, poFeature] : cpl::enumerate(features))
     115             :         {
     116             :             const OGRGeometry *poGeom =
     117         220 :                 poFeature->GetGeomFieldRef(geomFieldIndex);
     118             : 
     119         220 :             if (poGeom != nullptr && !poGeom->IsEmpty())
     120             :             {
     121         212 :                 poGeom->getEnvelope(&envelopes[i]);
     122         212 :                 oLayerExtent.Merge(envelopes[i]);
     123             :             }
     124             :         }
     125             : 
     126             :         std::vector<std::pair<std::size_t, std::uint32_t>> hilbertCodes(
     127          10 :             features.size());
     128         225 :         for (std::size_t i = 0; i < features.size(); i++)
     129             :         {
     130         220 :             hilbertCodes[i].first = i;
     131             : 
     132         220 :             if (envelopes[i].IsInit())
     133             :             {
     134             :                 double dfX, dfY;
     135         212 :                 envelopes[i].Center(dfX, dfY);
     136         212 :                 hilbertCodes[i].second =
     137         212 :                     GDALHilbertCode(&oLayerExtent, dfX, dfY);
     138             :             }
     139             :             else
     140             :             {
     141           8 :                 hilbertCodes[i].second =
     142           8 :                     std::numeric_limits<std::uint32_t>::max();
     143             :             }
     144             :         }
     145             : 
     146           5 :         std::sort(hilbertCodes.begin(), hilbertCodes.end(),
     147        1455 :                   [](const auto &a, const auto &b)
     148        1455 :                   { return a.second < b.second; });
     149             : 
     150          10 :         std::vector<size_t> sortedIndices;
     151         225 :         for (const auto &sItem : hilbertCodes)
     152             :         {
     153         220 :             sortedIndices.push_back(sItem.first);
     154             :         }
     155             : 
     156           5 :         const double dfProgressStart = nLayerFeatures > 0 ? 2.0 / 3.0 : 0.0;
     157             :         const double dfProgressRatio =
     158           5 :             (nLayerFeatures > 0 ? 1.0 / 3.0 : 1.0) /
     159           5 :             std::max(1.0, static_cast<double>(features.size()));
     160           5 :         return CreateDstFeatures(features, sortedIndices, dstLayer, pfnProgress,
     161             :                                  pProgressData, dfProgressStart,
     162           5 :                                  dfProgressRatio);
     163             :     }
     164             : 
     165             :   private:
     166             :     CPL_DISALLOW_COPY_ASSIGN(GDALVectorHilbertSortDataset)
     167             : };
     168             : 
     169             : #ifdef HAVE_GEOS
     170             : class GDALVectorSTRTreeSortDataset
     171             :     : public GDALVectorNonStreamingAlgorithmDataset
     172             : {
     173             :   public:
     174             :     using GDALVectorNonStreamingAlgorithmDataset::
     175             :         GDALVectorNonStreamingAlgorithmDataset;
     176             : 
     177           8 :     ~GDALVectorSTRTreeSortDataset() override
     178           4 :     {
     179           4 :         if (m_geosContext)
     180             :         {
     181           3 :             finishGEOS_r(m_geosContext);
     182             :         }
     183           8 :     }
     184             : 
     185           3 :     bool Process(OGRLayer &srcLayer, OGRLayer &dstLayer, int geomFieldIndex,
     186             :                  GDALProgressFunc pfnProgress, void *pProgressData) override
     187             :     {
     188           6 :         std::vector<std::unique_ptr<OGRFeature>> features;
     189           6 :         std::vector<size_t> sortedIndices;
     190             : 
     191             :         const GIntBig nLayerFeatures =
     192           3 :             srcLayer.TestCapability(OLCFastFeatureCount)
     193           3 :                 ? srcLayer.GetFeatureCount(false)
     194           3 :                 : -1;
     195             :         const double dfInvLayerFeatures =
     196           3 :             1.0 / std::max(1.0, static_cast<double>(nLayerFeatures));
     197           3 :         const double dfFirstPhaseProgressRatio =
     198             :             dfInvLayerFeatures * (2.0 / 3.0);
     199         203 :         for (auto &feature : srcLayer)
     200             :         {
     201         200 :             features.emplace_back(feature.release());
     202         336 :             if (pfnProgress && nLayerFeatures > 0 &&
     203         136 :                 !pfnProgress(static_cast<double>(features.size()) *
     204             :                                  dfFirstPhaseProgressRatio,
     205             :                              "", pProgressData))
     206             :             {
     207           0 :                 ReportError(CE_Failure, CPLE_UserInterrupt,
     208             :                             "Interrupted by user");
     209           0 :                 return false;
     210             :             }
     211             :         }
     212             :         // TODO: variant of this fn returning unique_ptr
     213           3 :         m_geosContext = OGRGeometry::createGEOSContext();
     214             : 
     215           3 :         auto TreeDeleter = [this](GEOSSTRtree *tree)
     216           3 :         { GEOSSTRtree_destroy_r(m_geosContext, tree); };
     217             : 
     218             :         std::unique_ptr<GEOSSTRtree, decltype(TreeDeleter)> poTree(
     219           6 :             GEOSSTRtree_create_r(m_geosContext, 10), TreeDeleter);
     220             : 
     221           3 :         OGREnvelope oGeomExtent;
     222           6 :         std::vector<size_t> nullIndices;
     223         203 :         for (const auto &[i, poFeature] : cpl::enumerate(features))
     224             :         {
     225             :             const OGRGeometry *poGeom =
     226         200 :                 poFeature->GetGeomFieldRef(geomFieldIndex);
     227             : 
     228         200 :             if (poGeom == nullptr || poGeom->IsEmpty())
     229             :             {
     230           8 :                 nullIndices.push_back(i);
     231           8 :                 continue;
     232             :             }
     233             : 
     234         192 :             poGeom->getEnvelope(&oGeomExtent);
     235         192 :             GEOSGeometry *poEnv = CreateGEOSEnvelope(oGeomExtent);
     236         192 :             if (poEnv == nullptr)
     237             :             {
     238           0 :                 return false;
     239             :             }
     240             : 
     241         192 :             GEOSSTRtree_insert_r(m_geosContext, poTree.get(), poEnv,
     242         192 :                                  reinterpret_cast<void *>(i));
     243         192 :             GEOSGeom_destroy_r(m_geosContext, poEnv);
     244             :         }
     245             : 
     246             : #if GEOS_VERSION_MAJOR > 3 ||                                                  \
     247             :     (GEOS_VERSION_MAJOR == 3 && GEOS_VERSION_MINOR >= 12)
     248           3 :         GEOSSTRtree_build_r(m_geosContext, poTree.get());
     249             : #else
     250             :         if (!features.empty())
     251             :         {
     252             :             // Perform a dummy query to force tree construction.
     253             :             GEOSGeometry *poEnv = CreateGEOSEnvelope(oGeomExtent);
     254             :             GEOSSTRtree_query_r(
     255             :                 m_geosContext, poTree.get(), poEnv, [](void *, void *) {},
     256             :                 nullptr);
     257             :             GEOSGeom_destroy_r(m_geosContext, poEnv);
     258             :         }
     259             : #endif
     260             : 
     261           3 :         GEOSSTRtree_iterate_r(
     262             :             m_geosContext, poTree.get(),
     263         192 :             [](void *item, void *userData)
     264             :             {
     265         192 :                 static_cast<std::vector<size_t> *>(userData)->push_back(
     266         192 :                     reinterpret_cast<size_t>(item));
     267         192 :             },
     268             :             &sortedIndices);
     269             : 
     270           3 :         sortedIndices.insert(sortedIndices.end(), nullIndices.begin(),
     271           6 :                              nullIndices.end());
     272           3 :         nullIndices.clear();
     273             : 
     274           3 :         const double dfProgressStart = nLayerFeatures > 0 ? 2.0 / 3.0 : 0.0;
     275             :         const double dfProgressRatio =
     276           3 :             (nLayerFeatures > 0 ? 1.0 / 3.0 : 1.0) /
     277           3 :             std::max(1.0, static_cast<double>(features.size()));
     278           3 :         return CreateDstFeatures(features, sortedIndices, dstLayer, pfnProgress,
     279             :                                  pProgressData, dfProgressStart,
     280           3 :                                  dfProgressRatio);
     281             :     }
     282             : 
     283             :   private:
     284             :     CPL_DISALLOW_COPY_ASSIGN(GDALVectorSTRTreeSortDataset)
     285             : 
     286             :     // FIXME: Duplicated from alg/zonal.cpp.
     287             :     // Put into OGRGeometryFactory?
     288         192 :     GEOSGeometry *CreateGEOSEnvelope(const OGREnvelope &oEnv) const
     289             :     {
     290         192 :         GEOSCoordSequence *seq = GEOSCoordSeq_create_r(m_geosContext, 2, 2);
     291         192 :         if (seq == nullptr)
     292             :         {
     293           0 :             return nullptr;
     294             :         }
     295         192 :         GEOSCoordSeq_setXY_r(m_geosContext, seq, 0, oEnv.MinX, oEnv.MinY);
     296         192 :         GEOSCoordSeq_setXY_r(m_geosContext, seq, 1, oEnv.MaxX, oEnv.MaxY);
     297         192 :         return GEOSGeom_createLineString_r(m_geosContext, seq);
     298             :     }
     299             : 
     300             :     GEOSContextHandle_t m_geosContext{nullptr};
     301             : };
     302             : #endif
     303             : 
     304             : }  // namespace
     305             : 
     306          12 : bool GDALVectorSortAlgorithm::RunStep(GDALPipelineStepRunContext &ctxt)
     307             : {
     308          12 :     auto poSrcDS = m_inputDataset[0].GetDatasetRef();
     309          12 :     std::unique_ptr<GDALVectorNonStreamingAlgorithmDataset> poDstDS;
     310             : 
     311          12 :     if (m_sortMethod == "hilbert")
     312             :     {
     313           8 :         poDstDS = std::make_unique<GDALVectorHilbertSortDataset>();
     314             :     }
     315             :     else
     316             :     {
     317             :         // Already checked for invalid method at arg parsing stage.
     318           4 :         CPLAssert(m_sortMethod == "strtree");
     319             : #ifdef HAVE_GEOS
     320           4 :         poDstDS = std::make_unique<GDALVectorSTRTreeSortDataset>();
     321             : #else
     322             :         CPLError(
     323             :             CE_Failure, CPLE_AppDefined,
     324             :             "--method strtree requires a GDAL build against the GEOS library.");
     325             :         return false;
     326             : #endif
     327             :     }
     328             : 
     329          24 :     GDALVectorAlgorithmLayerProgressHelper progressHelper(ctxt);
     330             : 
     331          24 :     for (auto &&poSrcLayer : poSrcDS->GetLayers())
     332             :     {
     333          12 :         if (m_inputLayerNames.empty() ||
     334           0 :             std::find(m_inputLayerNames.begin(), m_inputLayerNames.end(),
     335          12 :                       poSrcLayer->GetDescription()) != m_inputLayerNames.end())
     336             :         {
     337          12 :             const auto poSrcLayerDefn = poSrcLayer->GetLayerDefn();
     338          12 :             if (poSrcLayerDefn->GetGeomFieldCount() > 0)
     339             :             {
     340          10 :                 progressHelper.AddProcessedLayer(*poSrcLayer);
     341             :             }
     342             :             else
     343             :             {
     344           2 :                 progressHelper.AddPassThroughLayer(*poSrcLayer);
     345             :             }
     346             :         }
     347             :     }
     348             : 
     349          22 :     for (auto [poSrcLayer, bProcessed, layerProgressFunc, layerProgressData] :
     350          32 :          progressHelper)
     351             :     {
     352          12 :         if (bProcessed)
     353             :         {
     354          10 :             const auto poSrcLayerDefn = poSrcLayer->GetLayerDefn();
     355             :             const int geomFieldIndex =
     356          10 :                 m_geomField.empty()
     357          10 :                     ? 0
     358           3 :                     : poSrcLayerDefn->GetGeomFieldIndex(m_geomField.c_str());
     359             : 
     360          10 :             if (geomFieldIndex == -1)
     361             :             {
     362           2 :                 ReportError(CE_Failure, CPLE_AppDefined,
     363             :                             "Specified geometry field '%s' does not exist in "
     364             :                             "layer '%s'",
     365           2 :                             m_geomField.c_str(), poSrcLayer->GetDescription());
     366           2 :                 return false;
     367             :             }
     368             : 
     369          16 :             if (!poDstDS->AddProcessedLayer(
     370           8 :                     *poSrcLayer, *poSrcLayer->GetLayerDefn(), geomFieldIndex,
     371             :                     layerProgressFunc, layerProgressData.get()))
     372             :             {
     373           0 :                 return false;
     374             :             }
     375             :         }
     376             :         else
     377             :         {
     378           2 :             poDstDS->AddPassThroughLayer(*poSrcLayer);
     379             :         }
     380             :     }
     381             : 
     382          10 :     m_outputDataset.Set(std::move(poDstDS));
     383             : 
     384          10 :     return true;
     385             : }
     386             : 
     387             : GDALVectorSortAlgorithmStandalone::~GDALVectorSortAlgorithmStandalone() =
     388             :     default;
     389             : 
     390             : //! @endcond