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