fmt nested list in DataFrame with values

pola-rs · Sep 5, 2020 · 361166c · 361166c
1 parent a5505ab
commit 361166c
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 14 deletions.
diff --git a/polars/src/chunked_array/take.rs b/polars/src/chunked_array/take.rs
@@ -52,7 +52,7 @@ where
         impl_take_random_get_unchecked!(self, index, PrimitiveArray<T>)
     }
 }
-
+// extra trait such that it also works without extra reference.
 impl<'a> TakeRandomUtf8 for &'a Utf8Chunked {
     type Item = &'a str;
 

diff --git a/polars/src/fmt.rs b/polars/src/fmt.rs
@@ -274,12 +274,58 @@ impl Display for AnyType<'_> {
             }
             AnyType::IntervalDayTime(v) => write!(f, "{}", v),
             AnyType::IntervalYearMonth(v) => write!(f, "{}", v),
-            AnyType::LargeList(s) => write!(f, "list [{:?}]", s.dtype()),
+            AnyType::LargeList(s) => write!(f, "{:?}", s.fmt_largelist()),
             _ => unimplemented!(),
         }
     }
 }
 
+macro_rules! impl_fmt_largelist {
+    ($self:ident) => {{
+        match $self.len() {
+            1 => format!("[{:?}]", $self.get(0)),
+            2 => format!("[{:?}, {:?}]", $self.get(0), $self.get(1)),
+            3 => format!(
+                "[{:?}, {:?}, {:?}]",
+                $self.get(0),
+                $self.get(1),
+                $self.get(2)
+            ),
+            _ => format!(
+                "[{:?}, {:?}, ... {:?}]",
+                $self.get(0),
+                $self.get(1),
+                $self.get($self.len() - 1)
+            ),
+        }
+    }};
+}
+
+pub(crate) trait FmtLargeList {
+    fn fmt_largelist(&self) -> String;
+}
+
+impl<T> FmtLargeList for ChunkedArray<T>
+where
+    T: ArrowPrimitiveType,
+{
+    fn fmt_largelist(&self) -> String {
+        impl_fmt_largelist!(self)
+    }
+}
+
+impl FmtLargeList for Utf8Chunked {
+    fn fmt_largelist(&self) -> String {
+        impl_fmt_largelist!(self)
+    }
+}
+
+impl FmtLargeList for LargeListChunked {
+    fn fmt_largelist(&self) -> String {
+        impl_fmt_largelist!(self)
+    }
+}
+
 #[cfg(all(test, feature = "temporal"))]
 mod test {
     use crate::prelude::*;

diff --git a/polars/src/frame/group_by.rs b/polars/src/frame/group_by.rs
@@ -698,7 +698,7 @@ impl<'df, 'selection_str> GroupBy<'df, 'selection_str> {
     /// | date32     | u32        |
     /// +============+============+
     /// | 2020-08-23 | 1          |
-    /// +------------+------------T: PolarsNumericT+
+    /// +------------+------------+
     /// | 2020-08-22 | 2          |
     /// +------------+------------+
     /// | 2020-08-21 | 2          |
@@ -736,17 +736,17 @@ impl<'df, 'selection_str> GroupBy<'df, 'selection_str> {
     /// Returns:
     ///
     /// ```text
-    ///  +------------+---------------+
-    ///  | date       | temp_agg_list |
-    ///  | ---        | ---           |
-    ///  | date32     | list [i32]    |
-    ///  +============+===============+
-    ///  | 2020-08-23 | list [Int32]  |
-    ///  +------------+---------------+
-    ///  | 2020-08-22 | list [Int32]  |
-    ///  +------------+---------------+
-    ///  | 2020-08-21 | list [Int32]  |
-    ///  +------------+---------------+
+    /// +------------+------------------------+
+    /// | date       | temp_agg_list          |
+    /// | ---        | ---                    |
+    /// | date32     | list [i32]             |
+    /// +============+========================+
+    /// | 2020-08-23 | "[Some(9)]"            |
+    /// +------------+------------------------+
+    /// | 2020-08-22 | "[Some(7), Some(1)]"   |
+    /// +------------+------------------------+
+    /// | 2020-08-21 | "[Some(20), Some(10)]" |
+    /// +------------+------------------------+
     /// ```
     pub fn agg_list(&self) -> Result<DataFrame> {
         macro_rules! impl_gb {

diff --git a/polars/src/series/mod.rs b/polars/src/series/mod.rs
@@ -10,6 +10,7 @@ mod comparison;
 pub(crate) mod iterator;
 use arithmetic::{LhsNumOpsDispatch, NumOpsDispatch, NumOpsDispatchSeriesSingleNumber};
 // needed for enum dispatch
+use crate::fmt::FmtLargeList;
 use arrow::array::ArrayDataRef;
 use enum_dispatch::enum_dispatch;
 use num::{Num, NumCast};
@@ -749,6 +750,10 @@ impl Series {
     pub fn fill_none(&self, strategy: FillNoneStrategy) -> Result<Self> {
         Ok(apply_method_all_series_and_return!(self, fill_none, [strategy],?))
     }
+
+    pub(crate) fn fmt_largelist(&self) -> String {
+        apply_method_all_series!(self, fmt_largelist,)
+    }
 }
 
 fn pack_ca_to_series<N: PolarsDataType>(ca: ChunkedArray<N>) -> Series {