Repository: pola-rs/polars
File: crates/polars-python/src/lazyframe/general.rs (Rust)

Code before:
    }
}

pub struct LazyFrameStreamIterator {
    inner: Arc<Mutex<CollectBatches>>,
    dtype: ArrowDataType,
}

impl LazyFrameStreamIterator {
    fn new(inner: Arc<Mutex<CollectBatches>>, schema: ArrowDataType) -> Self {
        Self {
            inner,
            dtype: schema,
        }
    }

    fn field(&self) -> ArrowField {
        ArrowField::new(PlSmallStr::EMPTY, self.dtype.clone(), false)
    }
}

impl Iterator for LazyFrameStreamIterator {
    type Item = PolarsResult<ArrayRef>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.inner.lock().next() {
            None => None,
            Some(df) => {
                let mut df = df.unwrap();
                df.rechunk_mut();
                let batch_cols = df
                    .columns()
                    .iter()
                    .map(|v| v.as_materialized_series().clone())
                    .collect::<Vec<Series>>()
                    .iter()
                    .map(|s| s.to_arrow(0, CompatLevel::newest()))
                    .collect::<Vec<_>>();

                let array = arrow::array::StructArray::new(
                    self.dtype.clone(),
                    batch_cols[0].len(),
                    batch_cols,
                    None,
                );
                Some(Ok(Box::new(array)))
            },
        }
    }
}


Review comment (suggestion):
Can you do
```rust
let next = self.inner.lock().next()
```
before the match? I believe this holds the lock for the entirety of the match body.

Code after:
        let schema = ldf
            .collect_schema()
            .map_err(PyPolarsErr::from)?
            .to_arrow(CompatLevel::newest());

        let schema = ArrowDataType::Struct(schema.clone().into_iter_values().collect());

        let iter = Box::new(ArrowStreamIterator::new(self.inner.clone(), schema));
        let field = iter.field();
        let stream = export_iterator(iter, field);
        let stream_capsule_name = CString::new("arrow_array_stream").unwrap();
        PyCapsule::new(py, stream, Some(stream_capsule_name))
    }
}

pub struct ArrowStreamIterator {
    inner: Arc<Mutex<CollectBatches>>,
    dtype: ArrowDataType,
}

impl ArrowStreamIterator {
    fn new(inner: Arc<Mutex<CollectBatches>>, schema: ArrowDataType) -> Self {
        Self {
            inner,
            dtype: schema,
        }
    }

    fn field(&self) -> ArrowField {
        ArrowField::new(PlSmallStr::EMPTY, self.dtype.clone(), false)
    }
}

impl Iterator for ArrowStreamIterator {
    type Item = PolarsResult<ArrayRef>;

    fn next(&mut self) -> Option<Self::Item> {
        let next = self.inner.lock().next();
        match next {
            None => None,
            Some(Err(err)) => Some(Err(err)),
            Some(Ok(df)) => {
                let height = df.height();
                let arrays = df.rechunk_into_arrow(CompatLevel::newest());
                Some(Ok(Box::new(arrow::array::StructArray::new(
                    self.dtype.clone(),
                    height,
                    arrays,
                    None,
                ))))
            },