fix: fills in empty fields with NaN

hexa/datasets/migrations/0012_alter_datasetfilesample_sample.py

@@ -0,0 +1,24 @@
+# Generated by Django 4.2.16 on 2024-10-22 12:00
+
+from django.db import migrations, models
+
+import hexa.datasets.models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("datasets", "0011_alter_datasetfilesample_sample"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="datasetfilesample",
+            name="sample",
+            field=models.JSONField(
+                blank=True,
+                default=list,
+                encoder=hexa.datasets.models.DataframeJsonEncoder,
+                null=True,
+            ),
+        ),
+    ]

hexa/datasets/models.py

@@ -1,4 +1,5 @@
 import logging
+import math
 import secrets
 from functools import cached_property
 
@@ -350,6 +351,24 @@ class Meta:
         ordering = ["uri"]
 
 
+class DataframeJsonEncoder(DjangoJSONEncoder):
+    def encode(self, obj):
+        # Recursively replace NaN with None (since it's a float, it does not call 'default' method)
+        def replace_nan(item):
+            if isinstance(item, float) and math.isnan(item):
+                return None
+            elif isinstance(item, dict):
+                return {key: replace_nan(value) for key, value in item.items()}
+            elif isinstance(item, list):
+                return [replace_nan(element) for element in item]
+            return item
+
+        # Preprocess the object to replace NaN values
+        obj = replace_nan(obj)
+        # Use the superclass's encode method to serialize the preprocessed object
+        return super().encode(obj)
+
+
 class DatasetFileSample(Base):
     STATUS_PROCESSING = "PROCESSING"
     STATUS_FAILED = "FAILED"
@@ -360,7 +379,12 @@ class DatasetFileSample(Base):
         (STATUS_FAILED, _("Failed")),
         (STATUS_FINISHED, _("Finished")),
     ]
-    sample = JSONField(blank=True, default=list, null=True, encoder=DjangoJSONEncoder)
+    sample = JSONField(
+        blank=True,
+        default=list,
+        null=True,
+        encoder=DataframeJsonEncoder,
+    )
     status = models.CharField(
         max_length=10,
         choices=STATUS_CHOICES,

hexa/datasets/tests/fixtures/example_with_nan.csv

@@ -0,0 +1,4 @@
+name,surname,age,married
+Joe,Doe,10,True
+Liam,Smith,,False
+Emma,Johnson,,False

hexa/datasets/tests/test_generate_sample.py

@@ -7,7 +7,12 @@
 from pandas.errors import EmptyDataError
 
 from hexa.core.test import TestCase
-from hexa.datasets.models import Dataset, DatasetFileSample, DatasetVersionFile
+from hexa.datasets.models import (
+    DataframeJsonEncoder,
+    Dataset,
+    DatasetFileSample,
+    DatasetVersionFile,
+)
 from hexa.datasets.queue import (
     add_system_attributes,
     generate_sample,
@@ -21,6 +26,12 @@
 from hexa.workspaces.models import Workspace, WorkspaceMembershipRole
 
 
+class TestDataframeJsonEncoder(TestCase):
+    def test_default(self):
+        encoder = DataframeJsonEncoder()
+        self.assertEqual(encoder.encode({"a": float("nan")}), '{"a": null}')
+
+
 class TestCreateDatasetFileSampleTask(TestCase, DatasetTestMixin):
     @classmethod
     def setUpTestData(cls):
@@ -116,6 +127,21 @@ def test_generate_sample(
                 ],
                 None,
             ),
+            (
+                "example_with_nan.csv",
+                DatasetFileSample.STATUS_FINISHED,
+                [
+                    {
+                        "age": None,
+                        "name": "Liam",
+                        "married": False,
+                        "surname": "Smith",
+                    },
+                    {"age": 10.0, "name": "Joe", "married": True, "surname": "Doe"},
+                    {"age": 10.0, "name": "Joe", "married": True, "surname": "Doe"},
+                ],
+                None,
+            ),
         ]
         for (
             fixture_name,