Implement GetArrowStream() in OGRLayerDecorator, OGRMutexedLayer and …

…OGRLayerPool
OSGeo · Jun 14, 2022 · 87dbdd6 · 87dbdd6
1 parent 0946e21
commit 87dbdd6
Show file tree

Hide file tree

Showing 6 changed files with 69 additions and 0 deletions.
diff --git a/ogr/ogrsf_frmts/generic/ogrlayerdecorator.cpp b/ogr/ogrsf_frmts/generic/ogrlayerdecorator.cpp
@@ -29,6 +29,7 @@
 #ifndef DOXYGEN_SKIP
 
 #include "ogrlayerdecorator.h"
+#include "ogr_recordbatch.h"
 
 CPL_CVSID("$Id$")
 
@@ -97,6 +98,23 @@ OGRFeature *OGRLayerDecorator::GetNextFeature()
     return m_poDecoratedLayer->GetNextFeature();
 }
 
+GDALDataset *OGRLayerDecorator::GetDataset()
+{
+    if( !m_poDecoratedLayer ) return nullptr;
+    return m_poDecoratedLayer->GetDataset();
+}
+
+bool OGRLayerDecorator::GetArrowStream(struct ArrowArrayStream* out_stream,
+                                     CSLConstList papszOptions)
+{
+    if( !m_poDecoratedLayer )
+    {
+        memset(out_stream, 0, sizeof(*out_stream));
+        return false;
+    }
+    return m_poDecoratedLayer->GetArrowStream(out_stream, papszOptions);
+}
+
 OGRErr      OGRLayerDecorator::SetNextByIndex( GIntBig nIndex )
 {
     if( !m_poDecoratedLayer ) return OGRERR_FAILURE;

diff --git a/ogr/ogrsf_frmts/generic/ogrlayerdecorator.h b/ogr/ogrsf_frmts/generic/ogrlayerdecorator.h
@@ -66,6 +66,10 @@ class CPL_DLL OGRLayerDecorator : public OGRLayer
     virtual OGRErr      ICreateFeature( OGRFeature *poFeature ) override;
     virtual OGRErr      DeleteFeature( GIntBig nFID ) override;
 
+    virtual GDALDataset* GetDataset() override;
+    virtual bool         GetArrowStream(struct ArrowArrayStream* out_stream,
+                                        CSLConstList papszOptions = nullptr) override;
+
     virtual const char *GetName() override;
     virtual OGRwkbGeometryType GetGeomType() override;
     virtual OGRFeatureDefn *GetLayerDefn() override;

diff --git a/ogr/ogrsf_frmts/generic/ogrlayerpool.cpp b/ogr/ogrsf_frmts/generic/ogrlayerpool.cpp
@@ -29,6 +29,7 @@
 #ifndef DOXYGEN_SKIP
 
 #include "ogrlayerpool.h"
+#include "ogr_recordbatch.h"
 
 CPL_CVSID("$Id$")
 
@@ -287,6 +288,31 @@ OGRFeature *OGRProxiedLayer::GetNextFeature()
     return poUnderlyingLayer->GetNextFeature();
 }
 
+/************************************************************************/
+/*                            GDALDataset()                             */
+/************************************************************************/
+
+GDALDataset *OGRProxiedLayer::GetDataset()
+{
+    if( poUnderlyingLayer == nullptr && !OpenUnderlyingLayer() ) return nullptr;
+    return poUnderlyingLayer->GetDataset();
+}
+
+/************************************************************************/
+/*                          GetArrowStream()                            */
+/************************************************************************/
+
+bool OGRProxiedLayer::GetArrowStream(struct ArrowArrayStream* out_stream,
+                                     CSLConstList papszOptions)
+{
+    if( poUnderlyingLayer == nullptr && !OpenUnderlyingLayer() )
+    {
+        memset(out_stream, 0, sizeof(*out_stream));
+        return false;
+    }
+    return poUnderlyingLayer->GetArrowStream(out_stream, papszOptions);
+}
+
 /************************************************************************/
 /*                           SetNextByIndex()                           */
 /************************************************************************/

diff --git a/ogr/ogrsf_frmts/generic/ogrlayerpool.h b/ogr/ogrsf_frmts/generic/ogrlayerpool.h
@@ -132,6 +132,10 @@ class OGRProxiedLayer : public OGRAbstractProxiedLayer
     virtual OGRErr      ICreateFeature( OGRFeature *poFeature ) override;
     virtual OGRErr      DeleteFeature( GIntBig nFID ) override;
 
+    virtual GDALDataset* GetDataset() override;
+    virtual bool         GetArrowStream(struct ArrowArrayStream* out_stream,
+                                        CSLConstList papszOptions = nullptr) override;
+
     virtual const char *GetName() override;
     virtual OGRwkbGeometryType GetGeomType() override;
     virtual OGRFeatureDefn *GetLayerDefn() override;

diff --git a/ogr/ogrsf_frmts/generic/ogrmutexedlayer.cpp b/ogr/ogrsf_frmts/generic/ogrmutexedlayer.cpp
@@ -93,6 +93,19 @@ OGRFeature *OGRMutexedLayer::GetNextFeature()
     return OGRLayerDecorator::GetNextFeature();
 }
 
+GDALDataset *OGRMutexedLayer::GetDataset()
+{
+    CPLMutexHolderOptionalLockD(m_hMutex);
+    return OGRLayerDecorator::GetDataset();
+}
+
+bool OGRMutexedLayer::GetArrowStream(struct ArrowArrayStream* out_stream,
+                                     CSLConstList papszOptions)
+{
+    CPLMutexHolderOptionalLockD(m_hMutex);
+    return OGRLayerDecorator::GetArrowStream(out_stream, papszOptions);
+}
+
 OGRErr      OGRMutexedLayer::SetNextByIndex( GIntBig nIndex )
 {
     CPLMutexHolderOptionalLockD(m_hMutex);

diff --git a/ogr/ogrsf_frmts/generic/ogrmutexedlayer.h b/ogr/ogrsf_frmts/generic/ogrmutexedlayer.h
@@ -77,6 +77,10 @@ class CPL_DLL OGRMutexedLayer : public OGRLayerDecorator
     virtual OGRErr      ICreateFeature( OGRFeature *poFeature ) override;
     virtual OGRErr      DeleteFeature( GIntBig nFID ) override;
 
+    virtual GDALDataset* GetDataset() override;
+    virtual bool         GetArrowStream(struct ArrowArrayStream* out_stream,
+                                        CSLConstList papszOptions = nullptr) override;
+
     virtual const char *GetName() override;
     virtual OGRwkbGeometryType GetGeomType() override;
     virtual OGRFeatureDefn *GetLayerDefn() override;