[Python] Polars read memory usage issue

I'm trying to flatten a DynamoDB JSON file with Polars.

My JSON file is 1GB, 100mo after compressing to gz but in memory when I read it, it uses 16GB.

Because the data comes from DynamoDB, my columns can have multiple types. For example, 1 year ago it was a String, today it's a boolean.

So I made 4 experiments with only a pl.read_ndjson and write_parquet:

• Read a 5GB json, compressed in gz file with 5 levels nested and non fix schema between lines -> it explodes my 32GB memory
• Read 5GB json, compressed in gz file with 5 levels nested and fix schema between lines -> it explodes my memory too
• Read 5GB json, compressed in gz file without nested and fix schema between lines -> it uses 16GB to read it
• Read 5GB json file without nested and fix schema between lines -> it uses 6GB to read it

Here my script for simple data:

import json
import random
from uuid import uuid4
from datetime import datetime

DYNAMO_TYPES = ["S", "N", "BOOL", "M", "L"]

def random_dynamodb_value(i, depth=0):
if depth > 0:
return {"S": str(uuid4())}

value_type = DYNAMO_TYPES[i%5]

if value_type == "S":
return {"S": f"val-{uuid4().hex[:8]}"}
elif value_type == "N":
return {"N": str(random.randint(0, 1000))}
elif value_type == "BOOL":
return {"BOOL": random.choice([True, False])}

def generate_extra_easy_fields(n=50):
extra = {}
for i in range(n):
field_name = f"field_{i}"
extra[field_name] = random_dynamodb_value(i)
return extra

def generate_easy_entry():
base = {
"id": {"S": str(uuid4())},
"timestamp": {"S": datetime.now().isoformat()},
}

base.update(generate_extra_easy_fields(50))
return base

def generate_jsonl_file(filename: str, n: int = 20):
with open(filename, "w") as f:
for _ in range(n):
entry = generate_easy_entry()
f.write(json.dumps(entry) + "\n")

if __name__ == "__main__":
generate_jsonl_file("dynamodb_simple.json", n=2000000)


Here my script for complexe data :

import json
import random
from uuid import uuid4
from datetime import datetime

DYNAMO_TYPES = ["S", "N", "BOOL", "M", "L"]

def random_dynamodb_value(i, depth=0):
if depth > 2:
return {"S": str(uuid4())}

value_type = DYNAMO_TYPES[i%5]

if value_type == "S":
return {"S": f"val-{uuid4().hex[:8]}"}
elif value_type == "N":
return {"N": str(random.randint(0, 1000))}
elif value_type == "BOOL":
return {"BOOL": random.choice([True, False])}
elif value_type == "M":
return {
"M": {
f"key_{i}": random_dynamodb_value(i, depth + 1)
for i in range(random.randint(1, 3))
}
}
elif value_type == "L":
return {
"L": [random_dynamodb_value(depth + 1) for _ in range(random.randint(2, 4))]
}

def generate_multitype_field():
return random.choice([
{"S": "text"},
{"BOOL": random.choice([True, False])},
{"N": str(random.randint(0, 100))},
])

def generate_extra_fields(n=50):
extra = {}
for i in range(n):
field_name = f"field_{i}"
if i % 10 == 0:
extra[field_name] = random_dynamodb_value(i)
return extra

def generate_entry():
base = {
"id": {"S": str(uuid4())},
"timestamp": {"S": datetime.now().isoformat()},
"meta": {"M": {"level1": {"M": {"level2": {"M": {"level3": {"L": [random_dynamodb_value(3) for _ in range(2)]}}}}}}},
"listOfStructs": {
"L": [
{
"M": {
"subId": {"N": str(i)},
"flag": {"BOOL": random.choice([True, False])}
}
} for i in range(random.randint(2, 4))
]
},
"structOfStructOfStruct": {"M": {"a": {"M": {"b": {"M": {"c": random_dynamodb_value(2)}}}}}},
"structOfStructOfList": {"M": {"outer": {"M": {"innerList": {"L": [random_dynamodb_value(2) for _ in range(3)]}}}}},
"multiTypeField": generate_multitype_field()
}

base.update(generate_extra_fields(50))
return base

def generate_jsonl_file(filename: str, n: int = 20):
with open(filename, "w") as f:
for _ in range(n):
entry = generate_entry()
f.write(json.dumps(entry) + "\n")

if __name__ == "__main__":
generate_jsonl_file("dynamodb_complexe.json", n=2000000)

Continue reading...