feat: Add JSON encoding for polars list type: pl.Expr.list.json_encode()

crates/polars-plan/src/dsl/function_expr/list.rs

@@ -56,6 +56,8 @@ pub enum ListFunction {
     Join(bool),
     #[cfg(feature = "dtype-array")]
     ToArray(usize),
+    #[cfg(feature = "json")]
+    JsonEncode,
 }
 
 impl ListFunction {
@@ -103,6 +105,8 @@ impl ListFunction {
             #[cfg(feature = "dtype-array")]
             ToArray(width) => mapper.try_map_dtype(|dt| map_list_dtype_to_array_dtype(dt, *width)),
             NUnique => mapper.with_dtype(IDX_DTYPE),
+            #[cfg(feature = "json")]
+            JsonEncode => mapper.with_dtype(DataType::String),
         }
     }
 }
@@ -174,6 +178,8 @@ impl Display for ListFunction {
             Join(_) => "join",
             #[cfg(feature = "dtype-array")]
             ToArray(_) => "to_array",
+            #[cfg(feature = "json")]
+            JsonEncode => "to_json",
         };
         write!(f, "list.{name}")
     }
@@ -235,6 +241,8 @@ impl From<ListFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
             #[cfg(feature = "dtype-array")]
             ToArray(width) => map!(to_array, width),
             NUnique => map!(n_unique),
+            #[cfg(feature = "json")]
+            JsonEncode => map!(to_json),
         }
     }
 }
@@ -641,3 +649,17 @@ pub(super) fn to_array(s: &Series, width: usize) -> PolarsResult<Series> {
 pub(super) fn n_unique(s: &Series) -> PolarsResult<Series> {
     Ok(s.list()?.lst_n_unique()?.into_series())
 }
+
+#[cfg(feature = "json")]
+pub(super) fn to_json(s: &Series) -> PolarsResult<Series> {
+    let ca = s.list()?;
+
+    let dtype = ca.dtype().to_arrow(CompatLevel::newest());
+
+    let iter = ca.chunks().iter().map(|arr| {
+        let arr = arrow::compute::cast::cast_unchecked(arr.as_ref(), &dtype).unwrap();
+        polars_json::json::write::serialize_to_utf8(arr.as_ref())
+    });
+
+    Ok(StringChunked::from_chunk_iter(ca.name(), iter).into_series())
+}

crates/polars-plan/src/dsl/list.rs

@@ -406,4 +406,10 @@ impl ListNameSpace {
         let other = other.into();
         self.set_operation(other, SetOperation::SymmetricDifference)
     }
+
+    #[cfg(feature = "json")]
+    pub fn json_encode(self) -> Expr {
+        self.0
+            .map_private(FunctionExpr::ListExpr(ListFunction::JsonEncode))
+    }
 }

crates/polars-python/src/expr/list.rs

@@ -253,4 +253,9 @@ impl PyExpr {
         }
         .into()
     }
+
+    #[cfg(feature = "json")]
+    fn list_json_encode(&self) -> Self {
+        self.inner.clone().list().json_encode().into()
+    }
 }

py-polars/polars/expr/list.py

@@ -1358,3 +1358,40 @@ def set_symmetric_difference(self, other: IntoExpr) -> Expr:
         """  # noqa: W505.
         other = parse_into_expression(other, str_as_lit=False)
         return wrap_expr(self._pyexpr.list_set_operation(other, "symmetric_difference"))
+
+    def json_encode(self) -> Expr:
+        r"""
+        Convert this list to a string column with json values.
+
+        Examples
+        --------
+        >>> pl.DataFrame({"a": [[1, 2], [45], [9, 1, 3], None]}).with_columns(
+        ...     pl.col("a").list.json_encode().alias("encoded")
+        ... )
+        shape: (4, 2)
+        ┌───────────┬─────────┐
+        │ a         ┆ encoded │
+        │ ---       ┆ ---     │
+        │ list[i64] ┆ str     │
+        ╞═══════════╪═════════╡
+        │ [1, 2]    ┆ [1,2]   │
+        │ [45]      ┆ [45]    │
+        │ [9, 1, 3] ┆ [9,1,3] │
+        │ null      ┆ null    │
+        └───────────┴─────────┘
+
+        >>> pl.DataFrame({"a": [["\\", '"foo"'], [None, ""]]}).with_columns(
+        ...     pl.col("a").list.json_encode().alias("encoded")
+        ... )
+        shape: (2, 2)
+        ┌────────────────┬──────────────────┐
+        │ a              ┆ encoded          │
+        │ ---            ┆ ---              │
+        │ list[str]      ┆ str              │
+        ╞════════════════╪══════════════════╡
+        │ ["\", ""foo""] ┆ ["\\","\"foo\""] │
+        │ [null, ""]     ┆ [null,""]        │
+        └────────────────┴──────────────────┘
+
+        """
+        return wrap_expr(self._pyexpr.list_json_encode())

py-polars/polars/series/list.py

@@ -1052,3 +1052,21 @@ def set_symmetric_difference(self, other: Series) -> Series:
             [5, 7, 8]
         ]
         """  # noqa: W505
+
+    def json_encode(self) -> Series:
+        """
+        Convert this list Series into a string Series with json values.
+
+        Examples
+        --------
+        >>> a = pl.Series([[1, 2, 3], [], [None, 3], [5, 6, 7]])
+        >>> a.list.json_encode()
+        shape: (4,)
+        Series: '' [str]
+        [
+            "[1,2,3]"
+            "[]"
+            "[null,3]"
+            "[5,6,7]"
+        ]
+        """

py-polars/tests/unit/operations/namespaces/list/test_list.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import math
 from datetime import date, datetime
 
 import numpy as np
@@ -917,3 +918,64 @@ def test_list_eval_type_cast_11188() -> None:
     assert df.select(
         pl.col("a").list.eval(pl.element().cast(pl.String)).alias("a_str")
     ).schema == {"a_str": pl.List(pl.String)}
+
+
+def test_list_json_encode() -> None:
+    df = pl.DataFrame(
+        {
+            "a": [[1, None, 3], [4, 5, 6], None],
+            "b": [
+                [
+                    {
+                        "foo": 1,
+                    },
+                    {"foo": 2},
+                    {
+                        "foo": 3,
+                    },
+                ],
+                [
+                    {
+                        "foo": 3,
+                    },
+                    {"foo": 4},
+                    {
+                        "foo": 5,
+                    },
+                ],
+                [
+                    {
+                        "foo": 6,
+                    },
+                    {"foo": 7},
+                    {
+                        "foo": 8,
+                    },
+                ],
+            ],
+            "c": [[True, False, True], [False, True, False], [True, True, False]],
+            "d": [[1.6, 2.7, 3.8], [math.inf, -math.inf, 6.1], [7.2, 8.3, math.nan]],
+        }
+    )
+    df = df.with_columns(
+        pl.col("a").list.json_encode(),
+        pl.col("b").list.json_encode(),
+        pl.col("c").list.json_encode(),
+        pl.col("d").list.json_encode(),
+    )
+    assert df.schema == {
+        "a": pl.String,
+        "b": pl.String,
+        "c": pl.String,
+        "d": pl.String,
+    }
+    assert df.to_dict(as_series=False) == {
+        "a": ["[1,null,3]", "[4,5,6]", "null"],
+        "b": [
+            """[{"foo":1},{"foo":2},{"foo":3}]""",
+            """[{"foo":3},{"foo":4},{"foo":5}]""",
+            """[{"foo":6},{"foo":7},{"foo":8}]""",
+        ],
+        "c": ["[true,false,true]", "[false,true,false]", "[true,true,false]"],
+        "d": ["[1.6,2.7,3.8]", "[null,null,6.1]", "[7.2,8.3,null]"],
+    }