Line data    Source code

       1             : /******************************************************************************
       2             :  *
       3             :  * Project:  GDAL
       4             :  * Purpose:  "select" step of "vector pipeline"
       5             :  * Author:   Even Rouault <even dot rouault at spatialys.com>
       6             :  *
       7             :  ******************************************************************************
       8             :  * Copyright (c) 2024, Even Rouault <even dot rouault at spatialys.com>
       9             :  *
      10             :  * SPDX-License-Identifier: MIT
      11             :  ****************************************************************************/
      12             : 
      13             : #include "gdalalg_vector_select.h"
      14             : 
      15             : #include "gdal_priv.h"
      16             : #include "ogrsf_frmts.h"
      17             : #include "ogr_p.h"
      18             : 
      19             : #include <set>
      20             : 
      21             : //! @cond Doxygen_Suppress
      22             : 
      23             : #ifndef _
      24             : #define _(x) (x)
      25             : #endif
      26             : 
      27             : /************************************************************************/
      28             : /*        GDALVectorSelectAlgorithm::GDALVectorSelectAlgorithm()        */
      29             : /************************************************************************/
      30             : 
      31          51 : GDALVectorSelectAlgorithm::GDALVectorSelectAlgorithm(bool standaloneStep)
      32             :     : GDALVectorPipelineStepAlgorithm(NAME, DESCRIPTION, HELP_URL,
      33          51 :                                       standaloneStep)
      34             : {
      35          51 :     if (!standaloneStep)
      36             :     {
      37          35 :         AddOutputLayerNameArg(/* hiddenForCLI = */ false,
      38             :                               /* shortNameOutputLayerAllowed = */ false);
      39             :     }
      40          51 :     AddActiveLayerArg(&m_activeLayer);
      41             :     AddArg("fields", 0, _("Fields to select (or exclude if --exclude)"),
      42         102 :            &m_fields)
      43          51 :         .SetPositional()
      44          51 :         .SetRequired();
      45         102 :     AddArg("exclude", 0, _("Exclude specified fields"), &m_exclude)
      46          51 :         .SetMutualExclusionGroup("exclude-ignore");
      47             :     AddArg("ignore-missing-fields", 0, _("Ignore missing fields"),
      48         102 :            &m_ignoreMissingFields)
      49          51 :         .SetMutualExclusionGroup("exclude-ignore");
      50          51 : }
      51             : 
      52             : namespace
      53             : {
      54             : 
      55             : /************************************************************************/
      56             : /*                    GDALVectorSelectAlgorithmLayer                    */
      57             : /************************************************************************/
      58             : 
      59             : class GDALVectorSelectAlgorithmLayer final
      60             :     : public GDALVectorPipelineOutputLayer
      61             : {
      62             :   private:
      63             :     OGRFeatureDefn *const m_poFeatureDefn = nullptr;
      64             :     std::vector<int> m_anMapSrcFieldsToDstFields{};
      65             :     std::vector<int> m_anMapDstGeomFieldsToSrcGeomFields{};
      66             : 
      67             :     CPL_DISALLOW_COPY_ASSIGN(GDALVectorSelectAlgorithmLayer)
      68             : 
      69             :     std::unique_ptr<OGRFeature>
      70         102 :     TranslateFeature(std::unique_ptr<OGRFeature> poSrcFeature) const
      71             :     {
      72         102 :         auto poFeature = std::make_unique<OGRFeature>(m_poFeatureDefn);
      73         102 :         poFeature->SetFID(poSrcFeature->GetFID());
      74         102 :         const auto styleString = poSrcFeature->GetStyleString();
      75         102 :         if (styleString)
      76           0 :             poFeature->SetStyleString(styleString);
      77         102 :         poFeature->SetFieldsFrom(poSrcFeature.get(),
      78             :                                  m_anMapSrcFieldsToDstFields.data(), false,
      79             :                                  false);
      80         102 :         int iDstGeomField = 0;
      81         184 :         for (int nSrcGeomField : m_anMapDstGeomFieldsToSrcGeomFields)
      82             :         {
      83          82 :             poFeature->SetGeomFieldDirectly(
      84             :                 iDstGeomField, poSrcFeature->StealGeometry(nSrcGeomField));
      85          82 :             ++iDstGeomField;
      86             :         }
      87         102 :         return poFeature;
      88             :     }
      89             : 
      90         101 :     void TranslateFeature(
      91             :         std::unique_ptr<OGRFeature> poSrcFeature,
      92             :         std::vector<std::unique_ptr<OGRFeature>> &apoOutFeatures) override
      93             :     {
      94         101 :         apoOutFeatures.push_back(TranslateFeature(std::move(poSrcFeature)));
      95         101 :     }
      96             : 
      97             :   public:
      98          10 :     explicit GDALVectorSelectAlgorithmLayer(
      99             :         OGRLayer &oSrcLayer, const std::string &osOutputLayerName)
     100          10 :         : GDALVectorPipelineOutputLayer(oSrcLayer),
     101          10 :           m_poFeatureDefn(new OGRFeatureDefn(osOutputLayerName.empty()
     102          10 :                                                  ? oSrcLayer.GetName()
     103          20 :                                                  : osOutputLayerName.c_str()))
     104             :     {
     105          10 :         SetDescription(m_poFeatureDefn->GetName());
     106          10 :         SetMetadata(oSrcLayer.GetMetadata());
     107          10 :         m_poFeatureDefn->SetGeomType(wkbNone);
     108          10 :         m_poFeatureDefn->Reference();
     109          10 :     }
     110             : 
     111          20 :     ~GDALVectorSelectAlgorithmLayer() override
     112          10 :     {
     113          10 :         m_poFeatureDefn->Release();
     114          20 :     }
     115             : 
     116           6 :     bool IncludeFields(const std::vector<std::string> &selectedFields,
     117             :                        bool bStrict)
     118             :     {
     119          12 :         std::set<std::string> oSetSelFields;
     120          12 :         std::set<std::string> oSetSelFieldsUC;
     121          16 :         for (const std::string &osFieldName : selectedFields)
     122             :         {
     123          10 :             oSetSelFields.insert(osFieldName);
     124          10 :             oSetSelFieldsUC.insert(CPLString(osFieldName).toupper());
     125             :         }
     126             : 
     127          12 :         std::set<std::string> oSetUsedSetFieldsUC;
     128             : 
     129           6 :         const auto poSrcLayerDefn = m_srcLayer.GetLayerDefn();
     130          21 :         for (const auto poSrcFieldDefn : poSrcLayerDefn->GetFields())
     131             :         {
     132             :             const auto oIter = oSetSelFieldsUC.find(
     133          15 :                 CPLString(poSrcFieldDefn->GetNameRef()).toupper());
     134          15 :             if (oIter != oSetSelFieldsUC.end())
     135             :             {
     136          10 :                 m_anMapSrcFieldsToDstFields.push_back(
     137           5 :                     m_poFeatureDefn->GetFieldCount());
     138          10 :                 OGRFieldDefn oDstFieldDefn(*poSrcFieldDefn);
     139           5 :                 m_poFeatureDefn->AddFieldDefn(&oDstFieldDefn);
     140           5 :                 oSetUsedSetFieldsUC.insert(*oIter);
     141             :             }
     142             :             else
     143             :             {
     144          10 :                 m_anMapSrcFieldsToDstFields.push_back(-1);
     145             :             }
     146             :         }
     147             : 
     148          13 :         for (const auto poSrcFieldDefn : poSrcLayerDefn->GetGeomFields())
     149             :         {
     150             :             const auto oIter = oSetSelFieldsUC.find(
     151           7 :                 CPLString(poSrcFieldDefn->GetNameRef()).toupper());
     152           7 :             if (oIter != oSetSelFieldsUC.end())
     153             :             {
     154           2 :                 m_anMapDstGeomFieldsToSrcGeomFields.push_back(
     155           1 :                     m_poFeatureDefn->GetGeomFieldCount());
     156           2 :                 OGRGeomFieldDefn oDstFieldDefn(*poSrcFieldDefn);
     157           1 :                 m_poFeatureDefn->AddGeomFieldDefn(&oDstFieldDefn);
     158           1 :                 oSetUsedSetFieldsUC.insert(*oIter);
     159             :             }
     160             :         }
     161             : 
     162             :         const auto oIter = oSetSelFieldsUC.find(
     163           6 :             CPLString(OGR_GEOMETRY_DEFAULT_NON_EMPTY_NAME).toupper());
     164          11 :         if (m_poFeatureDefn->GetGeomFieldCount() == 0 &&
     165          11 :             oIter != oSetSelFieldsUC.end() &&
     166           2 :             poSrcLayerDefn->GetGeomFieldCount() == 1)
     167             :         {
     168           2 :             const auto poSrcFieldDefn = poSrcLayerDefn->GetGeomFieldDefn(0);
     169           2 :             m_anMapDstGeomFieldsToSrcGeomFields.push_back(0);
     170           4 :             OGRGeomFieldDefn oDstFieldDefn(*poSrcFieldDefn);
     171           2 :             m_poFeatureDefn->AddGeomFieldDefn(&oDstFieldDefn);
     172           2 :             oSetUsedSetFieldsUC.insert(*oIter);
     173             :         }
     174             : 
     175           6 :         if (oSetUsedSetFieldsUC.size() != oSetSelFields.size())
     176             :         {
     177           6 :             for (const std::string &osName : oSetSelFields)
     178             :             {
     179          15 :                 if (!cpl::contains(oSetUsedSetFieldsUC,
     180           5 :                                    CPLString(osName).toupper()))
     181             :                 {
     182           4 :                     CPLError(bStrict ? CE_Failure : CE_Warning, CPLE_AppDefined,
     183             :                              "Field '%s' does not exist in layer '%s'.%s",
     184           2 :                              osName.c_str(), m_srcLayer.GetDescription(),
     185             :                              bStrict ? " You may specify "
     186             :                                        "--ignore-missing-fields to skip it"
     187             :                                      : " It will be ignored");
     188           2 :                     if (bStrict)
     189           1 :                         return false;
     190             :                 }
     191             :             }
     192             :         }
     193             : 
     194           5 :         return true;
     195             :     }
     196             : 
     197           4 :     void ExcludeFields(const std::vector<std::string> &fields)
     198             :     {
     199           8 :         std::set<std::string> oSetSelFields;
     200           8 :         std::set<std::string> oSetSelFieldsUC;
     201           9 :         for (const std::string &osFieldName : fields)
     202             :         {
     203           5 :             oSetSelFields.insert(osFieldName);
     204           5 :             oSetSelFieldsUC.insert(CPLString(osFieldName).toupper());
     205             :         }
     206             : 
     207           4 :         const auto poSrcLayerDefn = m_srcLayer.GetLayerDefn();
     208          16 :         for (const auto poSrcFieldDefn : poSrcLayerDefn->GetFields())
     209             :         {
     210             :             const auto oIter = oSetSelFieldsUC.find(
     211          12 :                 CPLString(poSrcFieldDefn->GetNameRef()).toupper());
     212          12 :             if (oIter != oSetSelFieldsUC.end())
     213             :             {
     214           1 :                 m_anMapSrcFieldsToDstFields.push_back(-1);
     215             :             }
     216             :             else
     217             :             {
     218          22 :                 m_anMapSrcFieldsToDstFields.push_back(
     219          11 :                     m_poFeatureDefn->GetFieldCount());
     220          22 :                 OGRFieldDefn oDstFieldDefn(*poSrcFieldDefn);
     221          11 :                 m_poFeatureDefn->AddFieldDefn(&oDstFieldDefn);
     222             :             }
     223             :         }
     224             : 
     225           0 :         if (oSetSelFieldsUC.find(
     226           8 :                 CPLString(OGR_GEOMETRY_DEFAULT_NON_EMPTY_NAME).toupper()) !=
     227          13 :                 oSetSelFieldsUC.end() &&
     228           1 :             poSrcLayerDefn->GetGeomFieldCount() == 1)
     229             :         {
     230             :             // exclude default geometry field
     231             :         }
     232             :         else
     233             :         {
     234           6 :             for (const auto poSrcFieldDefn : poSrcLayerDefn->GetGeomFields())
     235             :             {
     236             :                 const auto oIter = oSetSelFieldsUC.find(
     237           3 :                     CPLString(poSrcFieldDefn->GetNameRef()).toupper());
     238           3 :                 if (oIter == oSetSelFieldsUC.end())
     239             :                 {
     240           4 :                     m_anMapDstGeomFieldsToSrcGeomFields.push_back(
     241           2 :                         m_poFeatureDefn->GetGeomFieldCount());
     242           4 :                     OGRGeomFieldDefn oDstFieldDefn(*poSrcFieldDefn);
     243           2 :                     m_poFeatureDefn->AddGeomFieldDefn(&oDstFieldDefn);
     244             :                 }
     245             :             }
     246             :         }
     247           4 :     }
     248             : 
     249          51 :     const OGRFeatureDefn *GetLayerDefn() const override
     250             :     {
     251          51 :         return m_poFeatureDefn;
     252             :     }
     253             : 
     254           3 :     GIntBig GetFeatureCount(int bForce) override
     255             :     {
     256           3 :         if (!m_poAttrQuery && !m_poFilterGeom)
     257           1 :             return m_srcLayer.GetFeatureCount(bForce);
     258           2 :         return OGRLayer::GetFeatureCount(bForce);
     259             :     }
     260             : 
     261           2 :     OGRErr IGetExtent(int iGeomField, OGREnvelope *psExtent,
     262             :                       bool bForce) override
     263             :     {
     264           2 :         return m_srcLayer.GetExtent(iGeomField, psExtent, bForce);
     265             :     }
     266             : 
     267           2 :     OGRFeature *GetFeature(GIntBig nFID) override
     268             :     {
     269             :         auto poSrcFeature =
     270           4 :             std::unique_ptr<OGRFeature>(m_srcLayer.GetFeature(nFID));
     271           2 :         if (!poSrcFeature)
     272           1 :             return nullptr;
     273           1 :         return TranslateFeature(std::move(poSrcFeature)).release();
     274             :     }
     275             : 
     276          12 :     int TestCapability(const char *pszCap) const override
     277             :     {
     278          12 :         if (EQUAL(pszCap, OLCRandomRead) || EQUAL(pszCap, OLCCurveGeometries) ||
     279          12 :             EQUAL(pszCap, OLCMeasuredGeometries) ||
     280          12 :             EQUAL(pszCap, OLCZGeometries) ||
     281          12 :             (EQUAL(pszCap, OLCFastFeatureCount) && !m_poAttrQuery &&
     282           2 :              !m_poFilterGeom) ||
     283          11 :             EQUAL(pszCap, OLCFastGetExtent) || EQUAL(pszCap, OLCStringsAsUTF8))
     284             :         {
     285           2 :             return m_srcLayer.TestCapability(pszCap);
     286             :         }
     287          10 :         return false;
     288             :     }
     289             : };
     290             : 
     291             : }  // namespace
     292             : 
     293             : /************************************************************************/
     294             : /*                 GDALVectorSelectAlgorithm::RunStep()                 */
     295             : /************************************************************************/
     296             : 
     297          10 : bool GDALVectorSelectAlgorithm::RunStep(GDALPipelineStepRunContext &)
     298             : {
     299          10 :     auto poSrcDS = m_inputDataset[0].GetDatasetRef();
     300          10 :     CPLAssert(poSrcDS);
     301             : 
     302          10 :     CPLAssert(m_outputDataset.GetName().empty());
     303          10 :     CPLAssert(!m_outputDataset.GetDatasetRef());
     304             : 
     305          20 :     auto outDS = std::make_unique<GDALVectorPipelineOutputDataset>(*poSrcDS);
     306             : 
     307          20 :     for (auto &&poSrcLayer : poSrcDS->GetLayers())
     308             :     {
     309          13 :         if (m_activeLayer.empty() ||
     310           2 :             m_activeLayer == poSrcLayer->GetDescription())
     311             :         {
     312             :             auto poLayer = std::make_unique<GDALVectorSelectAlgorithmLayer>(
     313          10 :                 *poSrcLayer, m_outputLayerName);
     314          10 :             if (m_exclude)
     315             :             {
     316           4 :                 poLayer->ExcludeFields(m_fields);
     317             :             }
     318             :             else
     319             :             {
     320           6 :                 if (!poLayer->IncludeFields(m_fields, !m_ignoreMissingFields))
     321           1 :                     return false;
     322             :             }
     323           9 :             outDS->AddLayer(*poSrcLayer, std::move(poLayer));
     324             :         }
     325             :         else
     326             :         {
     327           2 :             outDS->AddLayer(
     328             :                 *poSrcLayer,
     329           2 :                 std::make_unique<GDALVectorPipelinePassthroughLayer>(
     330             :                     *poSrcLayer));
     331             :         }
     332             :     }
     333             : 
     334           9 :     m_outputDataset.Set(std::move(outDS));
     335             : 
     336           9 :     return true;
     337             : }
     338             : 
     339             : GDALVectorSelectAlgorithmStandalone::~GDALVectorSelectAlgorithmStandalone() =
     340             :     default;
     341             : 
     342             : //! @endcond