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