diff --git a/crates/polars-core/src/config.rs b/crates/polars-core/src/config.rs index d94113724406..14d6130b3c6f 100644 --- a/crates/polars-core/src/config.rs +++ b/crates/polars-core/src/config.rs @@ -13,6 +13,7 @@ pub(crate) const FMT_TABLE_HIDE_DATAFRAME_SHAPE_INFORMATION: &str = pub(crate) const FMT_TABLE_INLINE_COLUMN_DATA_TYPE: &str = "POLARS_FMT_TABLE_INLINE_COLUMN_DATA_TYPE"; pub(crate) const FMT_TABLE_ROUNDED_CORNERS: &str = "POLARS_FMT_TABLE_ROUNDED_CORNERS"; +pub(crate) const FMT_TABLE_CELL_LIST_LEN: &str = "POLARS_FMT_TABLE_CELL_LIST_LEN"; // Other env vars #[cfg(feature = "dtype-decimal")] diff --git a/crates/polars-core/src/fmt.rs b/crates/polars-core/src/fmt.rs index 59c69f8cdc95..e480007cd8e0 100644 --- a/crates/polars-core/src/fmt.rs +++ b/crates/polars-core/src/fmt.rs @@ -1,6 +1,6 @@ #[cfg(any(feature = "fmt", feature = "fmt_no_tty"))] use std::borrow::Cow; -use std::fmt::{Debug, Display, Formatter}; +use std::fmt::{Debug, Display, Formatter, Write}; use std::sync::atomic::{AtomicU8, Ordering}; use std::{fmt, str}; @@ -949,12 +949,32 @@ impl Series { self.get(1).unwrap(), self.get(2).unwrap() ), - _ => format!( - "[{}, {}, … {}]", - self.get(0).unwrap(), - self.get(1).unwrap(), - self.get(self.len() - 1).unwrap() - ), + _ => { + let max_items = std::env::var(FMT_TABLE_CELL_LIST_LEN) + .as_deref() + .unwrap_or("") + .parse() + .map_or(3, |n: i64| if n < 0 { self.len() } else { n as usize }); + + let mut result = "[".to_owned(); + + for (i, item) in self.iter().enumerate() { + write!(result, "{item}").unwrap(); + + if i != self.len() - 1 { + result.push_str(", "); + } + + if i == max_items - 2 { + result.push_str("… "); + write!(result, "{}", self.get(self.len() - 1).unwrap()).unwrap(); + break; + } + } + result.push(']'); + + result + }, } } } @@ -1089,7 +1109,7 @@ mod test { fn test_fmt_list() { let mut builder = ListPrimitiveChunkedBuilder::::new("a", 10, 10, DataType::Int32); - builder.append_opt_slice(Some(&[1, 2, 3])); + builder.append_opt_slice(Some(&[1, 2, 3, 4, 5, 6])); builder.append_opt_slice(None); let list = builder.finish().into_series(); @@ -1097,11 +1117,23 @@ mod test { r#"shape: (2,) Series: 'a' [list[i32]] [ - [1, 2, 3] + [1, 2, … 6] null ]"#, format!("{:?}", list) ); + + std::env::set_var("POLARS_FMT_TABLE_CELL_LIST_LEN", "10"); + + assert_eq!( + r#"shape: (2,) +Series: 'a' [list[i32]] +[ + [1, 2, 3, 4, 5, 6] + null +]"#, + format!("{:?}", list) + ) } #[test] diff --git a/py-polars/polars/config.py b/py-polars/polars/config.py index 2ca55f2bd123..80f471714e38 100644 --- a/py-polars/polars/config.py +++ b/py-polars/polars/config.py @@ -69,6 +69,7 @@ def _get_float_fmt() -> str: # pragma: no cover "POLARS_FMT_TABLE_HIDE_DATAFRAME_SHAPE_INFORMATION", "POLARS_FMT_TABLE_INLINE_COLUMN_DATA_TYPE", "POLARS_FMT_TABLE_ROUNDED_CORNERS", + "POLARS_FMT_TABLE_CELL_LIST_LEN", "POLARS_STREAMING_CHUNK_SIZE", "POLARS_TABLE_WIDTH", "POLARS_VERBOSE", @@ -527,6 +528,54 @@ def set_fmt_str_lengths(cls, n: int | None) -> type[Config]: os.environ["POLARS_FMT_STR_LEN"] = str(n) return cls + @classmethod + def set_fmt_table_cell_list_len(cls, n: int | None) -> type[Config]: + """ + Set the number of elements to display for List values. + + Values less than 0 will result in all values being printed. + + Parameters + ---------- + n : int + Number of values to display. + + Examples + -------- + >>> df = pl.DataFrame( + ... { + ... "nums": [ + ... [1, 2, 3, 4, 5, 6], + ... ] + ... } + ... ) + >>> df + shape: (1, 1) + ┌─────────────┐ + │ nums │ + │ --- │ + │ list[i64] │ + ╞═════════════╡ + │ [1, 2, … 6] │ + └─────────────┘ + >>> with pl.Config(fmt_table_cell_list_len=10): + ... print(df) + ... + shape: (1, 1) + ┌────────────────────┐ + │ nums │ + │ --- │ + │ list[i64] │ + ╞════════════════════╡ + │ [1, 2, 3, 4, 5, 6] │ + └────────────────────┘ + """ + if n is None: + os.environ.pop("POLARS_FMT_TABLE_CELL_LIST_LEN", None) + else: + os.environ["POLARS_FMT_TABLE_CELL_LIST_LEN"] = str(n) + return cls + @classmethod def set_streaming_chunk_size(cls, size: int | None) -> type[Config]: """