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             :     friend class OGRADBCDataset;
      91             : 
      92             :     OGRADBCDataset *m_poDS = nullptr;
      93             :     std::unique_ptr<AdbcStatement> m_statement{};
      94             :     std::unique_ptr<OGRArrowArrayToOGRFeatureAdapterLayer> m_poAdapterLayer{};
      95             :     std::unique_ptr<OGRArrowArrayStream> m_stream{};
      96             : 
      97             :     struct ArrowSchema m_schema
      98             :     {
      99             :     };
     100             : 
     101             :     bool m_bEOF = false;
     102             :     size_t m_nIdx = 0;
     103             :     GIntBig m_nFeatureID = 0;
     104             :     bool m_bIsParquetLayer = false;
     105             : 
     106             :     OGRFeature *GetNextRawFeature();
     107             :     bool GetArrowStreamInternal(struct ArrowArrayStream *out_stream);
     108             :     GIntBig GetFeatureCountParquet();
     109             : 
     110             :     CPL_DISALLOW_COPY_ASSIGN(OGRADBCLayer)
     111             : 
     112             :   public:
     113             :     OGRADBCLayer(OGRADBCDataset *poDS, const char *pszName,
     114             :                  std::unique_ptr<AdbcStatement> poStatement,
     115             :                  std::unique_ptr<OGRArrowArrayStream> poStream,
     116             :                  ArrowSchema *schema);
     117             :     ~OGRADBCLayer() override;
     118             : 
     119           0 :     OGRFeatureDefn *GetLayerDefn() override
     120             :     {
     121           0 :         return m_poAdapterLayer->GetLayerDefn();
     122             :     }
     123             : 
     124             :     void ResetReading() override;
     125           0 :     DEFINE_GET_NEXT_FEATURE_THROUGH_RAW(OGRADBCLayer)
     126             :     int TestCapability(const char *) override;
     127             :     GDALDataset *GetDataset() override;
     128             :     bool GetArrowStream(struct ArrowArrayStream *out_stream,
     129             :                         CSLConstList papszOptions = nullptr) override;
     130             :     GIntBig GetFeatureCount(int bForce) override;
     131             : };
     132             : 
     133             : /************************************************************************/
     134             : /*                            OGRADBCDataset                            */
     135             : /************************************************************************/
     136             : 
     137             : class OGRADBCDataset final : public GDALDataset
     138             : {
     139             :     friend class OGRADBCLayer;
     140             : 
     141             :     AdbcDriver m_driver{};
     142             :     AdbcDatabase m_database{};
     143             :     std::unique_ptr<AdbcConnection> m_connection{};
     144             :     std::vector<std::unique_ptr<OGRLayer>> m_apoLayers{};
     145             :     std::string m_osParquetFilename{};
     146             : 
     147             :   public:
     148           0 :     OGRADBCDataset() = default;
     149             :     ~OGRADBCDataset() override;
     150             : 
     151             :     bool Open(const GDALOpenInfo *poOpenInfo);
     152             : 
     153           0 :     int GetLayerCount() override
     154             :     {
     155           0 :         return static_cast<int>(m_apoLayers.size());
     156             :     }
     157             : 
     158           0 :     OGRLayer *GetLayer(int idx) override
     159             :     {
     160           0 :         return (idx >= 0 && idx < GetLayerCount()) ? m_apoLayers[idx].get()
     161           0 :                                                    : nullptr;
     162             :     }
     163             : 
     164             :     OGRLayer *GetLayerByName(const char *pszName) override;
     165             : 
     166             :     std::unique_ptr<OGRADBCLayer> CreateLayer(const char *pszStatement,
     167             :                                               const char *pszLayerName);
     168             : 
     169             :     OGRLayer *ExecuteSQL(const char *pszStatement, OGRGeometry *poSpatialFilter,
     170             :                          const char *pszDialect) override;
     171             : };
     172             : 
     173             : /************************************************************************/
     174             : /*                            OGRADBCError                              */
     175             : /************************************************************************/
     176             : 
     177             : struct OGRADBCError
     178             : {
     179             :     AdbcError error{ADBC_ERROR_INIT};
     180             : 
     181             :     inline OGRADBCError() = default;
     182             : 
     183           0 :     inline ~OGRADBCError()
     184           0 :     {
     185           0 :         clear();
     186           0 :     }
     187             : 
     188           0 :     inline void clear()
     189             :     {
     190           0 :         if (error.release)
     191           0 :             error.release(&error);
     192           0 :         memset(&error, 0, sizeof(error));
     193           0 :     }
     194             : 
     195           0 :     inline const char *message() const
     196             :     {
     197           0 :         return error.message ? error.message : "";
     198             :     }
     199             : 
     200           0 :     inline operator AdbcError *()
     201             :     {
     202           0 :         return &error;
     203             :     }
     204             : 
     205             :     CPL_DISALLOW_COPY_ASSIGN(OGRADBCError)
     206             : };
     207             : 
     208             : #endif  // OGR_ADBC_INCLUDED