JVM: keep NaN, inf and -inf in to_json_list and in nested values of the columnar readers

src/arcadedb/results.py

@@ -39,6 +39,39 @@
 }
 
 
+# NaN and the infinities: JSON has no literal for them and the engine's writers make null of them, so the
+# bridge (RowBatcher, ColumnBatcher) writes each one as {"\u0001": ["f", "NaN"]} (the text is Double.toString).
+_NONFINITE_KIND = "f"
+_NONFINITE_TEXT = frozenset(("NaN", "Infinity", "-Infinity"))
+
+
+def _nonfinite_hook(obj):
+    """The json object_hook that turns the bridge's non-finite tag back into a float."""
+    tagged = obj.get(_TYPED_TAG)
+    if (
+        isinstance(tagged, list)
+        and len(tagged) == 2
+        and tagged[0] == _NONFINITE_KIND
+        and tagged[1] in _NONFINITE_TEXT
+    ):
+        return float(tagged[1])
+    return obj
+
+
+def _loads_nonfinite(text):
+    """``json.loads`` of bridge JSON (a str or UTF-8 bytes) that may hold the non-finite tag.
+
+    The hook runs only when the tag character appears in the text (the writer
+    escapes it as ``\\u0001``), so a batch with no NaN or infinity parses at
+    the speed of the plain C decoder.
+    """
+    if isinstance(text, (bytes, bytearray, memoryview)):
+        text = bytes(text).decode("utf-8")
+    if "\\u0001" in text or _TYPED_TAG in text:
+        return json.loads(text, object_hook=_nonfinite_hook)
+    return json.loads(text)
+
+
 def _typed_hook(side):
     """The json object_hook that restores the values TypedRows tagged."""
 
@@ -535,6 +568,12 @@ def to_json_list(self, batch_size: int = 10_000) -> List[Dict[str, Any]]:
         floats. Use ``to_list()`` when you want ``date``, ``datetime`` and
         ``Decimal`` values: it costs about the same.
 
+        NaN, Infinity and -Infinity arrive as Python floats (``float('nan')``,
+        ``float('inf')``), at the top level and inside a list or a map, as on
+        the native build. JSON has no literal for them, so the bridge writes
+        each as a tag that is turned back into a float here: the result is not
+        strict JSON, and ``json.dumps(..., allow_nan=False)`` of it raises.
+
         The encoding is one rule for every row, at the top level and inside a
         list or a map: a DATE is the epoch milliseconds of its UTC midnight (not
         a count of days) whatever the JVM time zone, the same number
@@ -571,8 +610,6 @@ def iter_json_batches(
         type semantics; bounds memory to one batch. Falls back to chunked
         per-row conversion when the bridge jar is unavailable.
         """
-        import json
-
         if int(batch_size) < 1:
             raise ValueError(f"batch_size must be at least 1, got {batch_size}")
         row_batcher = _bridge_class("RowBatcher")
@@ -595,7 +632,7 @@ def iter_json_batches(
                 java_batch = row_batcher.nextJsonBatch(self._java_result_set, size)
             except JException as exc:
                 raise _read_error(exc) from exc
-            batch = json.loads(str(java_batch))
+            batch = _loads_nonfinite(str(java_batch))
             if len(batch) < size:
                 # nextJsonBatch stops short only when the result set is drained, so a
                 # short batch is the last one: no second call (a JVM crossing, a

src/java/com/arcadedb/python/RowBatcher.java

@@ -32,13 +32,103 @@
 
 import com.arcadedb.query.sql.executor.Result;
 import com.arcadedb.query.sql.executor.ResultSet;
+import com.arcadedb.serializer.json.JSONArray;
 import com.arcadedb.serializer.json.JSONObject;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
 public final class RowBatcher {
 
+  /** The key of a tagged value; Python reads the same character (the tag TypedRows uses). */
+  static final String TAG = "\u0001";
+
   private RowBatcher() {
   }
 
+  /**
+   * JSON has no NaN or infinity, and the engine's writers turn them into null. This returns {@code value} with every
+   * NaN, Infinity and -Infinity (a Double or Float, or an element of a float[], a double[], a Collection or a Map, at any
+   * depth) replaced by the tag {@code {"\u0001": ["f", "NaN"]}} (the text is {@code Double.toString}), which Python turns
+   * back into a float. A value with nothing non-finite is returned as it is, with no copy, so the common row costs one
+   * type test per value.
+   */
+  static Object keepNonFinite(final Object value) {
+    if (value instanceof Double || value instanceof Float) {
+      final double d = ((Number) value).doubleValue();
+      return Double.isFinite(d) ? value : tag(d);
+    }
+    if (value instanceof float[] floats) {
+      for (final float f : floats)
+        if (!Float.isFinite(f)) {
+          final List<Object> out = new ArrayList<>(floats.length);
+          for (final float x : floats)
+            out.add(Float.isFinite(x) ? (Object) x : tag(x));
+          return out;
+        }
+      return value;
+    }
+    if (value instanceof double[] doubles) {
+      for (final double d : doubles)
+        if (!Double.isFinite(d)) {
+          final List<Object> out = new ArrayList<>(doubles.length);
+          for (final double x : doubles)
+            out.add(Double.isFinite(x) ? (Object) x : tag(x));
+          return out;
+        }
+      return value;
+    }
+    if (value instanceof Collection<?> collection) {
+      List<Object> out = null;
+      int i = 0;
+      for (final Object element : collection) {
+        final Object kept = keepNonFinite(element);
+        if (kept != element && out == null) {
+          out = new ArrayList<>(collection.size());
+          int j = 0;
+          for (final Object before : collection) {
+            if (j++ == i)
+              break;
+            out.add(before);
+          }
+        }
+        if (out != null)
+          out.add(kept);
+        i++;
+      }
+      return out != null ? out : value;
+    }
+    if (value instanceof Map<?, ?> map) {
+      Map<Object, Object> out = null;
+      int i = 0;
+      for (final Map.Entry<?, ?> entry : map.entrySet()) {
+        final Object element = entry.getValue();
+        final Object kept = keepNonFinite(element);
+        if (kept != element && out == null) {
+          out = new LinkedHashMap<>(map.size() * 2);
+          int j = 0;
+          for (final Map.Entry<?, ?> before : map.entrySet()) {
+            if (j++ == i)
+              break;
+            out.put(before.getKey(), before.getValue());
+          }
+        }
+        if (out != null)
+          out.put(entry.getKey(), kept);
+        i++;
+      }
+      return out != null ? out : value;
+    }
+    return value;
+  }
+
+  private static JSONObject tag(final double nonFinite) {
+    return new JSONObject().put(TAG, new JSONArray(List.of("f", Double.toString(nonFinite))));
+  }
+
   /**
    * Serialize up to {@code max} rows of the result set into a JSON array
    * string. Returns fewer than {@code max} rows only when the result set is
@@ -86,7 +176,8 @@ private static void appendRow(final StringBuilder sb, final Result row) {
       // java.util.Calendar as its epoch milliseconds like a Date (#9630), and writes a date past the range of a long of
       // milliseconds as its ISO text (#10483), where the encoder this replaced threw and failed the whole batch. The members
       // of a list or a map go the same way
-      final Object value = row.getProperty(property);
+      // The one exception: NaN and the infinities, which the writers would turn into null, travel as a tag
+      final Object value = keepNonFinite(row.getProperty(property));
       obj.put(property, value);
     }
     sb.append(obj);

tests/test_nonfinite_floats.py

@@ -0,0 +1,175 @@
+"""NaN, Infinity and -Infinity survive every reader, at the top level and inside a list or a map.
+
+JSON has no literal for them and the engine's JSON writers turn them into null, so the JVM build
+returned None from ``to_json_list()`` and, for a value inside a list or a map, from ``to_columns()``,
+``to_dataframe()`` and ``to_arrow()`` (the native build kept them). The bridge now writes them as a
+tag that Python turns back into a float, so both builds return a Python float where the object reader
+``to_list()`` always did. ``to_list()`` is the control in every test below.
+
+A top-level number is a float64 column in the columnar readers, which kept NaN already.
+Shared by both builds: ``tests/conftest.py`` runs this file against the native wheel as well.
+"""
+
+import math
+
+import numpy as np
+import pyarrow as pa
+import pytest
+
+import arcadedb
+
+VALUES = [
+    pytest.param(math.nan, id="nan"),
+    pytest.param(math.inf, id="inf"),
+    pytest.param(-math.inf, id="ninf"),
+]
+SHAPES = ["top", "list", "map"]
+
+
+def _wrap(shape, value):
+    if shape == "list":
+        return [value]
+    if shape == "map":
+        return {"x": value}
+    return value
+
+
+def _same(actual, expected):
+    """One float: NaN is NaN, an infinity is that infinity (with its sign), and None fails."""
+    assert actual is not None, "a non-finite float came back as None"
+    actual = float(actual)
+    if math.isnan(expected):
+        assert math.isnan(actual)
+    else:
+        assert math.isinf(actual) and actual == expected
+
+
+def _check(actual, shape, value):
+    if shape == "list":
+        assert len(actual) == 1
+        _same(actual[0], value)
+    elif shape == "map":
+        assert list(actual) == ["x"]
+        _same(actual["x"], value)
+    else:
+        _same(actual, value)
+
+
+@pytest.fixture
+def db(temp_db_path):
+    with arcadedb.create_database(temp_db_path) as database:
+        database.command("sql", "CREATE DOCUMENT TYPE Num")
+        yield database
+
+
+def _insert(db, shape, value, k=0):
+    with db.transaction():
+        db.command(
+            "sql",
+            "INSERT INTO Num SET k = :k, v = :v",
+            {"k": k, "v": _wrap(shape, value)},
+        )
+
+
+QUERY = "SELECT v FROM Num ORDER BY k"
+
+
+@pytest.mark.parametrize("shape", SHAPES)
+@pytest.mark.parametrize("value", VALUES)
+def test_to_list_keeps_it(db, shape, value):
+    """The control: the object reader always kept it."""
+    _insert(db, shape, value)
+    _check(db.query("sql", QUERY).to_list()[0]["v"], shape, value)
+
+
+@pytest.mark.parametrize("batch_size", [10_000, 1])
+@pytest.mark.parametrize("shape", SHAPES)
+@pytest.mark.parametrize("value", VALUES)
+def test_to_json_list_returns_a_float(db, shape, value, batch_size):
+    """to_json_list() gives Python floats, as the native build does (JSON-native types otherwise)."""
+    _insert(db, shape, value)
+    row = db.query("sql", QUERY).to_json_list(batch_size=batch_size)[0]
+    _check(row["v"], shape, value)
+
+
+@pytest.mark.parametrize("shape", SHAPES)
+@pytest.mark.parametrize("value", VALUES)
+def test_iter_json_batches_returns_a_float(db, shape, value):
+    _insert(db, shape, value)
+    (batch,) = list(db.query("sql", QUERY).iter_json_batches(batch_size=5))
+    _check(batch[0]["v"], shape, value)
+
+
+@pytest.mark.parametrize("batch_size", [25_000, 2, 1])
+@pytest.mark.parametrize("shape", SHAPES)
+@pytest.mark.parametrize("value", VALUES)
+def test_to_columns_keeps_it(db, shape, value, batch_size):
+    # a list of floats is a numpy matrix on the native build and a list of lists on the JVM build:
+    # tolist() reads both
+    _insert(db, shape, value)
+    column = db.query("sql", QUERY).to_columns(batch_size=batch_size)["v"]
+    cell = column[0]
+    _check(cell.tolist() if isinstance(cell, np.ndarray) else cell, shape, value)
+
+
+@pytest.mark.parametrize("shape", SHAPES)
+@pytest.mark.parametrize("value", VALUES)
+def test_to_dataframe_keeps_it(db, shape, value):
+    _insert(db, shape, value)
+    cell = db.query("sql", QUERY).to_dataframe()["v"].iloc[0]
+    _check(cell.tolist() if isinstance(cell, np.ndarray) else cell, shape, value)
+
+
+@pytest.mark.parametrize("batch_size", [25_000, 2, 1])
+@pytest.mark.parametrize("shape", SHAPES)
+@pytest.mark.parametrize("value", VALUES)
+def test_to_arrow_keeps_it(db, shape, value, batch_size):
+    _insert(db, shape, value)
+    table = db.query("sql", QUERY).to_arrow(batch_size=batch_size)
+    assert isinstance(table, pa.Table)
+    _check(table.column("v").to_pylist()[0], shape, value)
+
+
+@pytest.mark.parametrize("shape", ["list", "map"])
+def test_rows_around_the_non_finite_one_are_unchanged(db, shape):
+    """A row with no NaN, a row with one, a row whose text holds the tag character, in batches of one and two."""
+    plain = [1.5, 2.5] if shape == "list" else {"x": 1.5, "y": "a\x01b"}
+    holder = (
+        [math.nan, math.inf, -math.inf]
+        if shape == "list"
+        else {"x": math.nan, "y": "a\x01b"}
+    )
+    with db.transaction():
+        db.command("sql", "INSERT INTO Num SET k = 0, v = :v", {"v": plain})
+        db.command("sql", "INSERT INTO Num SET k = 1, v = :v", {"v": holder})
+        db.command("sql", "INSERT INTO Num SET k = 2, v = :v", {"v": plain})
+    for batch_size in (2, 1):
+        rows = db.query("sql", QUERY).to_json_list(batch_size=batch_size)
+        assert rows[0]["v"] == plain and rows[2]["v"] == plain
+        if shape == "list":
+            assert math.isnan(rows[1]["v"][0])
+            assert rows[1]["v"][1:] == [math.inf, -math.inf]
+        else:
+            assert math.isnan(rows[1]["v"]["x"]) and rows[1]["v"]["y"] == "a\x01b"
+    for batch_size in (2, 1):
+        arrow = (
+            db.query("sql", QUERY)
+            .to_arrow(batch_size=batch_size)
+            .column("v")
+            .to_pylist()
+        )
+        assert arrow[0] == plain
+        if shape == "list":
+            assert math.isnan(arrow[1][0]) and arrow[1][1:] == [math.inf, -math.inf]
+        else:
+            assert math.isnan(arrow[1]["x"])
+
+
+def test_a_string_holding_the_tag_character_is_not_a_float(db):
+    with db.transaction():
+        db.command("sql", "INSERT INTO Num SET k = 0, v = :v", {"v": "\x01"})
+        db.command(
+            "sql", "INSERT INTO Num SET k = 1, v = :v", {"v": ["\x01", "f", "NaN"]}
+        )
+    rows = db.query("sql", QUERY).to_json_list()
+    assert [r["v"] for r in rows] == ["\x01", ["\x01", "f", "NaN"]]