use polars_error::{PolarsResult, polars_bail};
use crate::frame::column::Column;
pub(super) fn infer_broadcast_height(columns: &[Column]) -> usize {
if columns.is_empty() {
return 0;
}
columns
.iter()
.map(|c| c.len())
.find(|len| *len != 1)
.unwrap_or(1)
}
pub(super) fn broadcast_columns(height: usize, columns: &mut [Column]) -> PolarsResult<()> {
for col in columns.iter_mut() {
let len = col.len();
if len != height {
if len != 1 {
let name = col.name().clone();
let extra_info = if let Some(c) = columns.iter().find(|c| c.len() == height) {
format!(" (matching column '{}')", c.name())
} else {
String::new()
};
polars_bail!(
ShapeMismatch:
"could not create a new DataFrame: \
series {name:?} has length {len} \
while trying to broadcast to length {height}{extra_info}",
);
}
*col = col.new_from_index(0, height);
}
}
Ok(())
}