Merge pull request #9570 from rouault/flatgeobuf_getdataset

FlatGeoBuf: implement OGRLayer::GetDataset(), and CreateFieldFromArrowSchema() subfield fix
OSGeo · Mar 27, 2024 · 634264c · 634264c
2 parents c54fd13 + 3baac2b
commit 634264c
Show file tree

Hide file tree

Showing 6 changed files with 183 additions and 30 deletions.
diff --git a/autotest/ogr/ogr_flatgeobuf.py b/autotest/ogr/ogr_flatgeobuf.py
@@ -1217,3 +1217,147 @@ def test_ogr_flatgeobuf_issue_7401():
 
     ogr.GetDriverByName("FlatGeobuf").DeleteDataSource("/vsimem/test.fgb")
     assert not gdal.VSIStatL("/vsimem/test.fgb")
+
+
+###############################################################################
+
+
+@gdaltest.enable_exceptions()
+def test_ogr_flatgeobuf_write_arrow(tmp_vsimem):
+
+    ds = ogr.GetDriverByName("Memory").CreateDataSource("")
+    src_lyr = ds.CreateLayer("src_lyr")
+
+    field_def = ogr.FieldDefn("field_bool", ogr.OFTInteger)
+    field_def.SetSubType(ogr.OFSTBoolean)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_integer", ogr.OFTInteger)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_int16", ogr.OFTInteger)
+    field_def.SetSubType(ogr.OFSTInt16)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_integer64", ogr.OFTInteger64)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_float32", ogr.OFTReal)
+    field_def.SetSubType(ogr.OFSTFloat32)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_real", ogr.OFTReal)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_string", ogr.OFTString)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_binary", ogr.OFTBinary)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_date", ogr.OFTDate)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_time", ogr.OFTTime)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_datetime", ogr.OFTDateTime)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_boollist", ogr.OFTIntegerList)
+    field_def.SetSubType(ogr.OFSTBoolean)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_integerlist", ogr.OFTIntegerList)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_int16list", ogr.OFTIntegerList)
+    field_def.SetSubType(ogr.OFSTInt16)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_integer64list", ogr.OFTInteger64List)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_float32list", ogr.OFTRealList)
+    field_def.SetSubType(ogr.OFSTFloat32)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_reallist", ogr.OFTRealList)
+    src_lyr.CreateField(field_def)
+
+    field_def = ogr.FieldDefn("field_stringlist", ogr.OFTStringList)
+    src_lyr.CreateField(field_def)
+
+    src_feature = ogr.Feature(src_lyr.GetLayerDefn())
+    src_feature.SetField("field_bool", True)
+    src_feature.SetField("field_integer", 17)
+    src_feature.SetField("field_int16", -17)
+    src_feature.SetField("field_integer64", 9876543210)
+    src_feature.SetField("field_float32", 1.5)
+    src_feature.SetField("field_real", 18.4)
+    src_feature.SetField("field_string", "abc def")
+    src_feature.SetFieldBinary("field_binary", b"\x00\x01")
+    src_feature.SetField("field_binary", b"\x01\x23\x46\x57\x89\xAB\xCD\xEF")
+    src_feature.SetField("field_date", "2011/11/11")
+    src_feature.SetField("field_time", "14:10:35")
+    src_feature.SetField("field_datetime", 2011, 11, 11, 14, 10, 35.123, 0)
+    src_feature.field_boollist = [False, True]
+    src_feature.field_integerlist = [10, 20, 30]
+    src_feature.field_int16list = [10, -20, 30]
+    src_feature.field_integer64list = [9876543210]
+    src_feature.field_float32list = [1.5, -1.5]
+    src_feature.field_reallist = [123.5, 567.0]
+    src_feature.field_stringlist = ["abc", "def"]
+    src_feature.SetGeometry(ogr.CreateGeometryFromWkt("POINT (1 2)"))
+
+    src_lyr.CreateFeature(src_feature)
+
+    filename = str(tmp_vsimem / "temp.fgb")
+    dst_ds = ogr.GetDriverByName("FlatGeoBuf").CreateDataSource(filename)
+    dst_lyr = dst_ds.CreateLayer("dst_lyr")
+
+    stream = src_lyr.GetArrowStream(["INCLUDE_FID=NO"])
+    schema = stream.GetSchema()
+
+    success, error_msg = dst_lyr.IsArrowSchemaSupported(schema)
+    assert success, error_msg
+
+    for i in range(schema.GetChildrenCount()):
+        if schema.GetChild(i).GetName() != "wkb_geometry":
+            dst_lyr.CreateFieldFromArrowSchema(schema.GetChild(i))
+
+    while True:
+        array = stream.GetNextRecordBatch()
+        if array is None:
+            break
+        assert dst_lyr.WriteArrowBatch(schema, array) == ogr.OGRERR_NONE
+
+    dst_ds.Close()
+    dst_ds = ogr.Open(filename)
+    dst_lyr = dst_ds.GetLayer(0)
+    dst_feature = dst_lyr.GetNextFeature()
+    assert (
+        str(dst_feature)
+        == """OGRFeature(dst_lyr):0
+  field_bool (Integer(Boolean)) = 1
+  field_integer (Integer) = 17
+  field_int16 (Integer(Int16)) = -17
+  field_integer64 (Integer64) = 9876543210
+  field_float32 (Real(Float32)) = 1.5
+  field_real (Real) = 18.4
+  field_string (String) = abc def
+  field_binary (Binary) = 0123465789ABCDEF
+  field_date (DateTime) = 2011/11/11 00:00:00
+  field_time (String) = 14:10:35
+  field_datetime (DateTime) = 2011/11/11 14:10:35.123
+  field_boollist (String) = [ 0, 1 ]
+  field_integerlist (String) = [ 10, 20, 30 ]
+  field_int16list (String) = [ 10, -20, 30 ]
+  field_integer64list (String) = [ 9876543210 ]
+  field_float32list (String) = [ 1.5, -1.5 ]
+  field_reallist (String) = [ 123.5, 567.0 ]
+  field_stringlist (String) = [ "abc", "def" ]
+  POINT (1 2)
+
+"""
+    )
diff --git a/ogr/ogrsf_frmts/flatgeobuf/ogr_flatgeobuf.h b/ogr/ogrsf_frmts/flatgeobuf/ogr_flatgeobuf.h
@@ -103,6 +103,7 @@ class OGRFlatGeobufLayer final : public OGRLayer,
     bool m_ignoreAttributeFilter = false;
 
     // creation
+    GDALDataset *m_poDS = nullptr;  // parent dataset to get metadata from it
     bool m_create = false;
     std::deque<FeatureItem> m_featureItems;  // feature item description used to
                                              // create spatial index
@@ -142,7 +143,8 @@ class OGRFlatGeobufLayer final : public OGRLayer,
     OGRFlatGeobufLayer(const FlatGeobuf::Header *, GByte *headerBuf,
                        const char *pszFilename, VSILFILE *poFp,
                        uint64_t offset);
-    OGRFlatGeobufLayer(const char *pszLayerName, const char *pszFilename,
+    OGRFlatGeobufLayer(GDALDataset *poDS, const char *pszLayerName,
+                       const char *pszFilename,
                        const OGRSpatialReference *poSpatialRef,
                        OGRwkbGeometryType eGType,
                        bool bCreateSpatialIndexAtClose, VSILFILE *poFpWrite,
@@ -163,7 +165,7 @@ class OGRFlatGeobufLayer final : public OGRLayer,
     static OGRFlatGeobufLayer *Open(const char *pszFilename, VSILFILE *fp,
                                     bool bVerifyBuffers);
     static OGRFlatGeobufLayer *
-    Create(const char *pszLayerName, const char *pszFilename,
+    Create(GDALDataset *poDS, const char *pszLayerName, const char *pszFilename,
            const OGRSpatialReference *poSpatialRef, OGRwkbGeometryType eGType,
            bool bCreateSpatialIndexAtClose, char **papszOptions);
 
@@ -192,6 +194,11 @@ class OGRFlatGeobufLayer final : public OGRLayer,
         m_bVerifyBuffers = CPL_TO_BOOL(bFlag);
     }
 
+    GDALDataset *GetDataset() override
+    {
+        return m_poDS;
+    }
+
     const std::string &GetFilename() const override
     {
         return m_osFilename;

diff --git a/ogr/ogrsf_frmts/flatgeobuf/ogrflatgeobufdataset.cpp b/ogr/ogrsf_frmts/flatgeobuf/ogrflatgeobufdataset.cpp
@@ -451,7 +451,7 @@ OGRLayer *OGRFlatGeobufDataset::ICreateLayer(
 
     auto poLayer =
         std::unique_ptr<OGRFlatGeobufLayer>(OGRFlatGeobufLayer::Create(
-            pszLayerName, osFilename, poSpatialRef, eGType,
+            this, pszLayerName, osFilename, poSpatialRef, eGType,
             bCreateSpatialIndexAtClose, papszOptions));
     if (poLayer == nullptr)
         return nullptr;

diff --git a/ogr/ogrsf_frmts/flatgeobuf/ogrflatgeobufeditablelayer.cpp b/ogr/ogrsf_frmts/flatgeobuf/ogrflatgeobufeditablelayer.cpp
@@ -89,8 +89,9 @@ OGRErr OGRFlatGeobufEditableLayerSynchronizer::EditableSyncToDisk(
     auto createIndex = m_poFlatGeobufLayer->GetIndexNodeSize() > 0;
 
     OGRFlatGeobufLayer *poFlatGeobufTmpLayer = OGRFlatGeobufLayer::Create(
-        osLayerName.c_str(), osTmpFilename.c_str(), spatialRef, gType,
-        createIndex, m_papszOpenOptions);
+        m_poFlatGeobufLayer->GetDataset(), osLayerName.c_str(),
+        osTmpFilename.c_str(), spatialRef, gType, createIndex,
+        m_papszOpenOptions);
     if (poFlatGeobufTmpLayer == nullptr)
         return OGRERR_FAILURE;
 

diff --git a/ogr/ogrsf_frmts/flatgeobuf/ogrflatgeobuflayer.cpp b/ogr/ogrsf_frmts/flatgeobuf/ogrflatgeobuflayer.cpp
@@ -168,14 +168,12 @@ OGRFlatGeobufLayer::OGRFlatGeobufLayer(const Header *poHeader, GByte *headerBuf,
     m_poFeatureDefn->Reference();
 }
 
-OGRFlatGeobufLayer::OGRFlatGeobufLayer(const char *pszLayerName,
-                                       const char *pszFilename,
-                                       const OGRSpatialReference *poSpatialRef,
-                                       OGRwkbGeometryType eGType,
-                                       bool bCreateSpatialIndexAtClose,
-                                       VSILFILE *poFpWrite,
-                                       std::string &osTempFile)
-    : m_eGType(eGType),
+OGRFlatGeobufLayer::OGRFlatGeobufLayer(
+    GDALDataset *poDS, const char *pszLayerName, const char *pszFilename,
+    const OGRSpatialReference *poSpatialRef, OGRwkbGeometryType eGType,
+    bool bCreateSpatialIndexAtClose, VSILFILE *poFpWrite,
+    std::string &osTempFile)
+    : m_eGType(eGType), m_poDS(poDS),
       m_bCreateSpatialIndexAtClose(bCreateSpatialIndexAtClose),
       m_poFpWrite(poFpWrite), m_osTempFile(osTempFile)
 {
@@ -2409,19 +2407,18 @@ VSILFILE *OGRFlatGeobufLayer::CreateOutputFile(const CPLString &osFilename,
     return poFpWrite;
 }
 
-OGRFlatGeobufLayer *
-OGRFlatGeobufLayer::Create(const char *pszLayerName, const char *pszFilename,
-                           const OGRSpatialReference *poSpatialRef,
-                           OGRwkbGeometryType eGType,
-                           bool bCreateSpatialIndexAtClose, char **papszOptions)
+OGRFlatGeobufLayer *OGRFlatGeobufLayer::Create(
+    GDALDataset *poDS, const char *pszLayerName, const char *pszFilename,
+    const OGRSpatialReference *poSpatialRef, OGRwkbGeometryType eGType,
+    bool bCreateSpatialIndexAtClose, char **papszOptions)
 {
     std::string osTempFile = GetTempFilePath(pszFilename, papszOptions);
     VSILFILE *poFpWrite =
         CreateOutputFile(pszFilename, papszOptions, bCreateSpatialIndexAtClose);
     if (poFpWrite == nullptr)
         return nullptr;
     OGRFlatGeobufLayer *layer = new OGRFlatGeobufLayer(
-        pszLayerName, pszFilename, poSpatialRef, eGType,
+        poDS, pszLayerName, pszFilename, poSpatialRef, eGType,
         bCreateSpatialIndexAtClose, poFpWrite, osTempFile);
     return layer;
 }

diff --git a/ogr/ogrsf_frmts/generic/ogrlayerarrow.cpp b/ogr/ogrsf_frmts/generic/ogrlayerarrow.cpp
@@ -5911,20 +5911,24 @@ bool OGRLayer::CreateFieldFromArrowSchemaInternal(
                     oFieldDefn.SetDefault(oIter.second.c_str());
                 else if (oIter.first == MD_GDAL_OGR_SUBTYPE)
                 {
-                    const auto &osSubType = oIter.second;
-                    for (auto eSubType = OFSTNone; eSubType <= OFSTMaxSubType;)
+                    if (eTypeIn == eTypeOut)
                     {
-                        if (OGRFieldDefn::GetFieldSubTypeName(eSubType) ==
-                            osSubType)
+                        const auto &osSubType = oIter.second;
+                        for (auto eSubType = OFSTNone;
+                             eSubType <= OFSTMaxSubType;)
                         {
-                            oFieldDefn.SetSubType(eSubType);
-                            break;
+                            if (OGRFieldDefn::GetFieldSubTypeName(eSubType) ==
+                                osSubType)
+                            {
+                                oFieldDefn.SetSubType(eSubType);
+                                break;
+                            }
+                            if (eSubType == OFSTMaxSubType)
+                                break;
+                            else
+                                eSubType =
+                                    static_cast<OGRFieldSubType>(eSubType + 1);
                         }
-                        if (eSubType == OFSTMaxSubType)
-                            break;
-                        else
-                            eSubType =
-                                static_cast<OGRFieldSubType>(eSubType + 1);
                     }
                 }
                 else if (oIter.first == MD_GDAL_OGR_WIDTH)