feat(python): Allow custom JSONEncoder for the json_normalize function, minor speedup

[alexander-beedie]

Closes #20950.

• Allows a custom json encoder function to be passed to json_normalize, enabling the caller to target specific performance/behaviour.
• Minor speedup (~6% faster on test data) from optimising Python loops 🚀
• No change to default behaviour (which uses json.dumps).
• Additional JSON test coverage.

Example

import polars as pl

data = [
    {
        "id": n,
        "name": {"first": "Some", "middle": None, "last": "Name"},
        "attrs": {"height": 185, "weight": 90},
        "lifts": {
            "max_deadlift": 190, 
            "max_squat": 140, 
            "max_bench": 100,
        },
    } for n in range(5)
]

Default:

(if not specified, uses json.dumps)

pl.json_normalize(data, max_level=0)
# shape: (5, 4)
# ┌─────┬─────────────────────────────────┬───────────────────────────────┬─────────────────────────────────┐
# │ id  ┆ name                            ┆ attrs                         ┆ lifts                           │
# │ --- ┆ ---                             ┆ ---                           ┆ ---                             │
# │ i64 ┆ str                             ┆ str                           ┆ str                             │
# ╞═════╪═════════════════════════════════╪═══════════════════════════════╪═════════════════════════════════╡
# │ 0   ┆ {"first": "Some", "middle": nu… ┆ {"height": 185, "weight": 90} ┆ {"max_deadlift": 190, "max_squ… │
# │ 1   ┆ {"first": "Some", "middle": nu… ┆ {"height": 185, "weight": 90} ┆ {"max_deadlift": 190, "max_squ… │
# │ 2   ┆ {"first": "Some", "middle": nu… ┆ {"height": 185, "weight": 90} ┆ {"max_deadlift": 190, "max_squ… │
# │ 3   ┆ {"first": "Some", "middle": nu… ┆ {"height": 185, "weight": 90} ┆ {"max_deadlift": 190, "max_squ… │
# │ 4   ┆ {"first": "Some", "middle": nu… ┆ {"height": 185, "weight": 90} ┆ {"max_deadlift": 190, "max_squ… │
# └─────┴─────────────────────────────────┴───────────────────────────────┴─────────────────────────────────┘

Custom encoder:

(use orjson.dumps as the encoder instead; note that output is now binary)

import orjson
pl.json_normalize(data, max_level=0, encoder=orjson.dumps)
# shape: (5, 4)
# ┌─────┬─────────────────────────────────┬───────────────────────────────┬─────────────────────────────────┐
# │ id  ┆ name                            ┆ attrs                         ┆ lifts                           │
# │ --- ┆ ---                             ┆ ---                           ┆ ---                             │
# │ i64 ┆ binary                          ┆ binary                        ┆ binary                          │
# ╞═════╪═════════════════════════════════╪═══════════════════════════════╪═════════════════════════════════╡
# │ 0   ┆ b"{"first":"Some","middle":nul… ┆ b"{"height":185,"weight":90}" ┆ b"{"max_deadlift":190,"max_squ… │
# │ 1   ┆ b"{"first":"Some","middle":nul… ┆ b"{"height":185,"weight":90}" ┆ b"{"max_deadlift":190,"max_squ… │
# │ 2   ┆ b"{"first":"Some","middle":nul… ┆ b"{"height":185,"weight":90}" ┆ b"{"max_deadlift":190,"max_squ… │
# │ 3   ┆ b"{"first":"Some","middle":nul… ┆ b"{"height":185,"weight":90}" ┆ b"{"max_deadlift":190,"max_squ… │
# │ 4   ┆ b"{"first":"Some","middle":nul… ┆ b"{"height":185,"weight":90}" ┆ b"{"max_deadlift":190,"max_squ… │
# └─────┴─────────────────────────────────┴───────────────────────────────┴─────────────────────────────────┘