Line data    Source code

       1             : /******************************************************************************
       2             :  *
       3             :  * Project:  GDAL
       4             :  * Purpose:  Arrow Database Connectivity driver
       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             : #ifndef OGR_ADBC_INCLUDED
      14             : #define OGR_ADBC_INCLUDED
      15             : 
      16             : #include "gdal_priv.h"
      17             : #include "ogrsf_frmts.h"
      18             : #include "ogrlayerarrow.h"
      19             : 
      20             : #include "ogr_adbc_internal.h"
      21             : 
      22             : /************************************************************************/
      23             : /*                 OGRArrowArrayToOGRFeatureAdapterLayer                */
      24             : /************************************************************************/
      25             : 
      26             : class OGRArrowArrayToOGRFeatureAdapterLayer final : public OGRLayer
      27             : {
      28             :     friend class OGRADBCLayer;
      29             :     OGRFeatureDefn *m_poLayerDefn = nullptr;
      30             :     std::vector<std::unique_ptr<OGRFeature>> m_apoFeatures{};
      31             : 
      32             :     CPL_DISALLOW_COPY_ASSIGN(OGRArrowArrayToOGRFeatureAdapterLayer)
      33             : 
      34             :   public:
      35           0 :     explicit OGRArrowArrayToOGRFeatureAdapterLayer(const char *pszName)
      36           0 :     {
      37           0 :         m_poLayerDefn = new OGRFeatureDefn(pszName);
      38           0 :         m_poLayerDefn->SetGeomType(wkbNone);
      39           0 :         m_poLayerDefn->Reference();
      40           0 :     }
      41             : 
      42           0 :     ~OGRArrowArrayToOGRFeatureAdapterLayer()
      43           0 :     {
      44           0 :         m_poLayerDefn->Release();
      45           0 :     }
      46             : 
      47           0 :     OGRFeatureDefn *GetLayerDefn() override
      48             :     {
      49           0 :         return m_poLayerDefn;
      50             :     }
      51             : 
      52           0 :     void ResetReading() override
      53             :     {
      54           0 :     }
      55             : 
      56           0 :     OGRFeature *GetNextFeature() override
      57             :     {
      58           0 :         return nullptr;
      59             :     }
      60             : 
      61           0 :     int TestCapability(const char *pszCap) override
      62             :     {
      63           0 :         return EQUAL(pszCap, OLCCreateField) ||
      64           0 :                EQUAL(pszCap, OLCSequentialWrite);
      65             :     }
      66             : 
      67           0 :     OGRErr CreateField(const OGRFieldDefn *poFieldDefn, int) override
      68             :     {
      69           0 :         m_poLayerDefn->AddFieldDefn(poFieldDefn);
      70           0 :         return OGRERR_NONE;
      71             :     }
      72             : 
      73           0 :     OGRErr ICreateFeature(OGRFeature *poFeature) override
      74             :     {
      75             :         m_apoFeatures.emplace_back(
      76           0 :             std::unique_ptr<OGRFeature>(poFeature->Clone()));
      77           0 :         return OGRERR_NONE;
      78             :     }
      79             : };
      80             : 
      81             : /************************************************************************/
      82             : /*                            OGRADBCLayer                              */
      83             : /************************************************************************/
      84             : 
      85             : class OGRADBCDataset;
      86             : 
      87             : class OGRADBCLayer final : public OGRLayer,
      88             :                            public OGRGetNextFeatureThroughRaw<OGRADBCLayer>
      89             : {
      90             :   public:
      91             :     //! Describe the bbox column of a geometry column
      92             :     struct GeomColBBOX
      93             :     {
      94             :         std::string osXMin{};  // empty if no bbox column
      95             :         std::string osYMin{};
      96             :         std::string osXMax{};
      97             :         std::string osYMax{};
      98             :     };
      99             : 
     100             :   private:
     101             :     friend class OGRADBCDataset;
     102             : 
     103             :     OGRADBCDataset *m_poDS = nullptr;
     104             :     const std::string m_osBaseStatement{};    // as provided by user
     105             :     std::string m_osModifiedBaseStatement{};  // above tuned to use ST_AsWKB()
     106             :     std::string m_osModifiedSelect{};         // SELECT part of above
     107             :     std::string m_osAttributeFilter{};
     108             :     std::unique_ptr<AdbcStatement> m_statement{};
     109             :     std::unique_ptr<OGRArrowArrayToOGRFeatureAdapterLayer> m_poAdapterLayer{};
     110             :     std::unique_ptr<OGRArrowArrayStream> m_stream{};
     111             : 
     112             :     struct ArrowSchema m_schema
     113             :     {
     114             :     };
     115             : 
     116             :     bool m_bEOF = false;
     117             :     size_t m_nIdx = 0;
     118             :     GIntBig m_nFeatureID = 0;
     119             :     bool m_bIsParquetLayer = false;
     120             : 
     121             :     std::vector<GeomColBBOX>
     122             :         m_geomColBBOX{};                     // same size as GetGeomFieldCount()
     123             :     std::vector<OGREnvelope3D> m_extents{};  // same size as GetGeomFieldCount()
     124             : 
     125             :     OGRFeature *GetNextRawFeature();
     126             :     bool GetArrowStreamInternal(struct ArrowArrayStream *out_stream);
     127             :     GIntBig GetFeatureCountParquet();
     128             : 
     129             :     void BuildLayerDefn(bool bInternalUse);
     130             :     bool ReplaceStatement(const char *pszNewStatement);
     131             :     bool UpdateStatement();
     132             :     std::string GetCurrentStatement() const;
     133             : 
     134             :     CPL_DISALLOW_COPY_ASSIGN(OGRADBCLayer)
     135             : 
     136             :   public:
     137             :     OGRADBCLayer(OGRADBCDataset *poDS, const char *pszName,
     138             :                  const char *pszStatement,
     139             :                  std::unique_ptr<AdbcStatement> poStatement,
     140             :                  std::unique_ptr<OGRArrowArrayStream> poStream,
     141             :                  ArrowSchema *schema, bool bInternalUse);
     142             :     ~OGRADBCLayer() override;
     143             : 
     144           0 :     OGRFeatureDefn *GetLayerDefn() override
     145             :     {
     146           0 :         return m_poAdapterLayer->GetLayerDefn();
     147             :     }
     148             : 
     149             :     void ResetReading() override;
     150           0 :     DEFINE_GET_NEXT_FEATURE_THROUGH_RAW(OGRADBCLayer)
     151             :     int TestCapability(const char *) override;
     152             :     GDALDataset *GetDataset() override;
     153             :     bool GetArrowStream(struct ArrowArrayStream *out_stream,
     154             :                         CSLConstList papszOptions = nullptr) override;
     155             :     GIntBig GetFeatureCount(int bForce) override;
     156             : 
     157           0 :     void SetSpatialFilter(OGRGeometry *poGeom) override
     158             :     {
     159           0 :         SetSpatialFilter(0, poGeom);
     160           0 :     }
     161             : 
     162             :     OGRErr SetAttributeFilter(const char *pszFilter) override;
     163             :     void SetSpatialFilter(int iGeomField, OGRGeometry *poGeom) override;
     164             : 
     165             :     OGRErr GetExtent(OGREnvelope *psExtent, int bForce = TRUE) override;
     166             :     OGRErr GetExtent(int iGeomField, OGREnvelope *psExtent,
     167             :                      int bForce = TRUE) override;
     168             :     OGRErr GetExtent3D(int iGeomField, OGREnvelope3D *psExtent,
     169             :                        int bForce = TRUE) override;
     170             : };
     171             : 
     172             : /************************************************************************/
     173             : /*                            OGRADBCDataset                            */
     174             : /************************************************************************/
     175             : 
     176             : class OGRADBCDataset final : public GDALDataset
     177             : {
     178             :     friend class OGRADBCLayer;
     179             : 
     180             :     AdbcDriver m_driver{};
     181             :     AdbcDatabase m_database{};
     182             :     std::unique_ptr<AdbcConnection> m_connection{};
     183             :     std::vector<std::unique_ptr<OGRLayer>> m_apoLayers{};
     184             :     std::string m_osParquetFilename{};
     185             :     bool m_bIsDuckDBDataset = false;
     186             :     bool m_bIsDuckDBDriver = false;
     187             :     bool m_bSpatialLoaded = false;
     188             : 
     189             :   public:
     190           0 :     OGRADBCDataset() = default;
     191             :     ~OGRADBCDataset() override;
     192             : 
     193             :     bool Open(const GDALOpenInfo *poOpenInfo);
     194             : 
     195           0 :     int GetLayerCount() override
     196             :     {
     197           0 :         return static_cast<int>(m_apoLayers.size());
     198             :     }
     199             : 
     200           0 :     OGRLayer *GetLayer(int idx) override
     201             :     {
     202           0 :         return (idx >= 0 && idx < GetLayerCount()) ? m_apoLayers[idx].get()
     203           0 :                                                    : nullptr;
     204             :     }
     205             : 
     206             :     OGRLayer *GetLayerByName(const char *pszName) override;
     207             : 
     208             :     std::unique_ptr<OGRADBCLayer> CreateLayer(const char *pszStatement,
     209             :                                               const char *pszLayerName,
     210             :                                               bool bInternalUse);
     211             : 
     212           0 :     std::unique_ptr<OGRADBCLayer> CreateInternalLayer(const char *pszStatement)
     213             :     {
     214           0 :         return CreateLayer(pszStatement, "temp", true);
     215             :     }
     216             : 
     217             :     OGRLayer *ExecuteSQL(const char *pszStatement, OGRGeometry *poSpatialFilter,
     218             :                          const char *pszDialect) override;
     219             : };
     220             : 
     221             : /************************************************************************/
     222             : /*                            OGRADBCError                              */
     223             : /************************************************************************/
     224             : 
     225             : struct OGRADBCError
     226             : {
     227             :     AdbcError error{ADBC_ERROR_INIT};
     228             : 
     229             :     inline OGRADBCError() = default;
     230             : 
     231           0 :     inline ~OGRADBCError()
     232           0 :     {
     233           0 :         clear();
     234           0 :     }
     235             : 
     236           0 :     inline void clear()
     237             :     {
     238           0 :         if (error.release)
     239           0 :             error.release(&error);
     240           0 :         memset(&error, 0, sizeof(error));
     241           0 :     }
     242             : 
     243           0 :     inline const char *message() const
     244             :     {
     245           0 :         return error.message ? error.message : "";
     246             :     }
     247             : 
     248           0 :     inline operator AdbcError *()
     249             :     {
     250           0 :         return &error;
     251             :     }
     252             : 
     253             :     CPL_DISALLOW_COPY_ASSIGN(OGRADBCError)
     254             : };
     255             : 
     256             : #endif  // OGR_ADBC_INCLUDED