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         102 :     std::unique_ptr<OGRFeature> TranslateFeature(OGRFeature *poSrcFeature) const
      65             :     {
      66         102 :         auto poFeature = std::make_unique<OGRFeature>(m_poFeatureDefn);
      67         102 :         poFeature->SetFID(poSrcFeature->GetFID());
      68         102 :         const auto styleString = poSrcFeature->GetStyleString();
      69         102 :         if (styleString)
      70           0 :             poFeature->SetStyleString(styleString);
      71         102 :         poFeature->SetFieldsFrom(
      72             :             poSrcFeature, m_anMapSrcFieldsToDstFields.data(), false, false);
      73         102 :         int iDstGeomField = 0;
      74         184 :         for (int nSrcGeomField : m_anMapDstGeomFieldsToSrcGeomFields)
      75             :         {
      76          82 :             poFeature->SetGeomFieldDirectly(
      77             :                 iDstGeomField, poSrcFeature->StealGeometry(nSrcGeomField));
      78          82 :             ++iDstGeomField;
      79             :         }
      80         102 :         return poFeature;
      81             :     }
      82             : 
      83         101 :     void TranslateFeature(
      84             :         std::unique_ptr<OGRFeature> poSrcFeature,
      85             :         std::vector<std::unique_ptr<OGRFeature>> &apoOutFeatures) override
      86             :     {
      87         101 :         apoOutFeatures.push_back(TranslateFeature(poSrcFeature.release()));
      88         101 :     }
      89             : 
      90             :   public:
      91           9 :     explicit GDALVectorSelectAlgorithmLayer(OGRLayer &oSrcLayer)
      92           9 :         : GDALVectorPipelineOutputLayer(oSrcLayer),
      93           9 :           m_poFeatureDefn(new OGRFeatureDefn(oSrcLayer.GetName()))
      94             :     {
      95           9 :         SetDescription(oSrcLayer.GetDescription());
      96           9 :         SetMetadata(oSrcLayer.GetMetadata());
      97           9 :         m_poFeatureDefn->SetGeomType(wkbNone);
      98           9 :         m_poFeatureDefn->Reference();
      99           9 :     }
     100             : 
     101          18 :     ~GDALVectorSelectAlgorithmLayer() override
     102           9 :     {
     103           9 :         m_poFeatureDefn->Release();
     104          18 :     }
     105             : 
     106           5 :     bool IncludeFields(const std::vector<std::string> &selectedFields,
     107             :                        bool bStrict)
     108             :     {
     109          10 :         std::set<std::string> oSetSelFields;
     110          10 :         std::set<std::string> oSetSelFieldsUC;
     111          14 :         for (const std::string &osFieldName : selectedFields)
     112             :         {
     113           9 :             oSetSelFields.insert(osFieldName);
     114           9 :             oSetSelFieldsUC.insert(CPLString(osFieldName).toupper());
     115             :         }
     116             : 
     117          10 :         std::set<std::string> oSetUsedSetFieldsUC;
     118             : 
     119           5 :         const auto poSrcLayerDefn = m_srcLayer.GetLayerDefn();
     120          17 :         for (int i = 0; i < poSrcLayerDefn->GetFieldCount(); ++i)
     121             :         {
     122          12 :             const auto poSrcFieldDefn = poSrcLayerDefn->GetFieldDefn(i);
     123             :             auto oIter = oSetSelFieldsUC.find(
     124          12 :                 CPLString(poSrcFieldDefn->GetNameRef()).toupper());
     125          12 :             if (oIter != oSetSelFieldsUC.end())
     126             :             {
     127           8 :                 m_anMapSrcFieldsToDstFields.push_back(
     128           4 :                     m_poFeatureDefn->GetFieldCount());
     129           8 :                 OGRFieldDefn oDstFieldDefn(*poSrcFieldDefn);
     130           4 :                 m_poFeatureDefn->AddFieldDefn(&oDstFieldDefn);
     131           4 :                 oSetUsedSetFieldsUC.insert(*oIter);
     132             :             }
     133             :             else
     134             :             {
     135           8 :                 m_anMapSrcFieldsToDstFields.push_back(-1);
     136             :             }
     137             :         }
     138             : 
     139          11 :         for (int i = 0; i < poSrcLayerDefn->GetGeomFieldCount(); ++i)
     140             :         {
     141           6 :             const auto poSrcFieldDefn = poSrcLayerDefn->GetGeomFieldDefn(i);
     142             :             auto oIter = oSetSelFieldsUC.find(
     143           6 :                 CPLString(poSrcFieldDefn->GetNameRef()).toupper());
     144           6 :             if (oIter != oSetSelFieldsUC.end())
     145             :             {
     146           1 :                 m_anMapDstGeomFieldsToSrcGeomFields.push_back(i);
     147           2 :                 OGRGeomFieldDefn oDstFieldDefn(*poSrcFieldDefn);
     148           1 :                 m_poFeatureDefn->AddGeomFieldDefn(&oDstFieldDefn);
     149           1 :                 oSetUsedSetFieldsUC.insert(*oIter);
     150             :             }
     151             :         }
     152             : 
     153             :         auto oIter = oSetSelFieldsUC.find(
     154           5 :             CPLString(OGR_GEOMETRY_DEFAULT_NON_EMPTY_NAME).toupper());
     155           9 :         if (m_poFeatureDefn->GetGeomFieldCount() == 0 &&
     156           9 :             oIter != oSetSelFieldsUC.end() &&
     157           2 :             poSrcLayerDefn->GetGeomFieldCount() == 1)
     158             :         {
     159           2 :             const auto poSrcFieldDefn = poSrcLayerDefn->GetGeomFieldDefn(0);
     160           2 :             m_anMapDstGeomFieldsToSrcGeomFields.push_back(0);
     161           4 :             OGRGeomFieldDefn oDstFieldDefn(*poSrcFieldDefn);
     162           2 :             m_poFeatureDefn->AddGeomFieldDefn(&oDstFieldDefn);
     163           2 :             oSetUsedSetFieldsUC.insert(*oIter);
     164             :         }
     165             : 
     166           5 :         if (oSetUsedSetFieldsUC.size() != oSetSelFields.size())
     167             :         {
     168           6 :             for (const std::string &osName : oSetSelFields)
     169             :             {
     170          15 :                 if (!cpl::contains(oSetUsedSetFieldsUC,
     171           5 :                                    CPLString(osName).toupper()))
     172             :                 {
     173           4 :                     CPLError(bStrict ? CE_Failure : CE_Warning, CPLE_AppDefined,
     174             :                              "Field '%s' does not exist in layer '%s'.%s",
     175           2 :                              osName.c_str(), m_srcLayer.GetDescription(),
     176             :                              bStrict ? " You may specify "
     177             :                                        "--ignore-missing-fields to skip it"
     178             :                                      : " It will be ignored");
     179           2 :                     if (bStrict)
     180           1 :                         return false;
     181             :                 }
     182             :             }
     183             :         }
     184             : 
     185           4 :         return true;
     186             :     }
     187             : 
     188           4 :     void ExcludeFields(const std::vector<std::string> &fields)
     189             :     {
     190           8 :         std::set<std::string> oSetSelFields;
     191           8 :         std::set<std::string> oSetSelFieldsUC;
     192           9 :         for (const std::string &osFieldName : fields)
     193             :         {
     194           5 :             oSetSelFields.insert(osFieldName);
     195           5 :             oSetSelFieldsUC.insert(CPLString(osFieldName).toupper());
     196             :         }
     197             : 
     198           4 :         const auto poSrcLayerDefn = m_srcLayer.GetLayerDefn();
     199          16 :         for (int i = 0; i < poSrcLayerDefn->GetFieldCount(); ++i)
     200             :         {
     201          12 :             const auto poSrcFieldDefn = poSrcLayerDefn->GetFieldDefn(i);
     202             :             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 (int i = 0; i < poSrcLayerDefn->GetGeomFieldCount(); ++i)
     227             :             {
     228           3 :                 const auto poSrcFieldDefn = poSrcLayerDefn->GetGeomFieldDefn(i);
     229             :                 auto oIter = oSetSelFieldsUC.find(
     230           3 :                     CPLString(poSrcFieldDefn->GetNameRef()).toupper());
     231           3 :                 if (oIter == oSetSelFieldsUC.end())
     232             :                 {
     233           2 :                     m_anMapDstGeomFieldsToSrcGeomFields.push_back(i);
     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(poSrcFeature.get()).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 :     CPLAssert(m_inputDataset.GetDatasetRef());
     292           9 :     CPLAssert(m_outputDataset.GetName().empty());
     293           9 :     CPLAssert(!m_outputDataset.GetDatasetRef());
     294             : 
     295           9 :     auto poSrcDS = m_inputDataset.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