ENH: support reading from in-memory buffers

pyogrio/_ogr.pxd

@@ -1,5 +1,6 @@
 # Contains declarations against GDAL / OGR API
 from libc.stdint cimport int64_t
+from libc.stdio cimport FILE
 
 
 cdef extern from "cpl_conv.h":
@@ -33,6 +34,38 @@ cdef extern from "cpl_string.h":
     char**      CSLAddString(char **list, const char *string)
 
 
+cdef extern from "sys/stat.h" nogil:
+    struct stat:
+        int st_mode
+
+
+cdef extern from "cpl_vsi.h" nogil:
+
+    ctypedef int vsi_l_offset
+    ctypedef FILE VSILFILE
+    ctypedef stat VSIStatBufL
+
+    unsigned char *VSIGetMemFileBuffer(const char *path,
+                                       vsi_l_offset *data_len,
+                                       int take_ownership)
+    VSILFILE *VSIFileFromMemBuffer(const char *path, void *data,
+                                   vsi_l_offset data_len, int take_ownership)
+    VSILFILE* VSIFOpenL(const char *path, const char *mode)
+    int VSIFCloseL(VSILFILE *fp)
+    int VSIUnlink(const char *path)
+    int VSIMkdir(const char *path, long mode)
+    int VSIRmdir(const char *path)
+    int VSIFFlushL(VSILFILE *fp)
+    size_t VSIFReadL(void *buffer, size_t nSize, size_t nCount, VSILFILE *fp)
+    char** VSIReadDir(const char* pszPath)
+    int VSIFSeekL(VSILFILE *fp, vsi_l_offset nOffset, int nWhence)
+    vsi_l_offset VSIFTellL(VSILFILE *fp)
+    int VSIFTruncateL(VSILFILE *fp, vsi_l_offset nNewSize)
+    size_t VSIFWriteL(void *buffer, size_t nSize, size_t nCount, VSILFILE *fp)
+    int VSIStatL(const char *pszFilename, VSIStatBufL *psStatBuf)
+    int VSI_ISDIR(int mode)
+
+
 cdef extern from "ogr_core.h":
     ctypedef enum OGRErr:
         OGRERR_NONE  # success

pyogrio/_ogr.pyx

@@ -1,3 +1,6 @@
+from uuid import uuid4
+
+
 cdef get_string(const char *c_str, str encoding="UTF-8"):
     """Get Python string from a char *
 
@@ -113,3 +116,23 @@ def ogr_list_drivers():
 
     return drivers
 
+
+def buffer_to_virtual_file(bytesbuf, ext=''):
+    """Maps a bytes buffer to a virtual file.
+    `ext` is empty or begins with a period and contains at most one period.
+    """
+
+    vsi_filename = '/vsimem/{}'.format(uuid4().hex + ext)
+
+    vsi_handle = VSIFileFromMemBuffer(vsi_filename.encode("utf8"), <unsigned char *>bytesbuf, len(bytesbuf), 0)
+
+    if vsi_handle == NULL:
+        raise OSError('failed to map buffer to file')
+    if VSIFCloseL(vsi_handle) != 0:
+        raise OSError('failed to close mapped file handle')
+
+    return vsi_filename
+
+
+def remove_virtual_file(vsi_filename):
+    return VSIUnlink(vsi_filename.encode("utf8"))

pyogrio/raw.py

@@ -5,6 +5,7 @@
 
 with GDALEnv():
     from pyogrio._io import ogr_read, ogr_read_info, ogr_list_layers, ogr_write
+    from pyogrio._ogr import buffer_to_virtual_file, remove_virtual_file
 
 
 DRIVERS = {
@@ -16,7 +17,7 @@
 
 
 def read(
-    path,
+    path_or_buffer,
     layer=None,
     encoding=None,
     columns=None,
@@ -34,7 +35,7 @@ def read(
 
     Parameters
     ----------
-    path : pathlib.Path or str
+    path_or_buffer : pathlib.Path or str
         data source path
     layer : int or str, optional (default: first layer)
         If an integer is provided, it corresponds to the index of the layer
@@ -84,19 +85,36 @@ def read(
             "geometry": "<geometry type>"
         }
     """
+    from_buffer = False
+    if isinstance(path_or_buffer, bytes):
+        from_buffer = True
+        ext = ""
+        if path_or_buffer[:4].startswith(b'PK\x03\x04'):
+            ext = ".zip"
+        path = buffer_to_virtual_file(path_or_buffer, ext=ext)
+        if path_or_buffer[:4].startswith(b'PK\x03\x04'):
+            path = "/vsizip/" + path
+    else:
+        path = str(path_or_buffer)
+
+    try:
+        result = ogr_read(
+            path,
+            layer=layer,
+            encoding=encoding,
+            columns=columns,
+            read_geometry=read_geometry,
+            force_2d=force_2d,
+            skip_features=skip_features,
+            max_features=max_features or 0,
+            where=where,
+            bbox=bbox,
+        )
+    finally:
+        if from_buffer:
+            remove_virtual_file(path)
 
-    return ogr_read(
-        str(path),
-        layer=layer,
-        encoding=encoding,
-        columns=columns,
-        read_geometry=read_geometry,
-        force_2d=force_2d,
-        skip_features=skip_features,
-        max_features=max_features or 0,
-        where=where,
-        bbox=bbox,
-    )
+    return result
 
 
 def write(

pyogrio/tests/conftest.py

@@ -34,7 +34,7 @@ def naturalearth_lowres_vsi(tmp_path, naturalearth_lowres):
             filename = f"{naturalearth_lowres.stem}.{ext}"
             out.write(naturalearth_lowres.parent / filename, filename)
 
-    return f"/vsizip/{path}/{naturalearth_lowres.name}"
+    return path, f"/vsizip/{path}/{naturalearth_lowres.name}"
 
 
 @pytest.fixture(scope="session")

pyogrio/tests/test_core.py

@@ -43,7 +43,7 @@ def test_list_layers(naturalearth_lowres, naturalearth_lowres_vsi, test_fgdb_vsi
     )
 
     assert array_equal(
-        list_layers(naturalearth_lowres_vsi), [["naturalearth_lowres", "Polygon"]]
+        list_layers(naturalearth_lowres_vsi[1]), [["naturalearth_lowres", "Polygon"]]
     )
 
     # Measured 3D is downgraded to 2.5D during read

pyogrio/tests/test_geopandas_io.py

@@ -40,7 +40,7 @@ def test_read_dataframe(naturalearth_lowres):
 
 
 def test_read_dataframe_vsi(naturalearth_lowres_vsi):
-    df = read_dataframe(naturalearth_lowres_vsi)
+    df = read_dataframe(naturalearth_lowres_vsi[1])
     assert len(df) == 177
 
 

pyogrio/tests/test_raw_io.py

@@ -34,6 +34,7 @@ def test_read(naturalearth_lowres):
 
 
 def test_vsi_read_layers(naturalearth_lowres_vsi):
+    _, naturalearth_lowres_vsi = naturalearth_lowres_vsi
     assert array_equal(
         list_layers(naturalearth_lowres_vsi), [["naturalearth_lowres", "Polygon"]]
     )
@@ -219,3 +220,49 @@ def test_write_unsupported(tmpdir, naturalearth_lowres):
 
     with pytest.raises(DriverError, match="does not support write functionality"):
         write(filename, geometry, field_data, driver="OpenFileGDB", **meta)
+
+
+def assert_equal_result(result1, result2):
+    meta1, geometry1, field_data1 = result1
+    meta2, geometry2, field_data2 = result2
+
+    assert np.array_equal(meta1["fields"], meta2["fields"])
+    # assert np.array_equal(geometry1, geometry2)
+    import pygeos
+    assert pygeos.equals_exact(
+        pygeos.from_wkb(geometry1), pygeos.from_wkb(geometry2), tolerance=0.00001
+    ).all()
+    assert all([np.array_equal(f1, f2) for f1, f2 in zip(field_data1, field_data2)])
+
+
+def test_read_from_bytes(tmpdir, naturalearth_lowres):
+    meta, geometry, field_data = read(naturalearth_lowres)
+    filename = os.path.join(str(tmpdir), "test.gpkg")
+    write(filename, geometry, field_data, driver="GPKG", **meta)
+
+    with open(filename, "rb") as f:
+        buffer = f.read()
+
+    meta2, geometry2, field_data2 = read(buffer)
+    assert_equal_result((meta, geometry, field_data), (meta2, geometry2, field_data2))
+
+
+    filename = os.path.join(str(tmpdir), "test.geojson")
+    write(filename, geometry, field_data, driver="GeoJSON", **meta)
+
+    with open(filename, "rb") as f:
+        buffer = f.read()
+
+    meta2, geometry2, field_data2 = read(buffer)
+    assert_equal_result((meta, geometry, field_data), (meta2, geometry2, field_data2))
+
+
+def test_read_from_bytes_zipped(tmpdir, naturalearth_lowres_vsi):
+    path, vsi_path = naturalearth_lowres_vsi
+    meta, geometry, field_data = read(vsi_path)
+
+    with open(path, "rb") as f:
+        buffer = f.read()
+
+    meta2, geometry2, field_data2 = read(buffer)
+    assert_equal_result((meta, geometry, field_data), (meta2, geometry2, field_data2))