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          42 : GDALVectorLimitAlgorithm::GDALVectorLimitAlgorithm(bool standaloneStep)
      34             :     : GDALVectorPipelineStepAlgorithm(NAME, DESCRIPTION, HELP_URL,
      35          42 :                                       standaloneStep)
      36             : {
      37             :     AddArg("limit", 0, _("Limit the number of features to read per layer"),
      38          84 :            &m_featureLimit)
      39          42 :         .SetPositional()
      40          42 :         .SetRequired();
      41          42 :     AddActiveLayerArg(&m_activeLayer);
      42          42 : }
      43             : 
      44             : namespace
      45             : {
      46             : 
      47             : /************************************************************************/
      48             : /*                      GDALVectorReadLimitedLayer                      */
      49             : /************************************************************************/
      50             : 
      51           8 : class GDALVectorReadLimitedLayer final
      52             :     : public OGRLayer,
      53             :       public OGRGetNextFeatureThroughRaw<GDALVectorReadLimitedLayer>
      54             : {
      55             :   public:
      56           4 :     GDALVectorReadLimitedLayer(OGRLayer &layer, int featureLimit)
      57           4 :         : m_srcLayer(layer), m_featureLimit(featureLimit), m_featuresRead(0)
      58             :     {
      59           4 :         SetDescription(layer.GetDescription());
      60           4 :     }
      61             : 
      62             :     ~GDALVectorReadLimitedLayer() override;
      63             : 
      64         367 :     OGRFeature *GetNextRawFeature()
      65             :     {
      66         367 :         if (m_featuresRead < m_featureLimit)
      67             :         {
      68         327 :             m_featuresRead++;
      69         327 :             return m_srcLayer.GetNextFeature();
      70             :         }
      71          40 :         return nullptr;
      72             :     }
      73             : 
      74         299 :     DEFINE_GET_NEXT_FEATURE_THROUGH_RAW(GDALVectorReadLimitedLayer)
      75             : 
      76         231 :     const OGRFeatureDefn *GetLayerDefn() const override
      77             :     {
      78         231 :         return m_srcLayer.GetLayerDefn();
      79             :     }
      80             : 
      81          13 :     GIntBig GetFeatureCount(int bForce) override
      82             :     {
      83          13 :         if (m_poFilterGeom || m_poAttrQuery)
      84           6 :             return OGRLayer::GetFeatureCount(bForce);
      85             :         else
      86           7 :             return std::min(m_featureLimit, m_srcLayer.GetFeatureCount(bForce));
      87             :     }
      88             : 
      89         115 :     void ResetReading() override
      90             :     {
      91         115 :         m_featuresRead = 0;
      92         115 :         m_srcLayer.ResetReading();
      93         115 :     }
      94             : 
      95          39 :     int TestCapability(const char *pszCap) const override
      96             :     {
      97          39 :         return m_srcLayer.TestCapability(pszCap);
      98             :     }
      99             : 
     100             :   private:
     101             :     OGRLayer &m_srcLayer;
     102             :     GIntBig m_featureLimit;
     103             :     GIntBig m_featuresRead;
     104             : };
     105             : 
     106             : GDALVectorReadLimitedLayer::~GDALVectorReadLimitedLayer() = default;
     107             : 
     108             : /************************************************************************/
     109             : /*                     GDALVectorReadLimitedDataset                     */
     110             : /************************************************************************/
     111             : 
     112             : class GDALVectorReadLimitedDataset final : public GDALVectorOutputDataset
     113             : {
     114             :   public:
     115           4 :     explicit GDALVectorReadLimitedDataset(GDALDataset *poSrcDS)
     116           4 :         : GDALVectorOutputDataset(poSrcDS)
     117             :     {
     118           4 :     }
     119             : 
     120           9 :     int TestCapability(const char *pszCap) const override
     121             :     {
     122           9 :         if (EQUAL(pszCap, ODsCMeasuredGeometries) ||
     123           7 :             EQUAL(pszCap, ODsCZGeometries))
     124             :         {
     125           4 :             return m_srcDS.TestCapability(pszCap);
     126             :         }
     127           5 :         return false;
     128             :     }
     129             : };
     130             : 
     131             : }  // namespace
     132             : 
     133             : /************************************************************************/
     134             : /*                 GDALVectorLimitAlgorithm::RunStep()                  */
     135             : /************************************************************************/
     136             : 
     137           4 : bool GDALVectorLimitAlgorithm::RunStep(GDALPipelineStepRunContext &)
     138             : {
     139           4 :     auto poSrcDS = m_inputDataset[0].GetDatasetRef();
     140           4 :     CPLAssert(poSrcDS);
     141             : 
     142           4 :     CPLAssert(m_outputDataset.GetName().empty());
     143           4 :     CPLAssert(!m_outputDataset.GetDatasetRef());
     144             : 
     145           4 :     auto outDS = std::make_unique<GDALVectorReadLimitedDataset>(poSrcDS);
     146             : 
     147           8 :     for (auto &&poSrcLayer : poSrcDS->GetLayers())
     148             :     {
     149           4 :         if (m_activeLayer.empty() ||
     150           0 :             m_activeLayer == poSrcLayer->GetDescription())
     151             :         {
     152           8 :             outDS->AddLayer(std::make_unique<GDALVectorReadLimitedLayer>(
     153           4 :                 *poSrcLayer, m_featureLimit));
     154             :         }
     155             :         else
     156             :         {
     157           0 :             outDS->AddLayer(
     158           0 :                 std::make_unique<GDALVectorPipelinePassthroughLayer>(
     159             :                     *poSrcLayer));
     160             :         }
     161             :     }
     162             : 
     163           4 :     m_outputDataset.Set(std::move(outDS));
     164             : 
     165           8 :     return true;
     166             : }
     167             : 
     168             : //! @endcond