Repository: pola-rs/polars
File: crates/polars-stream/src/nodes/io_sources/ndjson/line_batch_distributor.rs (Rust)

Code before:

        let fixed_read_size = std::env::var("POLARS_FORCE_NDJSON_CHUNK_SIZE")
            .map(|x| {
                x.parse::<NonZeroUsize>().ok().unwrap_or_else(|| {
                    panic!("invalid value for POLARS_FORCE_NDJSON_CHUNK_SIZE: {x}")
                })
            })
            .ok();

        if verbose {
            eprintln!(
                "\
                [NDJSON LineBatchDistributor]: \
                n_rows_to_skip: {}, \
                reverse: {reverse} \
                fixed_read_size: {fixed_read_size:?} \
                ",
                row_skipper.cfg_n_rows_to_skip,
            );
        }

        let mut full_input_opt = if reverse {
            // Since decompression doesn't support reverse decompression, we have to fully
            // decompress the input. It's crucial for the streaming property that this doesn't get
            // called in the non-reverse case.
            assert!(
                !reader.is_compressed(),
                "Negative slicing and decompression risk OOM, should be handled on higher level."
            );
            let (full_input, _) = reader.read_next_slice(&Buffer::new(), usize::MAX)?;
            let offset = full_input.len();
            Some((full_input, offset))
        } else {
            None
        };

        let mut read_size = fixed_read_size
            .map(NonZeroUsize::get)
            .unwrap_or_else(CompressedReader::initial_read_size);
        let mut prev_leftover = Buffer::new();
        let mut chunk_idx = 0;

        loop {
            let (mem_slice, bytes_read) = if reverse {
                let (full_input, offset) = full_input_opt.as_mut().unwrap();
                let new_offset = offset.saturating_sub(read_size);
                let bytes_read = *offset - new_offset;
                let new_slice = full_input
                    .clone()
                    .sliced(new_offset..(*offset + prev_leftover.len()));
                *offset = new_offset;

Review comment (bug):
```suggestion
            debug_assert!(
```
If we hit this state due to a logic bug it's better to at least attempt to continue instead of immediately panicking on user builds.


Code after:

        let fixed_read_size = std::env::var("POLARS_FORCE_NDJSON_CHUNK_SIZE")
            .map(|x| {
                x.parse::<NonZeroUsize>().ok().unwrap_or_else(|| {
                    panic!("invalid value for POLARS_FORCE_NDJSON_CHUNK_SIZE: {x}")
                })
            })
            .ok();

        if verbose {
            eprintln!(
                "\
                [NDJSON LineBatchDistributor]: \
                n_rows_to_skip: {}, \
                reverse: {reverse} \
                fixed_read_size: {fixed_read_size:?} \
                ",
                row_skipper.cfg_n_rows_to_skip,
            );
        }

        let mut full_input_opt = if reverse {
            // Since decompression doesn't support reverse decompression, we have to fully
            // decompress the input. It's crucial for the streaming property that this doesn't get
            // called in the non-reverse case.
            debug_assert!(
                !reader.is_compressed(),
                "Negative slicing and decompression risk OOM, should be handled on higher level."
            );
            let (full_input, _) = reader.read_next_slice(&Buffer::new(), usize::MAX)?;
            let offset = full_input.len();
            Some((full_input, offset))
        } else {
            None
        };

        let mut read_size = fixed_read_size
            .map(NonZeroUsize::get)
            .unwrap_or_else(CompressedReader::initial_read_size);
        let mut prev_leftover = Buffer::new();
        let mut chunk_idx = 0;

        loop {
            let (mem_slice, bytes_read) = if reverse {
                let (full_input, offset) = full_input_opt.as_mut().unwrap();
                let new_offset = offset.saturating_sub(read_size);
                let bytes_read = *offset - new_offset;
                let new_slice = full_input
                    .clone()
                    .sliced(new_offset..(*offset + prev_leftover.len()));
                *offset = new_offset;