qgsarrowiterator.h

Go to the documentation of this file.

/***************************************************************************
    qgsarrowiterator.h
    ---------------------
    begin                : November 2025
    copyright            : (C) 2025 by Dewey Dunnington
    email                : dewey at dunnington dot ca
 ***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
 
#ifndef QGSARROWITERATOR_H
#define QGSARROWITERATOR_H
 
#include "qgis_core.h"
#include "qgis_sip.h"
#include "qgsfeatureiterator.h"
#include "qgsvectorlayer.h"
 
#ifndef ARROW_C_DATA_INTERFACE
#define ARROW_C_DATA_INTERFACE
 
#define ARROW_FLAG_DICTIONARY_ORDERED 1
#define ARROW_FLAG_NULLABLE 2
#define ARROW_FLAG_MAP_KEYS_SORTED 4
 
#ifndef SIP_RUN
struct ArrowSchema
{
    // Array type description
    const char *format;
    const char *name;
    const char *metadata;
    int64_t flags;
    int64_t n_children;
    struct ArrowSchema **children;
    struct ArrowSchema *dictionary;
 
    // Release callback
    void ( *release )( struct ArrowSchema * );
    // Opaque producer-specific data
    void *private_data;
};
#endif
 
#ifndef SIP_RUN
struct ArrowArray
{
    // Array data description
    int64_t length;
    int64_t null_count;
    int64_t offset;
    int64_t n_buffers;
    int64_t n_children;
    const void **buffers;
    struct ArrowArray **children;
    struct ArrowArray *dictionary;
 
    // Release callback
    void ( *release )( struct ArrowArray * );
    // Opaque producer-specific data
    void *private_data;
};
#endif
 
#endif // ARROW_C_DATA_INTERFACE
 
#ifndef ARROW_C_STREAM_INTERFACE
#define ARROW_C_STREAM_INTERFACE
 
#ifndef SIP_RUN
struct ArrowArrayStream
{
    // Callbacks providing stream functionality
    int ( *get_schema )( struct ArrowArrayStream *, struct ArrowSchema *out );
    int ( *get_next )( struct ArrowArrayStream *, struct ArrowArray *out );
    const char *( *get_last_error )( struct ArrowArrayStream * );
 
    // Release callback
    void ( *release )( struct ArrowArrayStream * );
 
    // Opaque producer-specific data
    void *private_data;
};
#endif
 
#endif // ARROW_C_STREAM_INTERFACE

class CORE_EXPORT QgsArrowInferSchemaOptions
{
  public:
    QgsArrowInferSchemaOptions();

    void setGeometryColumnName( const QString &geometryColumnName );

    QString geometryColumnName() const;
 
  private:
    QString mGeometryColumnName;
};

class CORE_EXPORT QgsArrowSchema
{
  public:
    QgsArrowSchema();

    QgsArrowSchema( const QgsArrowSchema &other );

    QgsArrowSchema &operator=( const QgsArrowSchema &other );
 
    ~QgsArrowSchema();
 
#ifndef SIP_RUN
    struct ArrowSchema *schema();

    const struct ArrowSchema *schema() const;
#endif

    unsigned long long cSchemaAddress() const;

    void exportToAddress( unsigned long long otherAddress );

    bool isValid() const;

    int geometryColumnIndex() const;

    void setGeometryColumnIndex( int geometryColumnIndex );
 
    // clang-format off
#ifdef SIP_RUN
    SIP_PYOBJECT __arrow_c_schema__();
    % MethodCode
    struct ArrowSchema *exportedSchema = static_cast<struct ArrowSchema *>( malloc( sizeof( struct ArrowSchema ) ) );
    if ( !exportedSchema )
    {
      PyErr_SetString( PyExc_MemoryError, "Failed to allocate ArrowSchema" );
      sipIsErr = 1;
    }
    else
    {
      memcpy( exportedSchema, sipCpp->schema(), sizeof( struct ArrowSchema ) );
      sipCpp->schema()->release = nullptr;
      sipRes = PyCapsule_New( exportedSchema, "arrow_schema", []( PyObject *capsule )
      {
        struct ArrowSchema *schema = static_cast<struct ArrowSchema *>( PyCapsule_GetPointer( capsule, "arrow_schema" ) );
        if ( schema && schema->release )
        {
          schema->release( schema );
        }
        free( schema );
      } );
      if ( !sipRes )
      {
        if ( exportedSchema->release )
        {
          exportedSchema->release( exportedSchema );
        }
        free( exportedSchema );
        sipIsErr = 1;
      }
    }
    % End

    static SIP_PYOBJECT fromArrow( SIP_PYOBJECT obj ) SIP_TYPEHINT( QgsArrowSchema );
    % MethodCode
    if ( PyCapsule_CheckExact( a0 ) && PyCapsule_IsValid( a0, "arrow_schema" ) ) {
      struct ArrowSchema *capsuleSchema = static_cast<struct ArrowSchema *>( PyCapsule_GetPointer( a0, "arrow_schema" ) );
      QgsArrowSchema *newSchema = new QgsArrowSchema();
      memcpy(newSchema->schema(), capsuleSchema, sizeof(struct ArrowSchema));
      capsuleSchema->release = nullptr;
      sipRes = sipConvertFromNewType( newSchema, sipType_QgsArrowSchema, nullptr );
    }
    else if ( PyObject_HasAttrString( a0, "__arrow_c_schema__" ) )
    {
      PyObject *method = PyObject_GetAttrString( a0, "__arrow_c_schema__" );
      if ( method )
      {
        PyObject *capsule = PyObject_CallObject( method, nullptr );
        Py_DECREF( method );
        if ( capsule )
        {
          if ( PyCapsule_CheckExact( capsule ) && PyCapsule_IsValid( capsule, "arrow_schema" ) )
          {
            struct ArrowSchema *capsuleSchema = static_cast<struct ArrowSchema *>( PyCapsule_GetPointer( capsule, "arrow_schema" ) );
            QgsArrowSchema *newSchema = new QgsArrowSchema();
            memcpy(newSchema->schema(), capsuleSchema, sizeof(struct ArrowSchema));
            capsuleSchema->release = nullptr;
            sipRes = sipConvertFromNewType( newSchema, sipType_QgsArrowSchema, nullptr );
          }
          else
          {
            PyErr_SetString( PyExc_TypeError, "__arrow_c_schema__() did not return a valid arrow_schema PyCapsule" );
            sipIsErr = 1;
          }
          Py_DECREF( capsule );
        }
        else
        {
          sipIsErr = 1; // Exception already set by PyObject_CallObject
        }
      }
      else
      {
        sipIsErr = 1; // Exception already set
      }
    }
    else
    {
      PyErr_Format( PyExc_TypeError, "Expected an object implementing __arrow_c_schema__(), got %s", Py_TYPE( a0 )->tp_name );
      sipIsErr = 1;
    }
    % End
#endif
 
  private:
    struct ArrowSchema mSchema {};
    int mGeometryColumnIndex = -1;
 
    // clang-format on
};

class CORE_EXPORT QgsArrowArray
{
  public:
    QgsArrowArray() = default;
 
    QgsArrowArray &operator=( QgsArrowArray &other ) = delete;
    QgsArrowArray( const QgsArrowArray &other ) = delete;
 
#ifndef SIP_RUN
    QgsArrowArray( QgsArrowArray &&other );

    QgsArrowArray &operator=( QgsArrowArray &&other );
#endif
 
    ~QgsArrowArray();
 
#ifndef SIP_RUN
    struct ArrowArray *array();

    const struct ArrowArray *array() const;
#endif

    unsigned long long cArrayAddress() const;

    void exportToAddress( unsigned long long otherAddress );

    bool isValid() const;
 
  private:
    struct ArrowArray mArray {};
 
#ifdef SIP_RUN
    QgsArrowArray( const QgsArrowArray &other );
#endif
};

class CORE_EXPORT QgsArrowArrayStream
{
  public:
    QgsArrowArrayStream() = default;
 
    QgsArrowArrayStream &operator=( QgsArrowArrayStream &other ) = delete;
    QgsArrowArrayStream( const QgsArrowArrayStream &other ) = delete;
 
#ifndef SIP_RUN
    QgsArrowArrayStream( QgsArrowArrayStream &&other );

    QgsArrowArrayStream &operator=( QgsArrowArrayStream &&other );
#endif
 
    ~QgsArrowArrayStream();
 
#ifndef SIP_RUN
    struct ArrowArrayStream *arrayStream();
#endif

    unsigned long long cArrayStreamAddress() const;

    void exportToAddress( unsigned long long otherAddress );

    bool isValid() const;
 
    // clang-format off
#ifdef SIP_RUN
    SIP_PYOBJECT __arrow_c_stream__( SIP_PYOBJECT requested_schema = Py_None );
    % MethodCode
    Q_UNUSED( a0 ); // requested_schema is not used but required by the protocol signature
    struct ArrowArrayStream *exportedStream = static_cast<struct ArrowArrayStream *>( malloc( sizeof( struct ArrowArrayStream ) ) );
    if ( !exportedStream )
    {
      PyErr_SetString( PyExc_MemoryError, "Failed to allocate ArrowArrayStream" );
      sipIsErr = 1;
    }
    else
    {
      memcpy( exportedStream, sipCpp->arrayStream(), sizeof( struct ArrowArrayStream ) );
      sipCpp->arrayStream()->release = nullptr;
      sipRes = PyCapsule_New( exportedStream, "arrow_array_stream", []( PyObject *capsule )
      {
        struct ArrowArrayStream *stream = static_cast<struct ArrowArrayStream *>( PyCapsule_GetPointer( capsule, "arrow_array_stream" ) );
        if ( stream && stream->release )
        {
          stream->release( stream );
        }
        free( stream );
      } );
      if ( !sipRes )
      {
        if ( exportedStream->release )
        {
          exportedStream->release( exportedStream );
        }
        free( exportedStream );
        sipIsErr = 1;
      }
    }
    % End

    static SIP_PYOBJECT fromArrow( SIP_PYOBJECT obj ) SIP_TYPEHINT( QgsArrowArrayStream );
    % MethodCode
    if ( PyCapsule_CheckExact( a0 ) && PyCapsule_IsValid( a0, "arrow_array_stream" ) )
    {
      struct ArrowArrayStream *capsuleStream = static_cast<struct ArrowArrayStream *>( PyCapsule_GetPointer( a0, "arrow_array_stream" ) );
      QgsArrowArrayStream *newStream = new QgsArrowArrayStream();
      memcpy(newStream->arrayStream(), capsuleStream, sizeof(struct ArrowArrayStream));
      capsuleStream->release = nullptr;
      sipRes = sipConvertFromNewType( newStream, sipType_QgsArrowArrayStream, nullptr );
    }
    else if ( PyObject_HasAttrString( a0, "__arrow_c_stream__" ) )
    {
      PyObject *method = PyObject_GetAttrString( a0, "__arrow_c_stream__" );
      if ( method )
      {
        PyObject *capsule = PyObject_CallObject( method, nullptr );
        Py_DECREF( method );
        if ( capsule )
        {
          if ( PyCapsule_CheckExact( capsule ) && PyCapsule_IsValid( capsule, "arrow_array_stream" ) )
          {
            struct ArrowArrayStream *capsuleStream = static_cast<struct ArrowArrayStream *>( PyCapsule_GetPointer( capsule, "arrow_array_stream" ) );
            QgsArrowArrayStream *newStream = new QgsArrowArrayStream();
            memcpy(newStream->arrayStream(), capsuleStream, sizeof(struct ArrowArrayStream));
            capsuleStream->release = nullptr;
            sipRes = sipConvertFromNewType( newStream, sipType_QgsArrowArrayStream, nullptr );
          }
          else
          {
            PyErr_SetString( PyExc_TypeError, "__arrow_c_stream__() did not return a valid arrow_array_stream PyCapsule" );
            sipIsErr = 1;
          }
          Py_DECREF( capsule );
        }
        else
        {
          sipIsErr = 1; // Exception already set by PyObject_CallObject
        }
      }
      else
      {
        sipIsErr = 1; // Exception already set
      }
    }
    else
    {
      PyErr_Format( PyExc_TypeError, "Expected an object implementing __arrow_c_stream__(), got %s", Py_TYPE( a0 )->tp_name );
      sipIsErr = 1;
    }
    % End
#endif
 
  private:
    struct ArrowArrayStream mArrayStream {};
 
#ifdef SIP_RUN
    QgsArrowArrayStream( const QgsArrowArrayStream &other );
#endif
 
    // clang-format on
};

class CORE_EXPORT QgsArrowIterator
{
  public:
    QgsArrowIterator() = default;

    explicit QgsArrowIterator( QgsFeatureIterator featureIterator );
 
#ifndef SIP_RUN
    struct ArrowSchema *schema();
#endif

    void setSchema( const QgsArrowSchema &schema );

    QgsArrowArrayStream toArrayStream( int batchSize = 65536 ) const;

    QgsArrowArray nextFeatures( int n ) SIP_THROW( QgsException );

    static QgsArrowSchema inferSchema( const QgsVectorLayer &layer, const QgsArrowInferSchemaOptions &options = QgsArrowInferSchemaOptions() ) SIP_THROW( QgsException );

    static QgsArrowSchema inferSchema(
      const QgsFields &fields, bool hasGeometry = false, const QgsCoordinateReferenceSystem &crs = QgsCoordinateReferenceSystem(), const QgsArrowInferSchemaOptions &options = QgsArrowInferSchemaOptions()
    ) SIP_THROW( QgsException );
 
  private:
    QgsFeatureIterator mFeatureIterator;
    QgsArrowSchema mSchema;
};
 
#endif // QGSARROWITERATOR_H