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

Code before:
                            }
                        }
                    }
                    all_ready &= self.input_heads[recv_idx].ready_to_send();
                }

                if !all_ready {
                    // One or more of the input heads is exhausted (this phase).
                    break;
                }

                // TODO: recombine morsels to make sure the concatenation is
                // close to the ideal morsel size.

                // Compute common size and send a combined morsel.
                let Some(common_size) = self.input_heads.iter().flat_map(|h| h.min_len()).min()
                else {
                    // If all input heads are broadcasts we don't get a common size,
                    // we handle this below.
                    break;
                };

                for input_head in &mut self.input_heads {
                    out.push(input_head.take(common_size));
                }
                let strict_concat = matches!(self.extend_behavior, ExtendBehavior::Raise);
                let out_df = concat_df_horizontal(&out, false, strict_concat)?;
                out.clear();

                let morsel = Morsel::new(out_df, self.out_seq, source_token.clone());
                self.out_seq = self.out_seq.successor();
                if sender.send(morsel).await.is_err() {
                    // Our receiver is no longer interested in any data, no
                    // need store the rest of the incoming stream, can directly
                    // return.
                    return Ok(());
                }
            }

            // We can't continue because one or more input heads is empty or all
            // inputs are broadcasts. We must tell everyone to stop, unblock all
            // pipes by consuming all ConsumeTokens, and then store all data
            // that was still flowing through the pipelines into input_heads for
            // the next phase.
            for input_head in &mut self.input_heads {
                for morsel in &mut input_head.morsels {
                    morsel.source_token().stop();
                    drop(morsel.take_consume_token());
                }
            }


Review comment (suggestion):
This match is unnecessary, you can always enable strict concat here. Our inputs at this point in the code should always have the same length.

Code after:
                            }
                        }
                    }
                    all_ready &= self.input_heads[recv_idx].ready_to_send();
                }

                if !all_ready {
                    // One or more of the input heads is exhausted (this phase).
                    break;
                }

                // TODO: recombine morsels to make sure the concatenation is
                // close to the ideal morsel size.

                // Compute common size and send a combined morsel.
                let Some(common_size) = self.input_heads.iter().flat_map(|h| h.min_len()).min()
                else {
                    // If all input heads are broadcasts we don't get a common size,
                    // we handle this below.
                    break;
                };

                for input_head in &mut self.input_heads {
                    out.push(input_head.take(common_size));
                }
                let out_df = concat_df_horizontal(&out, false, true)?;
                out.clear();

                let morsel = Morsel::new(out_df, self.out_seq, source_token.clone());
                self.out_seq = self.out_seq.successor();
                if sender.send(morsel).await.is_err() {
                    // Our receiver is no longer interested in any data, no
                    // need store the rest of the incoming stream, can directly
                    // return.
                    return Ok(());
                }
            }

            // We can't continue because one or more input heads is empty or all
            // inputs are broadcasts. We must tell everyone to stop, unblock all
            // pipes by consuming all ConsumeTokens, and then store all data
            // that was still flowing through the pipelines into input_heads for
            // the next phase.
            for input_head in &mut self.input_heads {
                for morsel in &mut input_head.morsels {
                    morsel.source_token().stop();
                    drop(morsel.take_consume_token());
                }
            }

            for (recv_idx, opt_recv) in receivers.iter_mut().enumerate() {