File: gdalalg_vector_limit.cpp

package info (click to toggle)
gdal 3.12.1%2Bdfsg-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 92,292 kB
  • sloc: cpp: 1,223,498; ansic: 206,434; python: 26,278; java: 6,001; xml: 4,769; sh: 3,855; cs: 2,513; yacc: 1,306; makefile: 214
file content (137 lines) | stat: -rw-r--r-- 3,733 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
/******************************************************************************
 *
 * Project:  GDAL
 * Purpose:  "limit" step of "vector pipeline"
 * Author:   Dan Baston
 *
 ******************************************************************************
 * Copyright (c) 2025, ISciences LLC
 *
 * SPDX-License-Identifier: MIT
 ****************************************************************************/

#include "gdalalg_vector_limit.h"
#include "gdalalg_vector_pipeline.h"

#include "gdal_priv.h"
#include "ogrsf_frmts.h"
#include "ogr_p.h"

#include <algorithm>
#include <set>

//! @cond Doxygen_Suppress

#ifndef _
#define _(x) (x)
#endif

/************************************************************************/
/*         GDALVectorLimitAlgorithm::GDALVectorLimitAlgorithm()       */
/************************************************************************/

GDALVectorLimitAlgorithm::GDALVectorLimitAlgorithm(bool standaloneStep)
    : GDALVectorPipelineStepAlgorithm(NAME, DESCRIPTION, HELP_URL,
                                      standaloneStep)
{
    AddArg("limit", 0, _("Limit the number of features to read per layer"),
           &m_featureLimit)
        .SetPositional()
        .SetRequired();
    AddActiveLayerArg(&m_activeLayer);
}

namespace
{

/************************************************************************/
/*                      GDALVectorReadLimitedLayer                      */
/************************************************************************/

class GDALVectorReadLimitedLayer : public OGRLayer
{
  public:
    GDALVectorReadLimitedLayer(OGRLayer &layer, int featureLimit)
        : m_srcLayer(layer), m_featureLimit(featureLimit), m_featuresRead(0)
    {
    }

    ~GDALVectorReadLimitedLayer() override;

    OGRFeature *GetNextFeature() override
    {
        if (m_featuresRead < m_featureLimit)
        {
            m_featuresRead++;
            return m_srcLayer.GetNextFeature();
        }
        return nullptr;
    }

    const OGRFeatureDefn *GetLayerDefn() const override
    {
        return m_srcLayer.GetLayerDefn();
    }

    GIntBig GetFeatureCount(int bForce) override
    {
        return std::min(m_featureLimit, m_srcLayer.GetFeatureCount(bForce));
    }

    void ResetReading() override
    {
        m_featuresRead = 0;
        m_srcLayer.ResetReading();
    }

    int TestCapability(const char *pszCap) const override
    {
        return m_srcLayer.TestCapability(pszCap);
    }

  private:
    OGRLayer &m_srcLayer;
    GIntBig m_featureLimit;
    GIntBig m_featuresRead;
};

GDALVectorReadLimitedLayer::~GDALVectorReadLimitedLayer() = default;

}  // namespace

/************************************************************************/
/*               GDALVectorLimitAlgorithm::RunStep()                   */
/************************************************************************/

bool GDALVectorLimitAlgorithm::RunStep(GDALPipelineStepRunContext &)
{
    auto poSrcDS = m_inputDataset[0].GetDatasetRef();
    CPLAssert(poSrcDS);

    CPLAssert(m_outputDataset.GetName().empty());
    CPLAssert(!m_outputDataset.GetDatasetRef());

    auto outDS = std::make_unique<GDALVectorOutputDataset>();

    for (auto &&poSrcLayer : poSrcDS->GetLayers())
    {
        if (m_activeLayer.empty() ||
            m_activeLayer == poSrcLayer->GetDescription())
        {
            outDS->AddLayer(std::make_unique<GDALVectorReadLimitedLayer>(
                *poSrcLayer, m_featureLimit));
        }
        else
        {
            outDS->AddLayer(
                std::make_unique<GDALVectorPipelinePassthroughLayer>(
                    *poSrcLayer));
        }
    }

    m_outputDataset.Set(std::move(outDS));

    return true;
}

//! @endcond