[Enh]: Add `Expr.str.json_decode`

[thomasaarholt]

Note

Edited by @dangotbanned

We would like to learn about your use case. For example, if this feature is needed to adopt Narwhals in an open source project, could you please enter the link to it below?

I'm trying to introduce narwhals at Microsoft. I'm a data scientist working there, analyzing huge data in Azure ML.

Please describe the purpose of the new feature or describe the problem to solve.

We often get (heavily!) nested json data structures in the dataframes we get through various APIs.

Polars and pyspark both have implementations to deal with this, and I've written a function that incorporates both:

Show function

from typing import TypeVar

import polars as pl
import pyspark.sql.functions as F
from pyspark.sql import DataFrame as SparkDataFrame
from pyspark.sql.types import StructType

Frame = TypeVar("Frame", pl.DataFrame, pl.LazyFrame, SparkDataFrame)


def decode_json_columns(
    df: Frame,
    column: str,
    new_column: str,
    json_schema: pl.Struct | StructType | None = None,
) -> Frame:
    """
    Decodes the 'content' column from JSON string to a struct with a defined schema.
    Args:
        df: Input DataFrame with a 'content' column containing JSON strings.
        schema: Polars Struct dtype or PySpark StructType for the JSON schema.
    """
    if isinstance(df, (pl.DataFrame, pl.LazyFrame)) and isinstance(json_schema, pl.Struct):
        return df.with_columns(pl.col(column).str.json_decode(dtype=json_schema).alias(new_column))
    elif isinstance(df, SparkDataFrame) and isinstance(json_schema, StructType):
        return df.withColumn(
            new_column,
            F.from_json(
                col=column,
                schema=json_schema,
                options={"mode": "PERMISSIVE"}, # these aren't necessary, but I thought I'd share this as food for thought
            ),
        )
    else:
        raise TypeError("Unsupported DataFrame type. Expected Polars or PySpark DataFrame.")

Minimized

import polars as pl

import sqlframe.standalone as sql
import sqlframe.standalone.functions as sql_f
import sqlframe.standalone.types as sql_types


def str_json_decode_polars(name: str, dtype: pl.DataType) -> pl.Expr:
    return pl.col(name).str.json_decode(dtype)


def str_json_decode_pyspark(
    name: str, dtype: sql_types.ArrayType | sql_types.StructType
) -> sql.Column:
    return sql_f.from_json(name, dtype)

This is a feature request for narwhals to support a similar API as polars.

It would be great to be able to provide the schema once, using nw.Schema, and thus avoid this boilerplate and having to pass the schema in per-implementation.

Suggest a solution if possible.

I haven't gotten used to the narwhals codebase yet, so don't have a solution, but here is at least a small example using the above function that might prove useful.
Note that I've had to add some os.environ and spark.config in order to get spark working on my Mac where I've written this example.

Show suggested solution

from __future__ import annotations

from typing import TypeVar

import polars as pl
import pyspark.sql.functions as F
from pyspark.sql import DataFrame as SparkDataFrame
from pyspark.sql import SparkSession
from pyspark.sql.types import IntegerType, StringType, StructField, StructType

Frame = TypeVar("Frame", pl.DataFrame, pl.LazyFrame, SparkDataFrame)
import os
import sys

os.environ["PYSPARK_PYTHON"] = sys.executable
os.environ["PYSPARK_DRIVER_PYTHON"] = sys.executable


def decode_json_columns(
    df: Frame,
    column: str,
    new_column: str,
    json_schema: pl.Struct | StructType | None = None,
) -> Frame:
    if isinstance(df, (pl.DataFrame, pl.LazyFrame)) and isinstance(json_schema, pl.Struct):
        return df.with_columns(pl.col(column).str.json_decode(dtype=json_schema).alias(new_column))
    elif isinstance(df, SparkDataFrame) and isinstance(json_schema, StructType):
        return df.withColumn(
            new_column,
            F.from_json(col=column, schema=json_schema, options={"mode": "PERMISSIVE"}),
        )
    else:
        raise TypeError("Unsupported DataFrame type. Expected Polars or PySpark DataFrame.")


def polars_example() -> None:
    df: pl.DataFrame = pl.DataFrame(
        {
            "id": [1, 2],
            "content": [
                '{"name": "Alice", "age": 30}',
                '{"name": "Bob", "age": 25}',
            ],
        }
    )

    schema: pl.Struct = pl.Struct(
        [
            pl.Field("name", pl.String),
            pl.Field("age", pl.Int64),
        ]
    )

    decoded: pl.DataFrame = decode_json_columns(
        df=df,
        column="content",
        new_column="content_struct",
        json_schema=schema,
    )

    print("\nPolars result")
    print(decoded)


def spark_example() -> None:
    spark: SparkSession = (
        SparkSession.builder.appName("app")
        .config(
            "spark.driver.extraJavaOptions",
            "--add-opens=java.base/java.nio=ALL-UNNAMED "
            "--add-opens=java.base/sun.nio.ch=ALL-UNNAMED "
            "--add-opens=java.base/java.lang=ALL-UNNAMED",
        )
        .config(
            "spark.executor.extraJavaOptions",
            "--add-opens=java.base/java.nio=ALL-UNNAMED "
            "--add-opens=java.base/sun.nio.ch=ALL-UNNAMED "
            "--add-opens=java.base/java.lang=ALL-UNNAMED",
        )
        .getOrCreate()
    )
    data: list[tuple[int, str]] = [
        (1, '{"name": "Alice", "age": 30}'),
        (2, '{"name": "Bob", "age": 25}'),
    ]

    df: SparkDataFrame = spark.createDataFrame(data, ["id", "content"])

    schema: StructType = StructType(
        [
            StructField("name", StringType(), True),
            StructField("age", IntegerType(), True),
        ]
    )

    decoded: SparkDataFrame = decode_json_columns(
        df=df,
        column="content",
        new_column="content_struct",
        json_schema=schema,
    )

    print("\nPySpark result")
    decoded.show(truncate=False)


polars_example()
spark_example()

[MarcoGorelli]

thanks for the request! on board with the idea

did you want to try submitting a PR as well? (no pressure to do so of course, just checking)

[thomasaarholt]

I'll give myself the deadline of trying to write it today. I'll probably use copilot quite a bit though. Let me make a PR and we can see how it goes.

[dangotbanned]

@thomasaarholt if you haven't started yet, could you put together a list of what backends could support str.json_decode please? 🙏
(See #3402 for an example under Backend support)

It is good if pyspark and polars support can solve your use-case 😄
But it would be even better to know we'd be introducing something that has more than 2 greens

[thomasaarholt]

Using this example for polars:

import polars as pl

df: pl.DataFrame = pl.DataFrame(
    {
        "id": [1, 2, 3],
        "payload": [
            '{"name": "Alice", "age": 30}',
            '{"name": "Bob", "age": 25}',
            '{"name": "Charlie", "age": 40}',
        ],
    }
)

payload_dtype: pl.DataType = pl.Struct(
    {
        "name": pl.String,
        "age": pl.Int64,
    }
)

result: pl.DataFrame = df.with_columns(
    pl.col("payload").str.json_decode(payload_dtype)
)

Arrow:
No support

dask:
Only via .map(json.loads)

duckdb:
"from_json(payload, '{name: VARCHAR, age: BIGINT}') AS payload"

ibis:
Seems to have a unwrap_as method in their expressions api

Pandas:
Only via .map(json.loads)

pyspark:

schema = StructType(
    [
        StructField("name", StringType()),
        StructField("age", LongType()),
    ]
)

df = df.withColumn("payload", from_json(col("payload"), schema))

[thomasaarholt]

Any thoughts on how the API would look to support the backends that would use json.loads? Since they don't decode into a particular type, just whatever python defaults to. I guess for pandas we could decode and then cast the column to a particular type?

[MarcoGorelli]

If pandas/dask require map and a Python function which operates element-by-element, then I think it should not be supported for those backends

If this is supportable by polars/duckdb/pyspark/ibis then i think that's enough

[dangotbanned]

Agreed on not using json.loads

pyarrow does have a canonical JSON extension type though, so there might be something there for pyarrow and (partially) pandas?

• https://arrow.apache.org/docs/format/CanonicalExtensions.html#json

[thomasaarholt]

I think we need to support the nw.Schema.to_X before can implement the json_decode method like polars does it, since the method takes an argument dtype. Unless we want to allow the user to specify the dtype using Native syntax?

[dangotbanned]

I think we need to support the nw.Schema.to_X before can implement the json_decode method like polars does it, since the method takes an argument dtype

Ah my bad, I wasn't very clear on this on discord

We already have the functionality for the DType conversion internally for all backends.
I was asking more if you'd still need json_decode, if we exposed a convenient way to convert the schema?

I'm not against adding the feature btw 😅

Unless we want to allow the user to specify the dtype using Native syntax?

Probably best to stick to one type system (narwhals), otherwise we'd need to check a pretty big laundry list of types on every call

[dangotbanned]

#3506 (comment)

@thomasaarholt
duckdb:
"from_json(payload, '{name: VARCHAR, age: BIGINT}') AS payload"

Just spotted that duckdb support for from_json in sqlframe seems to be missing.
We use that in CI as it runs a lot quicker than pyspark natively, so we may want to make a request upstream for it?

The good news is that (eakmanrq/sqlframe#330 (comment)) makes it sound like it just needs a nudge to get moving again 😄

[thomasaarholt]

I added a feature request there now!

I was asking more if you'd still need json_decode, if we exposed a convenient way to convert the schema?

I might be mistaken here, but I think you're saying "if we have a way to convert from narwhals schema to pyspark schema, can you just add the type of the column to the schema and parse the column at read-time?"

No, json_decode is still required, here's an example that reads in, note that "foo" is has been json-dumped, so when write_json is called, that column is doubly json-nested.

from io import StringIO
import polars as pl

df = pl.DataFrame(
    {
        "name": ["alice", "bob"],
        "foo": [
            json.dumps({"id": 1, "active": True}),
            json.dumps({"id": 2, "active": False}),
        ],
    }
)


schema = {
    "name": pl.String,
    "foo": pl.String,
}
schematized_df = pl.read_json(StringIO(df.write_json()), schema=schema) # this works

from io import StringIO
import polars as pl

foo_dtype = pl.Struct(
    [
        pl.Field("id", pl.Int64),
        pl.Field("active", pl.Boolean),
    ]
)
schema = {
    "name": pl.String,
    "foo": foo_dtype, # <-- this doesn't work
}

schematized_df = pl.read_json(StringIO(df.write_json()), schema=schema) # error

At read-time the column is read as a string type, and we'd need some additional magic, like a "JSON" type, from polars. I could imagine something like:

foo_dtype = pl.Struct(
    [
        pl.Field("id", pl.Int64),
        pl.Field("active", pl.Boolean),
    ]
)

schema = {
    "name": pl.String,
    "foo": pl.JSON[foo_dtype], # <-- this doesn't exist
}