Parquet formatter coercion issue (#494)

* Normalize column types

* Update unit test

* Update changelog

* Change return type for _normalize_column_types()

Author: NathalieCharbel

CHANGELOG.md

@@ -14,6 +14,10 @@
 
 - `SchemaExtractionTemplate` prompt updated to explicitly instruct the LLM not to use `__` as a prefix or suffix in node labels or relationship types.
 
+### Fixed
+
+- Fixed `ValueError` in `Neo4jGraphParquetFormatter` when nodes of the same label have mixed property types (e.g. `str` and `int` for the same property), which caused `pa.Table.from_pylist()` to fail. Mixed-type columns are now coerced to a consistent type before Parquet table creation.
+
 ## 1.14.0
 
 ### Added

src/neo4j_graphrag/experimental/components/parquet_formatter.py

@@ -419,6 +419,46 @@ def _relationships_to_rows(
 
         return type_to_rows
 
+    @staticmethod
+    def _normalize_column_types(rows: list[dict[str, Any]]) -> None:
+        """Coerce mixed-type columns in *rows* in-place so PyArrow can build the table.
+
+        PyArrow infers the column type from the first row; if subsequent rows have a
+        different Python type for the same column the table creation fails.  This method
+        detects those mismatches and coerces:
+        - {int, float} -> float  (lossless numeric promotion)
+        - anything else mixed -> str  (universal safe fallback)
+        """
+        if len(rows) <= 1:
+            return
+
+        col_types: dict[str, set[type]] = defaultdict(set)
+        for row in rows:
+            for key, value in row.items():
+                if value is not None:
+                    col_types[key].add(type(value))
+
+        cols_to_coerce: dict[str, type] = {}
+        for col, types in col_types.items():
+            if len(types) <= 1:
+                continue
+            target: type = float if types <= {int, float} else str
+            cols_to_coerce[col] = target
+            logger.warning(
+                "Mixed types for property '%s': %s — coercing to %s",
+                col,
+                {t.__name__ for t in types},
+                target.__name__,
+            )
+
+        for row in rows:
+            for col, target_type in cols_to_coerce.items():
+                if col in row and row[col] is not None:
+                    try:
+                        row[col] = target_type(row[col])
+                    except (ValueError, TypeError):
+                        row[col] = str(row[col])
+
     def format_parquet(
         self,
         rows: list[dict[str, Any]],
@@ -440,6 +480,7 @@ def format_parquet(
             import pyarrow as pa
             import pyarrow.parquet as pq
 
+            self._normalize_column_types(rows)
             table = pa.Table.from_pylist(rows)
             # Write to BytesIO buffer
             buffer = BytesIO()

tests/unit/experimental/components/test_kg_writer.py

@@ -25,6 +25,7 @@
     FilenameCollisionHandler,
 )
 from neo4j_graphrag.experimental.components.parquet_formatter import (
+    Neo4jGraphParquetFormatter,
     sanitize_parquet_filestem,
 )
 from neo4j_graphrag.experimental.components.kg_writer import (
@@ -697,3 +698,72 @@ async def test_parquet_writer_run_empty_graph() -> None:
     assert stats["nodes_per_label"] == {}
     assert stats["rel_per_type"] == {}
     assert result.metadata["files"] == []
+
+
+# ---------------------------------------------------------------------------
+# Neo4jGraphParquetFormatter._normalize_column_types
+# ---------------------------------------------------------------------------
+
+
+def test_normalize_column_types_single_row() -> None:
+    rows = [{"age": 30, "name": "Alice"}]
+    Neo4jGraphParquetFormatter._normalize_column_types(rows)
+    assert rows == [{"age": 30, "name": "Alice"}]
+
+
+def test_normalize_column_types_homogeneous() -> None:
+    rows = [{"age": 30}, {"age": 25}]
+    Neo4jGraphParquetFormatter._normalize_column_types(rows)
+    assert rows == [{"age": 30}, {"age": 25}]
+
+
+def test_normalize_column_types_mixed_str_int() -> None:
+    rows: list[dict[str, Any]] = [{"age": "45"}, {"age": 30}]
+    Neo4jGraphParquetFormatter._normalize_column_types(rows)
+    assert rows == [{"age": "45"}, {"age": "30"}]
+
+
+def test_normalize_column_types_mixed_int_float() -> None:
+    rows: list[dict[str, Any]] = [{"score": 3}, {"score": 3.5}]
+    Neo4jGraphParquetFormatter._normalize_column_types(rows)
+    assert rows == [{"score": 3.0}, {"score": 3.5}]
+
+
+def test_normalize_column_types_none_ignored() -> None:
+    """None values should not influence type detection."""
+    rows: list[dict[str, Any]] = [{"age": None}, {"age": 30}]
+    Neo4jGraphParquetFormatter._normalize_column_types(rows)
+    assert rows == [{"age": None}, {"age": 30}]
+
+
+@pytest.mark.asyncio
+async def test_parquet_writer_mixed_property_types() -> None:
+    """ParquetWriter succeeds when nodes of the same label have mixed property types."""
+    pytest.importorskip("pyarrow")
+    import pyarrow.parquet as pq
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        out = Path(tmpdir)
+        dest = _LocalParquetDestination(out)
+        writer = ParquetWriter(
+            nodes_dest=dest,
+            relationships_dest=dest,
+            collision_handler=FilenameCollisionHandler(),
+        )
+
+        node1 = Neo4jNode(
+            id="p1", label="Patient", properties={"name": "John", "age": "45"}
+        )
+        node2 = Neo4jNode(
+            id="p2", label="Patient", properties={"name": "Jane", "age": 30}
+        )
+        graph = Neo4jGraph(nodes=[node1, node2], relationships=[])
+
+        result = await writer.run(graph=graph)
+
+        assert result.status == "SUCCESS"
+        table = pq.read_table(out / "Patient.parquet")
+        assert table.num_rows == 2
+        # Both ages should have been coerced to str
+        ages = {v.as_py() for v in table.column("age")}
+        assert ages == {"45", "30"}