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             :     ~OGRArrowArrayToOGRFeatureAdapterLayer() override;
      43             : 
      44             :     using OGRLayer::GetLayerDefn;
      45             : 
      46           0 :     const OGRFeatureDefn *GetLayerDefn() const override
      47             :     {
      48           0 :         return m_poLayerDefn;
      49             :     }
      50             : 
      51           0 :     void ResetReading() override
      52             :     {
      53           0 :     }
      54             : 
      55           0 :     OGRFeature *GetNextFeature() override
      56             :     {
      57           0 :         return nullptr;
      58             :     }
      59             : 
      60           0 :     int TestCapability(const char *pszCap) const override
      61             :     {
      62           0 :         return EQUAL(pszCap, OLCCreateField) ||
      63           0 :                EQUAL(pszCap, OLCSequentialWrite);
      64             :     }
      65             : 
      66           0 :     OGRErr CreateField(const OGRFieldDefn *poFieldDefn, int) override
      67             :     {
      68           0 :         m_poLayerDefn->AddFieldDefn(poFieldDefn);
      69           0 :         return OGRERR_NONE;
      70             :     }
      71             : 
      72           0 :     OGRErr CreateGeomField(const OGRGeomFieldDefn *poGeomFieldDefn,
      73             :                            int) override
      74             :     {
      75           0 :         m_poLayerDefn->AddGeomFieldDefn(poGeomFieldDefn);
      76           0 :         return OGRERR_NONE;
      77             :     }
      78             : 
      79           0 :     OGRErr ICreateFeature(OGRFeature *poFeature) override
      80             :     {
      81             :         m_apoFeatures.emplace_back(
      82           0 :             std::unique_ptr<OGRFeature>(poFeature->Clone()));
      83           0 :         return OGRERR_NONE;
      84             :     }
      85             : };
      86             : 
      87             : /************************************************************************/
      88             : /*                            OGRADBCLayer                              */
      89             : /************************************************************************/
      90             : 
      91             : class OGRADBCDataset;
      92             : 
      93             : class OGRADBCLayer /* non final */ : public OGRLayer,
      94             :                                      public OGRGetNextFeatureThroughRaw<
      95             :                                          OGRADBCLayer>
      96             : {
      97             :   public:
      98             :     //! Describe the bbox column of a geometry column
      99             :     struct GeomColBBOX
     100             :     {
     101             :         std::string osXMin{};  // empty if no bbox column
     102             :         std::string osYMin{};
     103             :         std::string osXMax{};
     104             :         std::string osYMax{};
     105             :     };
     106             : 
     107             :   protected:
     108             :     friend class OGRADBCDataset;
     109             : 
     110             :     OGRADBCDataset *m_poDS = nullptr;
     111             :     const std::string m_osBaseStatement{};    // as provided by user
     112             :     std::string m_osModifiedBaseStatement{};  // above tuned to use ST_AsWKB()
     113             :     std::string m_osModifiedSelect{};         // SELECT part of above
     114             :     std::string m_osAttributeFilter{};
     115             :     std::unique_ptr<AdbcStatement> m_statement{};
     116             :     std::unique_ptr<OGRArrowArrayToOGRFeatureAdapterLayer> m_poAdapterLayer{};
     117             :     std::unique_ptr<OGRArrowArrayStream> m_stream{};
     118             :     bool m_bInternalUse = false;
     119             :     bool m_bLayerDefinitionError = false;
     120             : 
     121             :     struct ArrowSchema m_schema
     122             :     {
     123             :     };
     124             : 
     125             :     bool m_bEOF = false;
     126             :     size_t m_nIdx = 0;
     127             :     GIntBig m_nFeatureID = 0;
     128             :     GIntBig m_nMaxFeatureID = -1;
     129             :     bool m_bIsParquetLayer = false;
     130             : 
     131             :     std::vector<GeomColBBOX>
     132             :         m_geomColBBOX{};                     // same size as GetGeomFieldCount()
     133             :     std::vector<OGREnvelope3D> m_extents{};  // same size as GetGeomFieldCount()
     134             :     std::string m_osFIDColName{};
     135             : 
     136             :     OGRFeature *GetNextRawFeature();
     137             :     bool GetArrowStreamInternal(struct ArrowArrayStream *out_stream);
     138             :     GIntBig GetFeatureCountSelectCountStar();
     139             :     GIntBig GetFeatureCountArrow();
     140             :     GIntBig GetFeatureCountParquet();
     141             : 
     142             :     bool BuildLayerDefnInit();
     143             :     virtual void BuildLayerDefn();
     144             :     bool ReplaceStatement(const char *pszNewStatement);
     145             :     bool UpdateStatement();
     146             :     virtual std::string GetCurrentStatement() const;
     147             : 
     148           0 :     virtual bool RunDeferredCreation()
     149             :     {
     150           0 :         return true;
     151             :     }
     152             : 
     153             :     CPL_DISALLOW_COPY_ASSIGN(OGRADBCLayer)
     154             : 
     155             :   public:
     156             :     OGRADBCLayer(OGRADBCDataset *poDS, const char *pszName,
     157             :                  const std::string &osStatement, bool bInternalUse);
     158             :     OGRADBCLayer(OGRADBCDataset *poDS, const char *pszName,
     159             :                  std::unique_ptr<OGRArrowArrayStream> poStream,
     160             :                  ArrowSchema *schema, bool bInternalUse);
     161             :     ~OGRADBCLayer() override;
     162             : 
     163             :     bool GotError();
     164             : 
     165             :     using OGRLayer::GetLayerDefn;
     166             :     const OGRFeatureDefn *GetLayerDefn() const override;
     167             : 
     168           0 :     const char *GetName() const override
     169             :     {
     170           0 :         return GetDescription();
     171             :     }
     172             : 
     173             :     void ResetReading() override;
     174           0 :     DEFINE_GET_NEXT_FEATURE_THROUGH_RAW(OGRADBCLayer)
     175             :     int TestCapability(const char *) const override;
     176             :     GDALDataset *GetDataset() override;
     177             :     bool GetArrowStream(struct ArrowArrayStream *out_stream,
     178             :                         CSLConstList papszOptions = nullptr) override;
     179             :     GIntBig GetFeatureCount(int bForce) override;
     180             : 
     181             :     OGRErr SetAttributeFilter(const char *pszFilter) override;
     182             :     OGRErr ISetSpatialFilter(int iGeomField,
     183             :                              const OGRGeometry *poGeom) override;
     184             : 
     185             :     OGRErr IGetExtent(int iGeomField, OGREnvelope *psExtent,
     186             :                       bool bForce) override;
     187             :     OGRErr IGetExtent3D(int iGeomField, OGREnvelope3D *psExtent,
     188             :                         bool bForce) override;
     189             : 
     190             :     const char *GetFIDColumn() const override;
     191             : };
     192             : 
     193             : /************************************************************************/
     194             : /*                      OGRADBCBigQueryLayer                            */
     195             : /************************************************************************/
     196             : 
     197             : class OGRADBCBigQueryLayer final : public OGRADBCLayer
     198             : {
     199             :   private:
     200             :     friend class OGRADBCDataset;
     201             : 
     202             :     bool m_bDeferredCreation = false;
     203             : 
     204             :     void BuildLayerDefn() override;
     205             :     bool RunDeferredCreation() override;
     206             :     std::string GetCurrentStatement() const override;
     207             :     bool GetBigQueryDatasetAndTableId(std::string &osDatasetId,
     208             :                                       std::string &osTableId) const;
     209             : 
     210             :   public:
     211             :     OGRADBCBigQueryLayer(OGRADBCDataset *poDS, const char *pszName,
     212             :                          const std::string &osStatement, bool bInternalUse);
     213             : 
     214             :     int TestCapability(const char *) const override;
     215             :     GIntBig GetFeatureCount(int bForce) override;
     216             :     OGRErr IGetExtent(int iGeomField, OGREnvelope *psExtent,
     217             :                       bool bForce) override;
     218             :     OGRErr SetAttributeFilter(const char *pszFilter) override;
     219             : 
     220             :     OGRErr CreateField(const OGRFieldDefn *poField, int bApproxOK) override;
     221             :     OGRErr ICreateFeature(OGRFeature *poFeature) override;
     222             :     OGRErr ISetFeature(OGRFeature *poFeature) override;
     223             :     OGRErr DeleteFeature(GIntBig nFID) override;
     224             : 
     225             :     void SetDeferredCreation(const char *pszFIDColName,
     226             :                              const OGRGeomFieldDefn *poGeomFieldDefn);
     227             : };
     228             : 
     229             : /************************************************************************/
     230             : /*                            OGRADBCDataset                            */
     231             : /************************************************************************/
     232             : 
     233             : class OGRADBCDataset final : public GDALDataset
     234             : {
     235             :     friend class OGRADBCLayer;
     236             : 
     237             :     AdbcDriver m_driver{};
     238             :     AdbcDatabase m_database{};
     239             :     std::unique_ptr<AdbcConnection> m_connection{};
     240             :     std::vector<std::unique_ptr<OGRLayer>> m_apoLayers{};
     241             :     std::string m_osParquetFilename{};
     242             :     bool m_bIsDuckDBDataset = false;
     243             :     bool m_bIsDuckDBDriver = false;
     244             :     bool m_bSpatialLoaded = false;
     245             :     bool m_bIsBigQuery = false;
     246             :     std::string m_osBigQueryDatasetId{};
     247             : 
     248             :   public:
     249           1 :     OGRADBCDataset() = default;
     250             :     ~OGRADBCDataset() override;
     251             : 
     252             :     CPLErr FlushCache(bool bAtClosing) override;
     253             : 
     254             :     bool Open(const GDALOpenInfo *poOpenInfo);
     255             : 
     256           0 :     int GetLayerCount() const override
     257             :     {
     258           0 :         return static_cast<int>(m_apoLayers.size());
     259             :     }
     260             : 
     261           0 :     const OGRLayer *GetLayer(int idx) const override
     262             :     {
     263           0 :         return (idx >= 0 && idx < GetLayerCount()) ? m_apoLayers[idx].get()
     264           0 :                                                    : nullptr;
     265             :     }
     266             : 
     267             :     OGRLayer *GetLayerByName(const char *pszName) override;
     268             : 
     269             :     std::unique_ptr<OGRADBCLayer> CreateLayer(const char *pszStatement,
     270             :                                               const char *pszLayerName,
     271             :                                               bool bInternalUse);
     272             : 
     273             :     std::unique_ptr<OGRADBCLayer>
     274             :     CreateInternalLayer(const char *pszStatement) CPL_WARN_UNUSED_RESULT;
     275             : 
     276             :     OGRLayer *ICreateLayer(const char *pszName,
     277             :                            const OGRGeomFieldDefn *poGeomFieldDefn,
     278             :                            CSLConstList papszOptions) override;
     279             :     OGRErr DeleteLayer(int iLayer) override;
     280             : 
     281             :     OGRLayer *ExecuteSQL(const char *pszStatement, OGRGeometry *poSpatialFilter,
     282             :                          const char *pszDialect) override;
     283             : 
     284             :     int TestCapability(const char *pszCap) const override;
     285             : };
     286             : 
     287             : /************************************************************************/
     288             : /*                            OGRADBCError                              */
     289             : /************************************************************************/
     290             : 
     291             : struct OGRADBCError
     292             : {
     293             :     AdbcError error{ADBC_ERROR_INIT};
     294             : 
     295             :     inline OGRADBCError() = default;
     296             : 
     297           2 :     inline ~OGRADBCError()
     298           2 :     {
     299           2 :         clear();
     300           2 :     }
     301             : 
     302           3 :     inline void clear()
     303             :     {
     304           3 :         if (error.release)
     305           0 :             error.release(&error);
     306           3 :         memset(&error, 0, sizeof(error));
     307           3 :     }
     308             : 
     309           1 :     inline const char *message() const
     310             :     {
     311           1 :         return error.message ? error.message : "";
     312             :     }
     313             : 
     314           1 :     inline operator AdbcError *()
     315             :     {
     316           1 :         return &error;
     317             :     }
     318             : 
     319             :     CPL_DISALLOW_COPY_ASSIGN(OGRADBCError)
     320             : };
     321             : 
     322             : #endif  // OGR_ADBC_INCLUDED