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 :     virtual 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             :     virtual 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             :     OGRFeatherLayer(OGRFeatherDataset *poDS, const char *pszLayerName,
     102             :                     std::shared_ptr<arrow::io::RandomAccessFile> poFile,
     103             :                     bool bSeekable, const arrow::ipc::IpcReadOptions &oOptions,
     104             :                     std::shared_ptr<arrow::ipc::RecordBatchStreamReader>
     105             :                         &poRecordBatchStreamReader);
     106             : 
     107             :     void ResetReading() override;
     108             :     int TestCapability(const char *pszCap) override;
     109             :     GIntBig GetFeatureCount(int bForce) override;
     110             :     const char *GetMetadataItem(const char *pszName,
     111             :                                 const char *pszDomain = "") override;
     112             :     char **GetMetadata(const char *pszDomain = "") override;
     113             : 
     114             :     GDALDataset *GetDataset() override;
     115             : 
     116             :     std::unique_ptr<OGRFieldDomain> BuildDomain(const std::string &osDomainName,
     117             :                                                 int iFieldIndex) const override;
     118             : };
     119             : 
     120             : /************************************************************************/
     121             : /*                         OGRFeatherDataset                            */
     122             : /************************************************************************/
     123             : 
     124             : class OGRFeatherDataset final : public OGRArrowDataset
     125             : {
     126             :   public:
     127             :     explicit OGRFeatherDataset(
     128             :         const std::shared_ptr<arrow::MemoryPool> &poMemoryPool);
     129             : 
     130             :     int TestCapability(const char *) override;
     131             : };
     132             : 
     133             : /************************************************************************/
     134             : /*                        OGRFeatherWriterLayer                         */
     135             : /************************************************************************/
     136             : 
     137             : class OGRFeatherWriterLayer final : public OGRArrowWriterLayer
     138             : 
     139             : {
     140             :     OGRFeatherWriterLayer(const OGRFeatherWriterLayer &) = delete;
     141             :     OGRFeatherWriterLayer &operator=(const OGRFeatherWriterLayer &) = delete;
     142             : 
     143             :     GDALDataset *m_poDS = nullptr;
     144             :     bool m_bStreamFormat = false;
     145             :     std::shared_ptr<arrow::ipc::RecordBatchWriter> m_poFileWriter{};
     146             :     std::shared_ptr<arrow::KeyValueMetadata> m_poFooterKeyValueMetadata{};
     147             : 
     148         626 :     virtual bool IsFileWriterCreated() const override
     149             :     {
     150         626 :         return m_poFileWriter != nullptr;
     151             :     }
     152             : 
     153             :     virtual void CreateWriter() override;
     154             :     virtual bool CloseFileWriter() override;
     155             : 
     156             :     virtual void CreateSchema() override;
     157             :     virtual void PerformStepsBeforeFinalFlushGroup() override;
     158             : 
     159             :     virtual bool FlushGroup() override;
     160             : 
     161         231 :     virtual std::string GetDriverUCName() const override
     162             :     {
     163         231 :         return ARROW_DRIVER_NAME_UC;
     164             :     }
     165             : 
     166             :     virtual bool
     167             :     IsSupportedGeometryType(OGRwkbGeometryType eGType) const override;
     168             : 
     169           0 :     virtual bool IsSRSRequired() const override
     170             :     {
     171           0 :         return true;
     172             :     }
     173             : 
     174             :   public:
     175             :     OGRFeatherWriterLayer(
     176             :         GDALDataset *poDS, arrow::MemoryPool *poMemoryPool,
     177             :         const std::shared_ptr<arrow::io::OutputStream> &poOutputStream,
     178             :         const char *pszLayerName);
     179             : 
     180             :     ~OGRFeatherWriterLayer() override;
     181             : 
     182             :     bool SetOptions(const std::string &osFilename, CSLConstList papszOptions,
     183             :                     const OGRSpatialReference *poSpatialRef,
     184             :                     OGRwkbGeometryType eGType);
     185             : 
     186             :     bool WriteArrowBatch(const struct ArrowSchema *schema,
     187             :                          struct ArrowArray *array,
     188             :                          CSLConstList papszOptions = nullptr) override;
     189             : 
     190          10 :     GDALDataset *GetDataset() override
     191             :     {
     192          10 :         return m_poDS;
     193             :     }
     194             : };
     195             : 
     196             : /************************************************************************/
     197             : /*                        OGRFeatherWriterDataset                       */
     198             : /************************************************************************/
     199             : 
     200             : class OGRFeatherWriterDataset final : public GDALPamDataset
     201             : {
     202             :     const std::string m_osFilename{};
     203             :     std::unique_ptr<arrow::MemoryPool> m_poMemoryPool{};
     204             :     std::unique_ptr<OGRFeatherWriterLayer> m_poLayer{};
     205             :     std::shared_ptr<arrow::io::OutputStream> m_poOutputStream{};
     206             : 
     207             :   public:
     208             :     explicit OGRFeatherWriterDataset(
     209             :         const char *pszFilename,
     210             :         const std::shared_ptr<arrow::io::OutputStream> &poOutputStream);
     211             : 
     212             :     arrow::MemoryPool *GetMemoryPool() const
     213             :     {
     214             :         return m_poMemoryPool.get();
     215             :     }
     216             : 
     217             :     int GetLayerCount() override;
     218             :     OGRLayer *GetLayer(int idx) override;
     219             :     int TestCapability(const char *pszCap) override;
     220             :     std::vector<std::string> GetFieldDomainNames(
     221             :         CSLConstList /*papszOptions*/ = nullptr) const override;
     222             :     const OGRFieldDomain *
     223             :     GetFieldDomain(const std::string &name) const override;
     224             :     bool AddFieldDomain(std::unique_ptr<OGRFieldDomain> &&domain,
     225             :                         std::string &failureReason) override;
     226             : 
     227             :   protected:
     228             :     OGRLayer *ICreateLayer(const char *pszName,
     229             :                            const OGRGeomFieldDefn *poGeomFieldDefn,
     230             :                            CSLConstList papszOptions) override;
     231             : };
     232             : 
     233             : #endif  // OGR_FEATHER_H