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