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

Code before:
                        lambda
                            .call1(py, (PyErr::from(err),))
                            .map_err(|err| err.restore(py))
                            .ok();
                    },
                });
            });
        })
    }

    #[cfg(feature = "async")]
    fn collect_batches(
        &self,
        py: Python<'_>,
        engine: Wrap<Engine>,
        maintain_order: bool,
        chunk_size: Option<NonZeroUsize>,
        lazy: bool,
    ) -> PyResult<PyCollectBatches> {
        py.enter_polars(|| {
            let mut ldf = self.ldf.read().clone();
            // let schema = ldf.collect_schema()?;
            let schema = ldf
                .collect_schema()
                .map_err(PyPolarsErr::from)?
                .to_arrow(CompatLevel::newest());
            let collect_batches = ldf
                .collect_batches(engine.0, maintain_order, chunk_size, lazy)
                .map_err(PyPolarsErr::from)?;

            PyResult::Ok(PyCollectBatches {
                inner: Arc::new(Mutex::new(collect_batches)),
                schema,
            })
        })
    }

    #[cfg(feature = "parquet")]
    #[pyo3(signature = (
        target, sink_options, compression, compression_level, statistics, row_group_size, data_page_size,
        metadata, field_overwrites,
    ))]
    fn sink_parquet(
        &self,
        py: Python<'_>,
        target: PyFileSinkDestination,
        sink_options: PySinkOptions,
        compression: &str,
        compression_level: Option<i32>,
        statistics: Wrap<StatisticsOptions>,
        row_group_size: Option<usize>,

Review comment (performance):
Please don't collect the schema here. Instead clone the `ldf` and store it in the `PyCollectBatches` so we can collect it later, if needed.

Code after:
                        lambda
                            .call1(py, (PyErr::from(err),))
                            .map_err(|err| err.restore(py))
                            .ok();
                    },
                });
            });
        })
    }

    #[cfg(feature = "async")]
    fn collect_batches(
        &self,
        py: Python<'_>,
        engine: Wrap<Engine>,
        maintain_order: bool,
        chunk_size: Option<NonZeroUsize>,
        lazy: bool,
    ) -> PyResult<PyCollectBatches> {
        py.enter_polars(|| {
            let ldf = self.ldf.read().clone();

            let collect_batches = ldf
                .clone()
                .collect_batches(engine.0, maintain_order, chunk_size, lazy)
                .map_err(PyPolarsErr::from)?;

            PyResult::Ok(PyCollectBatches {
                inner: Arc::new(Mutex::new(collect_batches)),
                ldf,
            })
        })
    }

    #[cfg(feature = "parquet")]
    #[pyo3(signature = (
        target, sink_options, compression, compression_level, statistics, row_group_size, data_page_size,
        metadata, field_overwrites,
    ))]
    fn sink_parquet(
        &self,
        py: Python<'_>,
        target: PyFileSinkDestination,
        sink_options: PySinkOptions,
        compression: &str,
        compression_level: Option<i32>,
        statistics: Wrap<StatisticsOptions>,
        row_group_size: Option<usize>,
        data_page_size: Option<usize>,
        metadata: Wrap<Option<KeyValueMetadata>>,
        field_overwrites: Vec<Wrap<ParquetFieldOverwrites>>,