Line data    Source code

       1             : /******************************************************************************
       2             :  *
       3             :  * Project:  GDAL
       4             :  * Purpose:  "limit" step of "vector pipeline"
       5             :  * Author:   Dan Baston
       6             :  *
       7             :  ******************************************************************************
       8             :  * Copyright (c) 2025, ISciences LLC
       9             :  *
      10             :  * SPDX-License-Identifier: MIT
      11             :  ****************************************************************************/
      12             : 
      13             : #include "gdalalg_vector_limit.h"
      14             : #include "gdalalg_vector_pipeline.h"
      15             : 
      16             : #include "gdal_priv.h"
      17             : #include "ogrsf_frmts.h"
      18             : #include "ogr_p.h"
      19             : 
      20             : #include <algorithm>
      21             : #include <set>
      22             : 
      23             : //! @cond Doxygen_Suppress
      24             : 
      25             : #ifndef _
      26             : #define _(x) (x)
      27             : #endif
      28             : 
      29             : /************************************************************************/
      30             : /*         GDALVectorLimitAlgorithm::GDALVectorLimitAlgorithm()       */
      31             : /************************************************************************/
      32             : 
      33          36 : GDALVectorLimitAlgorithm::GDALVectorLimitAlgorithm(bool standaloneStep)
      34             :     : GDALVectorPipelineStepAlgorithm(NAME, DESCRIPTION, HELP_URL,
      35          36 :                                       standaloneStep)
      36             : {
      37             :     AddArg("limit", 0, _("Limit the number of features to read per layer"),
      38          72 :            &m_featureLimit)
      39          36 :         .SetPositional()
      40          36 :         .SetRequired();
      41          36 :     AddActiveLayerArg(&m_activeLayer);
      42          36 : }
      43             : 
      44             : namespace
      45             : {
      46             : 
      47             : /************************************************************************/
      48             : /*                      GDALVectorReadLimitedLayer                      */
      49             : /************************************************************************/
      50             : 
      51           4 : class GDALVectorReadLimitedLayer : public OGRLayer
      52             : {
      53             :   public:
      54           2 :     GDALVectorReadLimitedLayer(OGRLayer &layer, int featureLimit)
      55           2 :         : m_srcLayer(layer), m_featureLimit(featureLimit), m_featuresRead(0)
      56             :     {
      57           2 :     }
      58             : 
      59             :     ~GDALVectorReadLimitedLayer() override;
      60             : 
      61           8 :     OGRFeature *GetNextFeature() override
      62             :     {
      63           8 :         if (m_featuresRead < m_featureLimit)
      64             :         {
      65           6 :             m_featuresRead++;
      66           6 :             return m_srcLayer.GetNextFeature();
      67             :         }
      68           2 :         return nullptr;
      69             :     }
      70             : 
      71          18 :     const OGRFeatureDefn *GetLayerDefn() const override
      72             :     {
      73          18 :         return m_srcLayer.GetLayerDefn();
      74             :     }
      75             : 
      76           0 :     GIntBig GetFeatureCount(int bForce) override
      77             :     {
      78           0 :         return std::min(m_featureLimit, m_srcLayer.GetFeatureCount(bForce));
      79             :     }
      80             : 
      81           4 :     void ResetReading() override
      82             :     {
      83           4 :         m_featuresRead = 0;
      84           4 :         m_srcLayer.ResetReading();
      85           4 :     }
      86             : 
      87           2 :     int TestCapability(const char *pszCap) const override
      88             :     {
      89           2 :         return m_srcLayer.TestCapability(pszCap);
      90             :     }
      91             : 
      92             :   private:
      93             :     OGRLayer &m_srcLayer;
      94             :     GIntBig m_featureLimit;
      95             :     GIntBig m_featuresRead;
      96             : };
      97             : 
      98             : GDALVectorReadLimitedLayer::~GDALVectorReadLimitedLayer() = default;
      99             : 
     100             : }  // namespace
     101             : 
     102             : /************************************************************************/
     103             : /*               GDALVectorLimitAlgorithm::RunStep()                   */
     104             : /************************************************************************/
     105             : 
     106           1 : bool GDALVectorLimitAlgorithm::RunStep(GDALPipelineStepRunContext &)
     107             : {
     108           1 :     auto poSrcDS = m_inputDataset[0].GetDatasetRef();
     109           1 :     CPLAssert(poSrcDS);
     110             : 
     111           1 :     CPLAssert(m_outputDataset.GetName().empty());
     112           1 :     CPLAssert(!m_outputDataset.GetDatasetRef());
     113             : 
     114           1 :     auto outDS = std::make_unique<GDALVectorOutputDataset>();
     115             : 
     116           3 :     for (auto &&poSrcLayer : poSrcDS->GetLayers())
     117             :     {
     118           2 :         if (m_activeLayer.empty() ||
     119           0 :             m_activeLayer == poSrcLayer->GetDescription())
     120             :         {
     121           4 :             outDS->AddLayer(std::make_unique<GDALVectorReadLimitedLayer>(
     122           2 :                 *poSrcLayer, m_featureLimit));
     123             :         }
     124             :         else
     125             :         {
     126           0 :             outDS->AddLayer(
     127           0 :                 std::make_unique<GDALVectorPipelinePassthroughLayer>(
     128             :                     *poSrcLayer));
     129             :         }
     130             :     }
     131             : 
     132           1 :     m_outputDataset.Set(std::move(outDS));
     133             : 
     134           2 :     return true;
     135             : }
     136             : 
     137             : //! @endcond