Merge pull request #13414 from rouault/fix_13401

ogr2ogr/VectorTranslate: fix (and improve) selectFields support in Arrow code path

Author: rouault

apps/ogr2ogr_lib.cpp

@@ -520,12 +520,6 @@ struct AssociatedLayers
 
 class SetupTargetLayer
 {
-    bool CanUseWriteArrowBatch(OGRLayer *poSrcLayer, OGRLayer *poDstLayer,
-                               bool bJustCreatedLayer,
-                               const GDALVectorTranslateOptions *psOptions,
-                               bool bPreserveFID, bool &bError,
-                               OGRArrowArrayStream &streamSrc);
-
   public:
     GDALDataset *m_poSrcDS = nullptr;
     GDALDataset *m_poDstDS = nullptr;
@@ -565,6 +559,15 @@ class SetupTargetLayer
     std::unique_ptr<TargetLayerInfo>
     Setup(OGRLayer *poSrcLayer, const char *pszNewLayerName,
           GDALVectorTranslateOptions *psOptions, GIntBig &nTotalEventsDone);
+
+  private:
+    bool CanUseWriteArrowBatch(OGRLayer *poSrcLayer, OGRLayer *poDstLayer,
+                               bool bJustCreatedLayer,
+                               const GDALVectorTranslateOptions *psOptions,
+                               bool bPreserveFID, bool &bError,
+                               OGRArrowArrayStream &streamSrc);
+
+    void SetIgnoredFields(OGRLayer *poSrcLayer);
 };
 
 class LayerTranslator
@@ -4253,8 +4256,8 @@ bool SetupTargetLayer::CanUseWriteArrowBatch(
         !psOptions->bUpsert && !psOptions->bSkipFailures &&
         !psOptions->poClipSrc && !psOptions->poClipDst &&
         psOptions->asGCPs.empty() && !psOptions->bWrapDateline &&
-        !m_papszSelFields && !m_bAddMissingFields &&
-        m_eGType == GEOMTYPE_UNCHANGED && psOptions->eGeomOp == GEOMOP_NONE &&
+        !m_bAddMissingFields && m_eGType == GEOMTYPE_UNCHANGED &&
+        psOptions->eGeomOp == GEOMOP_NONE &&
         m_eGeomTypeConversion == GTC_DEFAULT && m_nCoordDim < 0 &&
         !m_papszFieldTypesToString && !m_papszMapFieldType &&
         !m_bUnsetFieldWidth && !m_bExplodeCollections && !m_pszZField &&
@@ -4282,6 +4285,11 @@ bool SetupTargetLayer::CanUseWriteArrowBatch(
             }
         }
 
+        if (m_bSelFieldsSet)
+        {
+            SetIgnoredFields(poSrcLayer);
+        }
+
         const CPLStringList aosGetArrowStreamOptions(BuildGetArrowStreamOptions(
             poSrcLayer, poDstLayer, psOptions, bPreserveFID));
         if (poSrcLayer->GetArrowStream(streamSrc.get(),
@@ -4428,6 +4436,55 @@ bool SetupTargetLayer::CanUseWriteArrowBatch(
     return bUseWriteArrowBatch;
 }
 
+/************************************************************************/
+/*              SetupTargetLayer::SetIgnoredFields()                    */
+/************************************************************************/
+
+void SetupTargetLayer::SetIgnoredFields(OGRLayer *poSrcLayer)
+{
+    bool bUseIgnoredFields = true;
+    CPLStringList aosWHEREUsedFields;
+    const auto poSrcFDefn = poSrcLayer->GetLayerDefn();
+
+    if (m_pszWHERE)
+    {
+        /* We must not ignore fields used in the -where expression
+         * (#4015) */
+        OGRFeatureQuery oFeatureQuery;
+        if (oFeatureQuery.Compile(poSrcFDefn, m_pszWHERE, FALSE, nullptr) ==
+            OGRERR_NONE)
+        {
+            aosWHEREUsedFields = oFeatureQuery.GetUsedFields();
+        }
+        else
+        {
+            bUseIgnoredFields = false;
+        }
+    }
+
+    CPLStringList aosIgnoredFields;
+    for (int iSrcField = 0;
+         bUseIgnoredFields && iSrcField < poSrcFDefn->GetFieldCount();
+         iSrcField++)
+    {
+        const char *pszFieldName =
+            poSrcFDefn->GetFieldDefn(iSrcField)->GetNameRef();
+        bool bFieldRequested =
+            CSLFindString(m_papszSelFields, pszFieldName) >= 0;
+        bFieldRequested |= aosWHEREUsedFields.FindString(pszFieldName) >= 0;
+        bFieldRequested |=
+            (m_pszZField != nullptr && EQUAL(pszFieldName, m_pszZField));
+
+        // If the source field is not requested, add it to the list of ignored
+        // fields.
+        if (!bFieldRequested)
+            aosIgnoredFields.push_back(pszFieldName);
+    }
+    if (bUseIgnoredFields)
+        poSrcLayer->SetIgnoredFields(
+            const_cast<const char **>(aosIgnoredFields.List()));
+}
+
 /************************************************************************/
 /*                   SetupTargetLayer::Setup()                          */
 /************************************************************************/
@@ -5301,58 +5358,9 @@ SetupTargetLayer::Setup(OGRLayer *poSrcLayer, const char *pszNewLayerName,
         /* Use SetIgnoredFields() on source layer if available */
         /* --------------------------------------------------------------------
          */
-        if (poSrcLayer->TestCapability(OLCIgnoreFields))
+        if (m_bSelFieldsSet && poSrcLayer->TestCapability(OLCIgnoreFields))
         {
-            bool bUseIgnoredFields = true;
-            CPLStringList aosWHEREUsedFields;
-
-            if (m_pszWHERE)
-            {
-                /* We must not ignore fields used in the -where expression
-                 * (#4015) */
-                OGRFeatureQuery oFeatureQuery;
-                if (oFeatureQuery.Compile(poSrcLayer->GetLayerDefn(),
-                                          m_pszWHERE, FALSE,
-                                          nullptr) == OGRERR_NONE)
-                {
-                    aosWHEREUsedFields = oFeatureQuery.GetUsedFields();
-                }
-                else
-                {
-                    bUseIgnoredFields = false;
-                }
-            }
-
-            CPLStringList aosIgnoredFields;
-            for (int iSrcField = 0;
-                 bUseIgnoredFields && iSrcField < poSrcFDefn->GetFieldCount();
-                 iSrcField++)
-            {
-                const char *pszFieldName =
-                    poSrcFDefn->GetFieldDefn(iSrcField)->GetNameRef();
-                bool bFieldRequested = false;
-                for (int iField = 0;
-                     m_papszSelFields && m_papszSelFields[iField]; iField++)
-                {
-                    if (EQUAL(pszFieldName, m_papszSelFields[iField]))
-                    {
-                        bFieldRequested = true;
-                        break;
-                    }
-                }
-                bFieldRequested |=
-                    aosWHEREUsedFields.FindString(pszFieldName) >= 0;
-                bFieldRequested |= (m_pszZField != nullptr &&
-                                    EQUAL(pszFieldName, m_pszZField));
-
-                /* If source field not requested, add it to ignored files list
-                 */
-                if (!bFieldRequested)
-                    aosIgnoredFields.push_back(pszFieldName);
-            }
-            if (bUseIgnoredFields)
-                poSrcLayer->SetIgnoredFields(
-                    const_cast<const char **>(aosIgnoredFields.List()));
+            SetIgnoredFields(poSrcLayer);
         }
     }
     else if (!bAppend || m_bAddMissingFields)

autotest/utilities/test_ogr2ogr_lib.py

@@ -166,6 +166,25 @@ def test_ogr2ogr_lib_6():
     assert feat.GetFieldAsString("PRFEDEA") == "35043411"
 
 
+###############################################################################
+# Test selectFields with arrow optimization
+
+
+@pytest.mark.require_driver("GPKG")
+def test_ogr2ogr_lib_selectFields_gpkg(tmp_vsimem):
+
+    srcDS = gdal.VectorTranslate(tmp_vsimem / "in.gpkg", "../ogr/data/poly.shp")
+    gdal.VectorTranslate(
+        tmp_vsimem / "out.gpkg", srcDS, selectFields=["eas_id", "prfedea"]
+    )
+    ds = ogr.Open(tmp_vsimem / "out.gpkg")
+    lyr = ds.GetLayer(0)
+    assert lyr.GetLayerDefn().GetFieldCount() == 2
+    feat = lyr.GetNextFeature()
+    assert feat.GetFieldAsDouble("EAS_ID") == 168
+    assert feat.GetFieldAsString("PRFEDEA") == "35043411"
+
+
 ###############################################################################
 # Test selectFields to []
 
@@ -176,6 +195,24 @@ def test_ogr2ogr_lib_sel_fields_empty():
     ds = gdal.VectorTranslate("", srcDS, format="MEM", selectFields=[])
     lyr = ds.GetLayer(0)
     assert lyr.GetLayerDefn().GetFieldCount() == 0
+    feat = lyr.GetNextFeature()
+    assert feat.GetGeometryRef() is not None
+
+
+###############################################################################
+# Test selectFields to []
+
+
+@pytest.mark.require_driver("GPKG")
+def test_ogr2ogr_lib_sel_fields_empty_with_arow_optimization(tmp_vsimem):
+
+    srcDS = gdal.VectorTranslate(tmp_vsimem / "in.gpkg", "../ogr/data/poly.shp")
+    gdal.VectorTranslate(tmp_vsimem / "out.gpkg", srcDS, selectFields=[])
+    ds = ogr.Open(tmp_vsimem / "out.gpkg")
+    lyr = ds.GetLayer(0)
+    assert lyr.GetLayerDefn().GetFieldCount() == 0
+    feat = lyr.GetNextFeature()
+    assert feat.GetGeometryRef() is not None
 
 
 ###############################################################################