From dc4ab3d7ba472716933e8700e0ef6fb37b9bf8be Mon Sep 17 00:00:00 2001
From: Joshua Arnott <josh@snorfalorpagus.net>
Date: Sat, 24 Mar 2018 16:02:21 +0000
Subject: [PATCH] Support for WKB as geometry type in read_vectorized

---
 fiona/_vectorized.pyx    | 12 +++++++++++-
 tests/test_vectorized.py | 10 ++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/fiona/_vectorized.pyx b/fiona/_vectorized.pyx
index fdc86f170..84fa37e02 100644
--- a/fiona/_vectorized.pyx
+++ b/fiona/_vectorized.pyx
@@ -1,6 +1,7 @@
 from .ogrext cimport Session, _deleteOgrFeature
 from .ogrext import FIELD_TYPES, FIELD_TYPES_MAP, OGRERR_NONE
 from ._shim cimport *
+from libc.stdlib cimport malloc, free
 from fiona.rfc3339 import FionaDateType, FionaDateTimeType, FionaTimeType
 
 import logging
@@ -12,7 +13,7 @@ cimport numpy as np
 
 log = logging.getLogger(__name__)
 
-def read_vectorized(collection):
+def read_vectorized(collection, use_wkb=False):
     cdef Session session
     cdef void * cogr_feature
     cdef void * cogr_geometry
@@ -35,6 +36,7 @@ def read_vectorized(collection):
     cdef long long [:] arr_int
     cdef double [:] arr_double
     cdef char * wkt
+    cdef char * wkb
 
     session = collection.session
     encoding = session._fileencoding
@@ -161,6 +163,14 @@ def read_vectorized(collection):
             cogr_geometry = OGR_F_GetGeometryRef(cogr_feature)
             if cogr_geometry == NULL:
                 data_geometry[feature_index] = None
+            elif use_wkb:
+                length = OGR_G_WkbSize(cogr_geometry)
+                wkb = <char*>malloc(sizeof(char)*length)
+                result = OGR_G_ExportToWkb(cogr_geometry, 1, wkb)
+                if result != OGRERR_NONE:
+                    raise ValueError("Failed to export geometry to WKB")
+                data_geometry[feature_index] = wkb[:length]
+                free(wkb)
             else:
                 result = OGR_G_ExportToWkt(cogr_geometry, &wkt)
                 if result != OGRERR_NONE:
diff --git a/tests/test_vectorized.py b/tests/test_vectorized.py
index 17adf6053..248f724ea 100644
--- a/tests/test_vectorized.py
+++ b/tests/test_vectorized.py
@@ -103,3 +103,13 @@ def test_datetime_fields(tmpdir):
         assert features["properties"]["date"][0] == np.datetime64("2018-03-24")
         assert features["properties"]["datetime"][0] == np.datetime64("2018-03-24T15:06:01")
         assert str(features["properties"]["nulldt"][0]) == "NaT"
+
+def test_wkb(path_coutwildrnp_shp):
+    with fiona.open(path_coutwildrnp_shp, "r") as collection:
+        features = read_vectorized(collection, use_wkb=True)
+
+    geometry = features["geometry"][0]
+    assert geometry[0:1] == b"\x01"  # little endian
+    assert geometry[1:5] == b"\x03\x00\x00\x00"  # polygon
+    assert geometry[5:9] == b"\x01\x00\x00\x00"  # 1 ring
+    assert len(geometry) == 1325