Repository: pola-rs/polars
File: crates/polars-stream/src/nodes/io_sinks2/writers/ipc/record_batch_encoder.rs (Rust)

Code before:
    pub arrow_converters: Vec<ToArrowConverter>,
    pub dictionary_id_offsets: Arc<[usize]>,
    pub write_options: WriteOptions,
    // Unstable.
    pub write_statistics_flags: bool,
}

impl RecordBatchEncoder {
    pub async fn run(self) -> PolarsResult<()> {
        let RecordBatchEncoder {
            mut morsel_rx,
            ipc_batch_tx,
            mut arrow_converters,
            dictionary_id_offsets,
            write_options,
            write_statistics_flags,
        } = self;

        let mut record_batch_arrow_arrays: Vec<Box<dyn Array>> =
            Vec::with_capacity(arrow_converters.len());

        while let Ok(morsel) = morsel_rx.recv().await {
            let (df, permit) = morsel.into_inner();
            let height = df.height();
            let columns = df.into_columns();
            let flags = write_statistics_flags.then_some(
                columns
                    .iter()
                    .map(|c| c.get_flags().bits())
                    .collect::<Vec<_>>(),
            );
            let custom_metadata = flags.map(|flags| {
                vec![schema::key_value(
                    IPC_RW_RECORD_BATCH_FLAGS_KEY,
                    serde_json::to_string(&flags).unwrap(),
                )]
            });

            assert!(record_batch_arrow_arrays.is_empty());
            assert_eq!(arrow_converters.len(), columns.len());

            // Rechunk and convert to arrow in parallel.
            for fut in parallelize_first_to_local(
                TaskPriority::High,
                columns.into_iter().zip(arrow_converters.drain(..)).map(
                    |(column, mut arrow_converter)| async move {
                        let rechunked = column.take_materialized_series().rechunk();
                        let dtype = rechunked.dtype();

                        let array: Box<dyn Array> =
                            arrow_converter.array_to_arrow(rechunked.chunks()[0].as_ref(), dtype);

Review comment (suggestion):
I missed this earlier, but `then_some()` will always run - this should be `then(|| )` instead

Code after:
    pub arrow_converters: Vec<ToArrowConverter>,
    pub dictionary_id_offsets: Arc<[usize]>,
    pub write_options: WriteOptions,
    // Unstable.
    pub write_statistics_flags: bool,
}

impl RecordBatchEncoder {
    pub async fn run(self) -> PolarsResult<()> {
        let RecordBatchEncoder {
            mut morsel_rx,
            ipc_batch_tx,
            mut arrow_converters,
            dictionary_id_offsets,
            write_options,
            write_statistics_flags,
        } = self;

        let mut record_batch_arrow_arrays: Vec<Box<dyn Array>> =
            Vec::with_capacity(arrow_converters.len());

        while let Ok(morsel) = morsel_rx.recv().await {
            let (df, permit) = morsel.into_inner();
            let height = df.height();
            let columns = df.into_columns();
            let flags = write_statistics_flags.then(|| {
                columns
                    .iter()
                    .map(|c| c.get_flags().bits())
                    .collect::<Vec<_>>()
            });
            let custom_metadata = flags.map(|flags| {
                vec![schema::key_value(
                    IPC_RW_RECORD_BATCH_FLAGS_KEY,
                    serde_json::to_string(&flags).unwrap(),
                )]
            });

            assert!(record_batch_arrow_arrays.is_empty());
            assert_eq!(arrow_converters.len(), columns.len());

            // Rechunk and convert to arrow in parallel.
            for fut in parallelize_first_to_local(
                TaskPriority::High,
                columns.into_iter().zip(arrow_converters.drain(..)).map(
                    |(column, mut arrow_converter)| async move {
                        let rechunked = column.take_materialized_series().rechunk();
                        let dtype = rechunked.dtype();

                        let array: Box<dyn Array> =
                            arrow_converter.array_to_arrow(rechunked.chunks()[0].as_ref(), dtype);