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