Line data    Source code

       1             : /******************************************************************************
       2             :  *
       3             :  * Project:  Feather Translator
       4             :  * Purpose:  Implements OGRFeatherDriver.
       5             :  * Author:   Even Rouault, <even.rouault at spatialys.com>
       6             :  *
       7             :  ******************************************************************************
       8             :  * Copyright (c) 2022, Planet Labs
       9             :  *
      10             :  * SPDX-License-Identifier: MIT
      11             :  ****************************************************************************/
      12             : 
      13             : #ifndef OGR_FEATHER_H
      14             : #define OGR_FEATHER_H
      15             : 
      16             : #include "ogrsf_frmts.h"
      17             : 
      18             : #include <map>
      19             : 
      20             : #include "../arrow_common/ogr_arrow.h"
      21             : 
      22             : #ifdef _MSC_VER
      23             : #pragma warning(push)
      24             : // warning 4244: 'initializing': conversion from 'int32_t' to 'int16_t',
      25             : // possible loss of data
      26             : #pragma warning(disable : 4244)
      27             : // warning 4458: declaration of 'type_id' hides class member
      28             : #pragma warning(disable : 4458)
      29             : #endif
      30             : 
      31             : #include "arrow/ipc/writer.h"
      32             : 
      33             : #ifdef _MSC_VER
      34             : #pragma warning(pop)
      35             : #endif
      36             : 
      37             : constexpr const char *GDAL_GEO_FOOTER_KEY = "gdal:geo";
      38             : constexpr const char *ARROW_DRIVER_NAME_UC = "ARROW";
      39             : 
      40             : /************************************************************************/
      41             : /*                        OGRFeatherLayer                               */
      42             : /************************************************************************/
      43             : 
      44             : class OGRFeatherDataset;
      45             : 
      46             : class OGRFeatherLayer final : public OGRArrowLayer
      47             : 
      48             : {
      49             :     OGRFeatherLayer(const OGRFeatherLayer &) = delete;
      50             :     OGRFeatherLayer &operator=(const OGRFeatherLayer &) = delete;
      51             : 
      52             :     OGRFeatherDataset *m_poDS = nullptr;
      53             : 
      54             :     // Variable only for seekable file format
      55             :     std::shared_ptr<arrow::ipc::RecordBatchFileReader>
      56             :         m_poRecordBatchFileReader{};
      57             : 
      58             :     // Variables only for streamable IPC format
      59             :     std::shared_ptr<arrow::io::RandomAccessFile> m_poFile{};
      60             :     bool m_bSeekable = true;
      61             :     arrow::ipc::IpcReadOptions m_oOptions{};
      62             :     std::shared_ptr<arrow::ipc::RecordBatchStreamReader>
      63             :         m_poRecordBatchReader{};
      64             :     bool m_bResetRecordBatchReaderAsked = false;
      65             :     bool m_bSingleBatch = false;
      66             :     std::shared_ptr<arrow::RecordBatch> m_poBatchIdx0{};
      67             :     std::shared_ptr<arrow::RecordBatch> m_poBatchIdx1{};
      68             : 
      69             :     CPLStringList m_aosFeatherMetadata{};
      70             : 
      71         186 :     std::string GetDriverUCName() const override
      72             :     {
      73         186 :         return ARROW_DRIVER_NAME_UC;
      74             :     }
      75             : 
      76             :     bool ResetRecordBatchReader();
      77             : 
      78             :     void EstablishFeatureDefn();
      79             :     void LoadGeoMetadata(const arrow::KeyValueMetadata *kv_metadata,
      80             :                          const std::string &key);
      81             :     OGRwkbGeometryType ComputeGeometryColumnType(int iGeomCol, int iCol) const;
      82             :     bool ReadNextBatch() override;
      83             : 
      84             :     void InvalidateCachedBatches() override;
      85             : 
      86             :     OGRFeature *GetNextRawFeature();
      87             : 
      88             :     bool CanRunNonForcedGetExtent() override;
      89             : 
      90             :     bool
      91             :     CanPostFilterArrowArray(const struct ArrowSchema *schema) const override;
      92             : 
      93             :     bool ReadNextBatchFile();
      94             :     bool ReadNextBatchStream();
      95             :     void TryToCacheFirstTwoBatches();
      96             : 
      97             :   public:
      98             :     OGRFeatherLayer(OGRFeatherDataset *poDS, const char *pszLayerName,
      99             :                     std::shared_ptr<arrow::ipc::RecordBatchFileReader>
     100             :                         &poRecordBatchFileReader,
     101             :                     CSLConstList papszOpenOptions);
     102             :     OGRFeatherLayer(OGRFeatherDataset *poDS, const char *pszLayerName,
     103             :                     std::shared_ptr<arrow::io::RandomAccessFile> poFile,
     104             :                     bool bSeekable, const arrow::ipc::IpcReadOptions &oOptions,
     105             :                     std::shared_ptr<arrow::ipc::RecordBatchStreamReader>
     106             :                         &poRecordBatchStreamReader,
     107             :                     CSLConstList papszOpenOptions);
     108             : 
     109             :     void ResetReading() override;
     110             :     int TestCapability(const char *pszCap) const override;
     111             :     GIntBig GetFeatureCount(int bForce) override;
     112             :     const char *GetMetadataItem(const char *pszName,
     113             :                                 const char *pszDomain = "") override;
     114             :     char **GetMetadata(const char *pszDomain = "") override;
     115             : 
     116             :     GDALDataset *GetDataset() override;
     117             : 
     118             :     std::unique_ptr<OGRFieldDomain> BuildDomain(const std::string &osDomainName,
     119             :                                                 int iFieldIndex) const override;
     120             : };
     121             : 
     122             : /************************************************************************/
     123             : /*                         OGRFeatherDataset                            */
     124             : /************************************************************************/
     125             : 
     126             : class OGRFeatherDataset final : public OGRArrowDataset
     127             : {
     128             :   public:
     129             :     explicit OGRFeatherDataset(
     130             :         const std::shared_ptr<arrow::MemoryPool> &poMemoryPool);
     131             : 
     132             :     int TestCapability(const char *) const override;
     133             : };
     134             : 
     135             : /************************************************************************/
     136             : /*                        OGRFeatherWriterLayer                         */
     137             : /************************************************************************/
     138             : 
     139             : class OGRFeatherWriterLayer final : public OGRArrowWriterLayer
     140             : 
     141             : {
     142             :     OGRFeatherWriterLayer(const OGRFeatherWriterLayer &) = delete;
     143             :     OGRFeatherWriterLayer &operator=(const OGRFeatherWriterLayer &) = delete;
     144             : 
     145             :     GDALDataset *m_poDS = nullptr;
     146             :     bool m_bStreamFormat = false;
     147             :     std::shared_ptr<arrow::ipc::RecordBatchWriter> m_poFileWriter{};
     148             :     std::shared_ptr<arrow::KeyValueMetadata> m_poFooterKeyValueMetadata{};
     149             : 
     150         626 :     bool IsFileWriterCreated() const override
     151             :     {
     152         626 :         return m_poFileWriter != nullptr;
     153             :     }
     154             : 
     155             :     void CreateWriter() override;
     156             :     bool CloseFileWriter() override;
     157             : 
     158             :     void CreateSchema() override;
     159             :     void PerformStepsBeforeFinalFlushGroup() override;
     160             : 
     161             :     bool FlushGroup() override;
     162             : 
     163         231 :     std::string GetDriverUCName() const override
     164             :     {
     165         231 :         return ARROW_DRIVER_NAME_UC;
     166             :     }
     167             : 
     168             :     virtual bool
     169             :     IsSupportedGeometryType(OGRwkbGeometryType eGType) const override;
     170             : 
     171           0 :     bool IsSRSRequired() const override
     172             :     {
     173           0 :         return true;
     174             :     }
     175             : 
     176             :     friend class OGRFeatherWriterDataset;
     177             :     bool Close();
     178             : 
     179             :   public:
     180             :     OGRFeatherWriterLayer(
     181             :         GDALDataset *poDS, arrow::MemoryPool *poMemoryPool,
     182             :         const std::shared_ptr<arrow::io::OutputStream> &poOutputStream,
     183             :         const char *pszLayerName);
     184             : 
     185             :     bool SetOptions(const std::string &osFilename, CSLConstList papszOptions,
     186             :                     const OGRSpatialReference *poSpatialRef,
     187             :                     OGRwkbGeometryType eGType);
     188             : 
     189             :     bool WriteArrowBatch(const struct ArrowSchema *schema,
     190             :                          struct ArrowArray *array,
     191             :                          CSLConstList papszOptions = nullptr) override;
     192             : 
     193         117 :     GDALDataset *GetDataset() override
     194             :     {
     195         117 :         return m_poDS;
     196             :     }
     197             : };
     198             : 
     199             : /************************************************************************/
     200             : /*                        OGRFeatherWriterDataset                       */
     201             : /************************************************************************/
     202             : 
     203             : class OGRFeatherWriterDataset final : public GDALPamDataset
     204             : {
     205             :     const std::string m_osFilename{};
     206             :     std::unique_ptr<arrow::MemoryPool> m_poMemoryPool{};
     207             :     std::unique_ptr<OGRFeatherWriterLayer> m_poLayer{};
     208             :     std::shared_ptr<arrow::io::OutputStream> m_poOutputStream{};
     209             : 
     210             :   public:
     211             :     explicit OGRFeatherWriterDataset(
     212             :         const char *pszFilename,
     213             :         const std::shared_ptr<arrow::io::OutputStream> &poOutputStream);
     214             : 
     215             :     ~OGRFeatherWriterDataset() override;
     216             : 
     217             :     CPLErr Close() override;
     218             : 
     219             :     arrow::MemoryPool *GetMemoryPool() const
     220             :     {
     221             :         return m_poMemoryPool.get();
     222             :     }
     223             : 
     224             :     int GetLayerCount() const override;
     225             :     const OGRLayer *GetLayer(int idx) const override;
     226             :     int TestCapability(const char *pszCap) const override;
     227             :     std::vector<std::string> GetFieldDomainNames(
     228             :         CSLConstList /*papszOptions*/ = nullptr) const override;
     229             :     const OGRFieldDomain *
     230             :     GetFieldDomain(const std::string &name) const override;
     231             :     bool AddFieldDomain(std::unique_ptr<OGRFieldDomain> &&domain,
     232             :                         std::string &failureReason) override;
     233             : 
     234             :   protected:
     235             :     OGRLayer *ICreateLayer(const char *pszName,
     236             :                            const OGRGeomFieldDefn *poGeomFieldDefn,
     237             :                            CSLConstList papszOptions) override;
     238             : };
     239             : 
     240             : #endif  // OGR_FEATHER_H