Handle premature loop termination

Author: gsakkis

pdal/StreamableExecutor.cpp

@@ -138,6 +138,13 @@ void PythonPointTable::reset()
     }
 }
 
+void PythonPointTable::disable()
+{
+    // TODO: uncomment the next line when/if StreamPointTable.m_capacity
+    // changes from private to protected
+    // m_capacity = 0;
+}
+
 void PythonPointTable::done()
 {
     m_arrays.push(nullptr);
@@ -198,14 +205,28 @@ PyArrayObject *StreamableExecutor::executeNext()
     // Blocks until something is ready.
     PyArrayObject *arr = m_table.fetchArray();
     if (arr == nullptr)
+        done();
+    return arr;
+}
+
+void StreamableExecutor::stop()
+{
+    if (m_thread)
     {
-        Py_BEGIN_ALLOW_THREADS
-        m_thread->join();
-        Py_END_ALLOW_THREADS
-        m_thread.reset();
-        m_executed = true;
+        m_table.disable();
+        while (PyArrayObject* arr = m_table.fetchArray())
+            Py_XDECREF(arr);
+        done();
     }
-    return arr;
+}
+
+void StreamableExecutor::done()
+{
+    Py_BEGIN_ALLOW_THREADS
+    m_thread->join();
+    Py_END_ALLOW_THREADS
+    m_thread.reset();
+    m_executed = true;
 }
 
 } // namespace python

pdal/StreamableExecutor.hpp

@@ -59,6 +59,7 @@ class PythonPointTable : public StreamPointTable
     ~PythonPointTable();
 
     virtual void finalize();
+    void disable();
     void done();
     PyArrayObject *fetchArray();
 
@@ -88,8 +89,11 @@ class StreamableExecutor : public PipelineExecutor
     ~StreamableExecutor();
 
     PyArrayObject* executeNext();
+    void stop();
 
 private:
+    void done();
+
     PythonPointTable m_table;
     std::unique_ptr<std::thread> m_thread;
 };

pdal/libpdalpython.pyx

@@ -107,6 +107,7 @@ cdef extern from "StreamableExecutor.hpp" namespace "pdal::python":
     cdef cppclass StreamableExecutor(PipelineExecutor):
         StreamableExecutor(string, int, int) except +
         np.PyArrayObject* executeNext() except +
+        void stop() except +
 
 
 cdef extern from "PyArray.hpp" namespace "pdal::python":
@@ -216,13 +217,16 @@ cdef class Pipeline:
 
     def __iter__(self):
         cdef StreamableExecutor* executor = self._get_executor()
-        while True:
-            arr_ptr = executor.executeNext()
-            if arr_ptr is NULL:
-                break
-            arr = <object>arr_ptr
-            Py_DECREF(arr)
-            yield arr
+        try:
+            while True:
+                arr_ptr = executor.executeNext()
+                if arr_ptr is NULL:
+                    break
+                arr = <object>arr_ptr
+                Py_DECREF(arr)
+                yield arr
+        finally:
+            executor.stop()
 
     #========= non-public properties & methods ===========================================
 

test/test_pipeline.py

@@ -484,3 +484,15 @@ def test_merged_arrays(self):
                                                   chunk_size=chunk_size))
             np.testing.assert_array_equal(np.concatenate(streaming_arrays),
                                           non_streaming_array)
+
+    def test_premature_exit(self):
+        """Can we stop iterating before all arrays are fetched"""
+        r = get_pipeline("range.json")
+        r.execute()
+        assert len(r.arrays) == 1
+        array = r.arrays[0]
+
+        ri = get_pipeline("range.json", chunk_size=100)
+        for array2 in ri:
+            np.testing.assert_array_equal(array2, array[:len(array2)])
+            break