lasta/AI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
9e88f4c7ad724b42ea25482ed74448fe58726314
AI/참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_kg_writer.py

1498 lines
52 KiB
Python
Raw Normal View History

참고소스 수정본
2026-05-12 19:40:31 +09:00
# Copyright (c) "Neo4j"
# Neo4j Sweden AB [https://neo4j.com]
# #
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# #
# https://www.apache.org/licenses/LICENSE-2.0
# #
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import tempfile
from io import BytesIO
from pathlib import Path
from typing import Any
from unittest import mock
from unittest.mock import MagicMock, Mock
import pytest
from neo4j_graphrag.experimental.components.filename_collision_handler import (
FilenameCollisionHandler,
)
from neo4j_graphrag.experimental.components.parquet_formatter import (
INTERNAL_ID_PROPERTY,
Neo4jGraphParquetFormatter,
get_unique_properties_for_node_type,
sanitize_parquet_filestem,
)
from neo4j_graphrag.experimental.components.kg_writer import (
Neo4jWriter,
ParquetWriter,
batched,
)
from neo4j_graphrag.experimental.components.schema import GraphSchema
from neo4j_graphrag.experimental.components.types import (
LexicalGraphConfig,
Neo4jGraph,
Neo4jNode,
Neo4jRelationship,
)
from neo4j_graphrag.neo4j_queries import (
upsert_node_query,
upsert_relationship_query,
)
def test_batched() -> None:
assert list(batched([1, 2, 3, 4], batch_size=2)) == [
[1, 2],
[3, 4],
]
assert list(batched([1, 2, 3], batch_size=2)) == [
[1, 2],
[3],
]
assert list(batched([1, 2, 3], batch_size=4)) == [
[1, 2, 3],
]
# --- sanitize_parquet_filestem tests ---
def test_sanitize_parquet_filestem_empty_returns_fallback() -> None:
assert sanitize_parquet_filestem("") == "unnamed"
def test_sanitize_parquet_filestem_safe_chars_unchanged() -> None:
assert sanitize_parquet_filestem("Person") == "Person"
assert sanitize_parquet_filestem("Person_KNOWS_Person") == "Person_KNOWS_Person"
assert sanitize_parquet_filestem("Label123") == "Label123"
assert sanitize_parquet_filestem("a_z_9") == "a_z_9"
def test_sanitize_parquet_filestem_unicode_transliterated() -> None:
assert sanitize_parquet_filestem("Zürich") == "Zurich"
assert sanitize_parquet_filestem("café") == "cafe"
assert sanitize_parquet_filestem("naïve") == "naive"
def test_sanitize_parquet_filestem_disallowed_replaced_with_underscore() -> None:
assert sanitize_parquet_filestem("a b") == "a_b"
assert sanitize_parquet_filestem("a-b") == "a_b"
assert sanitize_parquet_filestem("a.b") == "a_b"
def test_sanitize_parquet_filestem_all_disallowed_replaced() -> None:
# All disallowed chars become underscores (result non-empty, so no fallback)
assert sanitize_parquet_filestem("...") == "___"
assert sanitize_parquet_filestem(" ") == "__"
def test_get_unique_properties_for_node_type_deprecation_warning() -> None:
with pytest.warns(DeprecationWarning, match="get_unique_properties_for_node_type"):
assert get_unique_properties_for_node_type(None, "Person") == [
INTERNAL_ID_PROPERTY
]
# --- FilenameCollisionHandler tests ---
def test_filename_collision_handler_first_call_returns_unchanged() -> None:
FilenameCollisionHandler.reset()
handler = FilenameCollisionHandler()
out = Path("/some/output")
assert handler.get_unique_filename("Person.parquet", out) == "Person.parquet"
def test_filename_collision_handler_collisions_get_suffix() -> None:
FilenameCollisionHandler.reset()
handler = FilenameCollisionHandler()
out = Path("/some/output")
assert handler.get_unique_filename("Person.parquet", out) == "Person.parquet"
assert handler.get_unique_filename("Person.parquet", out) == "Person_1.parquet"
assert handler.get_unique_filename("Person.parquet", out) == "Person_2.parquet"
def test_filename_collision_handler_different_paths_no_collision() -> None:
FilenameCollisionHandler.reset()
handler = FilenameCollisionHandler()
out1 = Path("/out/a")
out2 = Path("/out/b")
assert handler.get_unique_filename("Person.parquet", out1) == "Person.parquet"
assert handler.get_unique_filename("Person.parquet", out2) == "Person.parquet"
def test_filename_collision_handler_reset_clears_state() -> None:
FilenameCollisionHandler.reset()
handler = FilenameCollisionHandler()
out = Path("/out")
handler.get_unique_filename("Person.parquet", out)
handler.get_unique_filename("Person.parquet", out)
FilenameCollisionHandler.reset()
assert handler.get_unique_filename("Person.parquet", out) == "Person.parquet"
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.get_version",
return_value=((5, 22, 0), False, False),
)
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.Neo4jWriter._db_setup",
return_value=None,
)
def test_upsert_nodes(_: Mock, driver: MagicMock) -> None:
driver.execute_query.return_value = (
[{"element_id": "#1"}],
None,
None,
)
neo4j_writer = Neo4jWriter(driver=driver)
node = Neo4jNode(id="1", label="Label", properties={"key": "value"})
neo4j_writer._upsert_nodes(nodes=[node], lexical_graph_config=LexicalGraphConfig())
driver.execute_query.assert_called_once_with(
upsert_node_query(False),
parameters_={
"rows": [
{
"label": "Label",
"labels": ["Label", "__Entity__"],
"id": "1",
"properties": {"key": "value"},
"embedding_properties": {},
}
]
},
database_=None,
)
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.get_version",
return_value=((5, 22, 0), False, False),
)
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.Neo4jWriter._db_setup",
return_value=None,
)
def test_upsert_nodes_with_embedding(
_: Mock,
driver: MagicMock,
) -> None:
driver.execute_query.return_value = (
[{"element_id": "#1"}],
None,
None,
)
neo4j_writer = Neo4jWriter(driver=driver)
node = Neo4jNode(
id="1",
label="Label",
properties={"key": "value"},
embedding_properties={"embeddingProp": [1.0, 2.0, 3.0]},
)
neo4j_writer._upsert_nodes(nodes=[node], lexical_graph_config=LexicalGraphConfig())
driver.execute_query.assert_any_call(
upsert_node_query(False),
parameters_={
"rows": [
{
"label": "Label",
"labels": ["Label", "__Entity__"],
"id": "1",
"properties": {"key": "value"},
"embedding_properties": {"embeddingProp": [1.0, 2.0, 3.0]},
}
]
},
database_=None,
)
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.get_version",
return_value=((5, 22, 0), False, False),
)
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.Neo4jWriter._db_setup",
return_value=None,
)
def test_upsert_relationship(_: Mock, driver: MagicMock) -> None:
neo4j_writer = Neo4jWriter(driver=driver)
rel = Neo4jRelationship(
start_node_id="1",
end_node_id="2",
type="RELATIONSHIP",
properties={"key": "value"},
)
neo4j_writer._upsert_relationships(
rels=[rel],
)
parameters = {
"rows": [
{
"type": "RELATIONSHIP",
"start_node_id": "1",
"end_node_id": "2",
"properties": {"key": "value"},
"embedding_properties": {},
}
]
}
driver.execute_query.assert_called_once_with(
upsert_relationship_query(False),
parameters_=parameters,
database_=None,
)
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.get_version",
return_value=((5, 22, 0), False, False),
)
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.Neo4jWriter._db_setup",
return_value=None,
)
def test_upsert_relationship_with_embedding(_: Mock, driver: MagicMock) -> None:
neo4j_writer = Neo4jWriter(driver=driver)
rel = Neo4jRelationship(
start_node_id="1",
end_node_id="2",
type="RELATIONSHIP",
properties={"key": "value"},
embedding_properties={"embeddingProp": [1.0, 2.0, 3.0]},
)
driver.execute_query.return_value.records = [{"elementId(r)": "rel_elem_id"}]
neo4j_writer._upsert_relationships(
rels=[rel],
)
parameters = {
"rows": [
{
"type": "RELATIONSHIP",
"start_node_id": "1",
"end_node_id": "2",
"properties": {"key": "value"},
"embedding_properties": {"embeddingProp": [1.0, 2.0, 3.0]},
}
]
}
driver.execute_query.assert_any_call(
upsert_relationship_query(False),
parameters_=parameters,
database_=None,
)
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.get_version",
return_value=((5, 22, 0), False, False),
)
@pytest.mark.asyncio
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.Neo4jWriter._db_setup",
return_value=None,
)
async def test_run(_: Mock, driver: MagicMock) -> None:
driver.execute_query.return_value = (
[
{"element_id": "#1"},
{"element_id": "#2"},
],
None,
None,
)
neo4j_writer = Neo4jWriter(driver=driver)
node = Neo4jNode(id="1", label="Label")
rel = Neo4jRelationship(start_node_id="1", end_node_id="2", type="RELATIONSHIP")
graph = Neo4jGraph(nodes=[node], relationships=[rel])
await neo4j_writer.run(graph=graph)
driver.execute_query.assert_any_call(
upsert_node_query(False),
parameters_={
"rows": [
{
"label": "Label",
"labels": ["Label", "__Entity__"],
"id": "1",
"properties": {},
"embedding_properties": {},
}
]
},
database_=None,
)
parameters_: dict[str, Any] = {
"rows": [
{
"type": "RELATIONSHIP",
"start_node_id": "1",
"end_node_id": "2",
"properties": {},
"embedding_properties": {},
}
]
}
driver.execute_query.assert_any_call(
upsert_relationship_query(False),
parameters_=parameters_,
database_=None,
)
@pytest.mark.asyncio
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.Neo4jWriter._db_setup",
return_value=None,
)
async def test_run_is_version_below_5_23(_: Mock) -> None:
driver = MagicMock()
driver.execute_query = Mock(
side_effect=(
# get_version
([{"versions": ["5.22.0"], "edition": "enterpise"}], None, None),
# upsert nodes
([{"_internal_id": "1", "element_id": "#1"}], None, None),
# upsert relationships
(None, None, None),
)
)
neo4j_writer = Neo4jWriter(driver=driver)
node = Neo4jNode(id="1", label="Label")
rel = Neo4jRelationship(start_node_id="1", end_node_id="2", type="RELATIONSHIP")
graph = Neo4jGraph(nodes=[node], relationships=[rel])
await neo4j_writer.run(graph=graph)
driver.execute_query.assert_any_call(
upsert_node_query(False),
parameters_={
"rows": [
{
"label": "Label",
"labels": ["Label", "__Entity__"],
"id": "1",
"properties": {},
"embedding_properties": {},
}
]
},
database_=None,
)
parameters_: dict[str, Any] = {
"rows": [
{
"type": "RELATIONSHIP",
"start_node_id": "1",
"end_node_id": "2",
"properties": {},
"embedding_properties": {},
}
]
}
driver.execute_query.assert_any_call(
upsert_relationship_query(False),
parameters_=parameters_,
database_=None,
)
@pytest.mark.asyncio
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.Neo4jWriter._db_setup",
return_value=None,
)
async def test_run_is_version_5_23_or_above(_: Mock) -> None:
driver = MagicMock()
driver.execute_query = Mock(
side_effect=(
# get_version
([{"versions": ["5.23.0"], "edition": "enterpise"}], None, None),
# upsert nodes
([{"element_id": "#1"}], None, None),
# upsert relationships
(None, None, None),
)
)
neo4j_writer = Neo4jWriter(driver=driver)
neo4j_writer.is_version_5_23_or_above = True
node = Neo4jNode(id="1", label="Label")
rel = Neo4jRelationship(start_node_id="1", end_node_id="2", type="RELATIONSHIP")
graph = Neo4jGraph(nodes=[node], relationships=[rel])
await neo4j_writer.run(graph=graph)
driver.execute_query.assert_any_call(
upsert_node_query(True, False),
parameters_={
"rows": [
{
"label": "Label",
"labels": ["Label", "__Entity__"],
"id": "1",
"properties": {},
"embedding_properties": {},
}
]
},
database_=None,
)
parameters_: dict[str, Any] = {
"rows": [
{
"type": "RELATIONSHIP",
"start_node_id": "1",
"end_node_id": "2",
"properties": {},
"embedding_properties": {},
}
]
}
driver.execute_query.assert_any_call(
upsert_relationship_query(True),
parameters_=parameters_,
database_=None,
)
@pytest.mark.asyncio
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.Neo4jWriter._db_setup",
return_value=None,
)
async def test_run_is_version_5_24_or_above(_: Mock) -> None:
driver = MagicMock()
driver.execute_query = Mock(
side_effect=(
# get_version
([{"versions": ["5.24.0"], "edition": "enterprise"}], None, None),
# upsert nodes
([{"element_id": "#1"}], None, None),
# upsert relationships
(None, None, None),
)
)
neo4j_writer = Neo4jWriter(driver=driver)
node = Neo4jNode(id="1", label="Label")
rel = Neo4jRelationship(start_node_id="1", end_node_id="2", type="RELATIONSHIP")
graph = Neo4jGraph(nodes=[node], relationships=[rel])
await neo4j_writer.run(graph=graph)
driver.execute_query.assert_any_call(
upsert_node_query(True, True),
parameters_={
"rows": [
{
"label": "Label",
"labels": ["Label", "__Entity__"],
"id": "1",
"properties": {},
"embedding_properties": {},
}
]
},
database_=None,
)
parameters_: dict[str, Any] = {
"rows": [
{
"type": "RELATIONSHIP",
"start_node_id": "1",
"end_node_id": "2",
"properties": {},
"embedding_properties": {},
}
]
}
driver.execute_query.assert_any_call(
upsert_relationship_query(True),
parameters_=parameters_,
database_=None,
)
@pytest.mark.parametrize(
"description, version, is_5_23_or_above, is_5_24_or_above",
[
("SemVer, < 5.23", "5.22.0", False, False),
("SemVer, == 5.23", "5.23.0", True, False),
("SemVer, > 5.23", "5.24.0", True, True),
("SemVer, < 5.23, Aura", "5.22-aura", False, False),
("SemVer, > 5.23, Aura", "5.24-aura", True, True),
("CalVer", "2025.01.0", True, True),
("CalVer, Aura", "2025.01-aura", True, True),
],
)
@mock.patch(
"neo4j_graphrag.experimental.components.kg_writer.Neo4jWriter._db_setup",
return_value=None,
)
def test_get_version(
_: Mock,
driver: MagicMock,
description: str,
version: str,
is_5_23_or_above: bool,
is_5_24_or_above: bool,
) -> None:
execute_query_mock = MagicMock(
return_value=(
[
{"versions": [version], "edition": "enterprise"},
],
None,
None,
)
)
driver.execute_query = execute_query_mock
neo4j_writer = Neo4jWriter(driver=driver)
assert (
neo4j_writer.is_version_5_23_or_above is is_5_23_or_above
), f"Failed is_version_5_23_or_above test case: {description}"
assert (
neo4j_writer.is_version_5_24_or_above is is_5_24_or_above
), f"Failed is_version_5_24_or_above test case: {description}"
# --- ParquetWriter tests ---
class _LocalParquetDestination:
"""Test-only implementation of ParquetOutputDestination for a local directory."""
def __init__(self, path: Path) -> None:
self._path = Path(path)
self._path.mkdir(parents=True, exist_ok=True)
@property
def output_path(self) -> str:
return str(self._path.resolve())
async def write(self, data: bytes, filename: str) -> None:
(self._path / filename).write_bytes(data)
@pytest.mark.asyncio
async def test_parquet_writer_missing_pyarrow_raises() -> None:
"""When pyarrow is not installed, run() returns FAILURE with error mentioning pyarrow."""
import builtins
real_import = builtins.__import__
def fake_import(name: str, *args: Any, **kwargs: Any) -> Any:
if name == "pyarrow":
raise ImportError("No module named 'pyarrow'")
return real_import(name, *args, **kwargs)
with tempfile.TemporaryDirectory() as tmpdir:
dest = _LocalParquetDestination(Path(tmpdir))
writer = ParquetWriter(
nodes_dest=dest,
relationships_dest=dest,
collision_handler=FilenameCollisionHandler(),
)
# Use non-empty graph so formatter calls format_parquet and triggers pyarrow import
node = Neo4jNode(id="n1", label="Person", properties={})
graph = Neo4jGraph(nodes=[node], relationships=[])
with mock.patch("builtins.__import__", side_effect=fake_import):
result = await writer.run(graph=graph)
assert result.status == "FAILURE"
assert result.metadata is not None and "error" in result.metadata
assert "pyarrow" in result.metadata["error"].lower()
@pytest.mark.asyncio
async def test_parquet_writer_run_success() -> None:
"""ParquetWriter uses formatter and writes one file per node label and per (head, type, tail)."""
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="n1", label="Person", properties={"name": "Alice"})
node2 = Neo4jNode(id="n2", label="Person", properties={"name": "Bob"})
rel = Neo4jRelationship(
start_node_id="n1", end_node_id="n2", type="KNOWS", properties={}
)
graph = Neo4jGraph(nodes=[node1, node2], relationships=[rel])
result = await writer.run(graph=graph)
assert result.status == "SUCCESS"
assert result.metadata is not None
stats = result.metadata.get("statistics") or {}
assert stats["node_count"] == 2
assert stats["relationship_count"] == 1
assert stats["nodes_per_label"] == {"Person": 2}
assert "KNOWS" in stats["rel_per_type"]
assert "input_files_count" in stats
assert "input_files_total_size_bytes" in stats
assert (out / "Person.parquet").exists()
assert (out / "Person_KNOWS_Person.parquet").exists()
# Check "files" metadata (file_path, columns, source mapping for rels)
assert "files" in result.metadata
assert len(result.metadata["files"]) == 2
node_file_info = next(f for f in result.metadata["files"] if f["is_node"])
assert node_file_info["name"] == "Person"
assert (
"file_path" in node_file_info
and "Person.parquet" in node_file_info["file_path"]
)
assert "columns" in node_file_info
assert any(
c["name"] == INTERNAL_ID_PROPERTY
and c["is_primary_key"]
and c["is_unique"] is False
for c in node_file_info["columns"]
)
assert {"type": "KEY", "properties": [INTERNAL_ID_PROPERTY]} in (
node_file_info.get("constraints") or []
)
rel_file_info = next(f for f in result.metadata["files"] if not f["is_node"])
assert rel_file_info["relationship_type"] == "KNOWS"
assert rel_file_info["start_node_source"] == "Person"
assert rel_file_info["end_node_source"] == "Person"
assert rel_file_info["start_node_primary_keys"] == [INTERNAL_ID_PROPERTY]
assert rel_file_info["end_node_primary_keys"] == [INTERNAL_ID_PROPERTY]
# Read back and sanity-check (formatter uses internal id, labels, and flat properties)
nodes_table = pq.read_table(out / "Person.parquet")
assert nodes_table.num_rows == 2
assert INTERNAL_ID_PROPERTY in nodes_table.column_names
assert "labels" in nodes_table.column_names
assert "name" in nodes_table.column_names
rels_table = pq.read_table(out / "Person_KNOWS_Person.parquet")
assert rels_table.num_rows == 1
assert "from" in rels_table.column_names
assert "to" in rels_table.column_names
assert rels_table.column("type")[0].as_py() == "KNOWS"
rel_cols = {c["name"]: c for c in rel_file_info["columns"]}
assert rel_cols["from"]["is_primary_key"] is True
assert rel_cols["from"]["is_unique"] is False
assert rel_cols["to"]["is_primary_key"] is True
assert rel_cols["to"]["is_unique"] is False
@pytest.mark.asyncio
async def test_parquet_writer_relationship_joins_on_single_property_key() -> None:
"""With a single-property KEY in schema, rel from/to use that property and metadata matches."""
pytest.importorskip("pyarrow")
import pyarrow.parquet as pq
schema_dict: dict[str, Any] = {
"node_types": [
{
"label": "Person",
"properties": [
{"name": "email", "type": "STRING"},
{"name": "name", "type": "STRING"},
],
}
],
"constraints": [
{
"type": "KEY",
"node_type": "Person",
"property_names": ["email"],
"relationship_type": None,
}
],
}
with tempfile.TemporaryDirectory() as tmpdir:
out = Path(tmpdir)
dest = _LocalParquetDestination(out)
writer = ParquetWriter(
nodes_dest=dest,
relationships_dest=dest,
collision_handler=FilenameCollisionHandler(),
)
n1 = Neo4jNode(
id="n1",
label="Person",
properties={"email": "a@b.c", "name": "Alice"},
)
n2 = Neo4jNode(
id="n2",
label="Person",
properties={"email": "b@b.c", "name": "Bob"},
)
rel = Neo4jRelationship(
start_node_id="n1", end_node_id="n2", type="KNOWS", properties={}
)
graph = Neo4jGraph(nodes=[n1, n2], relationships=[rel])
result = await writer.run(
graph=graph, schema=GraphSchema.model_validate(schema_dict)
)
assert result.status == "SUCCESS"
assert result.metadata is not None
rel_file = next(f for f in result.metadata["files"] if not f["is_node"])
assert rel_file["start_node_primary_keys"] == ["email"]
assert rel_file["end_node_primary_keys"] == ["email"]
rels_table = pq.read_table(Path(rel_file["file_path"]))
assert rels_table.column("from")[0].as_py() == "a@b.c"
assert rels_table.column("to")[0].as_py() == "b@b.c"
@pytest.mark.asyncio
async def test_parquet_writer_columns_uniqueness_sets_is_unique() -> None:
"""UNIQUENESS maps to is_unique; synthetic single-property KEY on internal id when no schema KEY."""
pytest.importorskip("pyarrow")
with tempfile.TemporaryDirectory() as tmpdir:
out = Path(tmpdir)
dest = _LocalParquetDestination(out)
writer = ParquetWriter(
nodes_dest=dest,
relationships_dest=dest,
collision_handler=FilenameCollisionHandler(),
)
schema_dict: dict[str, Any] = {
"node_types": [
{
"label": "Person",
"properties": [
{"name": "email", "type": "STRING"},
{"name": "name", "type": "STRING"},
],
}
],
"constraints": [
{
"type": "UNIQUENESS",
"node_type": "Person",
"property_names": ["email"],
"relationship_type": None,
}
],
}
node = Neo4jNode(
id="n1",
label="Person",
properties={"email": "a@b.c", "name": "Alice"},
)
graph = Neo4jGraph(nodes=[node], relationships=[])
result = await writer.run(
graph=graph, schema=GraphSchema.model_validate(schema_dict)
)
assert result.status == "SUCCESS"
assert result.metadata is not None
node_file = next(f for f in result.metadata["files"] if f["is_node"])
cols = {c["name"]: c for c in node_file["columns"]}
assert cols["email"]["is_unique"] is True
assert cols["email"]["is_primary_key"] is False
assert cols[INTERNAL_ID_PROPERTY]["is_primary_key"] is True
assert cols[INTERNAL_ID_PROPERTY]["is_unique"] is False
key_cs = [c for c in node_file["constraints"] if c["type"] == "KEY"]
assert key_cs == [{"type": "KEY", "properties": [INTERNAL_ID_PROPERTY]}]
@pytest.mark.asyncio
async def test_parquet_writer_columns_key_sets_is_primary_key() -> None:
"""KEY maps to is_primary_key on that property; is_unique stays false."""
pytest.importorskip("pyarrow")
with tempfile.TemporaryDirectory() as tmpdir:
out = Path(tmpdir)
dest = _LocalParquetDestination(out)
writer = ParquetWriter(
nodes_dest=dest,
relationships_dest=dest,
collision_handler=FilenameCollisionHandler(),
)
schema_dict: dict[str, Any] = {
"node_types": [
{
"label": "Person",
"properties": [
{"name": "email", "type": "STRING"},
{"name": "name", "type": "STRING"},
],
}
],
"constraints": [
{
"type": "KEY",
"node_type": "Person",
"property_names": ["email"],
"relationship_type": None,
}
],
}
node = Neo4jNode(
id="n1",
label="Person",
properties={"email": "a@b.c", "name": "Alice"},
)
graph = Neo4jGraph(nodes=[node], relationships=[])
result = await writer.run(
graph=graph, schema=GraphSchema.model_validate(schema_dict)
)
assert result.status == "SUCCESS"
assert result.metadata is not None
node_file = next(f for f in result.metadata["files"] if f["is_node"])
cols = {c["name"]: c for c in node_file["columns"]}
assert cols["email"]["is_primary_key"] is True
assert cols["email"]["is_unique"] is False
@pytest.mark.asyncio
async def test_parquet_writer_composite_key_constraint() -> None:
"""Composite KEY marks all constituent properties as is_primary_key and emits structured constraints metadata."""
pytest.importorskip("pyarrow")
with tempfile.TemporaryDirectory() as tmpdir:
out = Path(tmpdir)
dest = _LocalParquetDestination(out)
writer = ParquetWriter(
nodes_dest=dest,
relationships_dest=dest,
collision_handler=FilenameCollisionHandler(),
)
schema_dict: dict[str, Any] = {
"node_types": [
{
"label": "Actor",
"properties": [
{"name": "firstname", "type": "STRING"},
{"name": "surname", "type": "STRING"},
{"name": "age", "type": "INTEGER"},
],
}
],
"constraints": [
{
"type": "KEY",
"node_type": "Actor",
"property_names": ["firstname", "surname"],
"relationship_type": None,
}
],
}
node = Neo4jNode(
id="a1",
label="Actor",
properties={"firstname": "John", "surname": "Smith", "age": 42},
)
graph = Neo4jGraph(nodes=[node], relationships=[])
result = await writer.run(
graph=graph, schema=GraphSchema.model_validate(schema_dict)
)
assert result.status == "SUCCESS"
assert result.metadata is not None
node_file = next(f for f in result.metadata["files"] if f["is_node"])
cols = {c["name"]: c for c in node_file["columns"]}
# Both properties in the composite KEY should be marked as primary key
assert cols["firstname"]["is_primary_key"] is True
assert cols["surname"]["is_primary_key"] is True
assert cols["age"]["is_primary_key"] is False
assert cols[INTERNAL_ID_PROPERTY]["is_primary_key"] is True
# Structured constraints metadata should preserve composite grouping
assert "constraints" in node_file
key_constraints = [c for c in node_file["constraints"] if c["type"] == "KEY"]
assert len(key_constraints) == 2
assert key_constraints[0]["properties"] == ["firstname", "surname"]
assert key_constraints[1]["properties"] == [INTERNAL_ID_PROPERTY]
@pytest.mark.asyncio
async def test_parquet_writer_composite_uniqueness_constraint() -> None:
"""Composite UNIQUENESS marks all constituent properties as is_unique and emits structured constraints."""
pytest.importorskip("pyarrow")
with tempfile.TemporaryDirectory() as tmpdir:
out = Path(tmpdir)
dest = _LocalParquetDestination(out)
writer = ParquetWriter(
nodes_dest=dest,
relationships_dest=dest,
collision_handler=FilenameCollisionHandler(),
)
schema_dict: dict[str, Any] = {
"node_types": [
{
"label": "Book",
"properties": [
{"name": "title", "type": "STRING"},
{"name": "year", "type": "INTEGER"},
{"name": "isbn", "type": "STRING"},
],
}
],
"constraints": [
{
"type": "UNIQUENESS",
"node_type": "Book",
"property_names": ["title", "year"],
"relationship_type": None,
}
],
}
node = Neo4jNode(
id="b1",
label="Book",
properties={"title": "Neo4j in Action", "year": 2024, "isbn": "123"},
)
graph = Neo4jGraph(nodes=[node], relationships=[])
result = await writer.run(
graph=graph, schema=GraphSchema.model_validate(schema_dict)
)
assert result.status == "SUCCESS"
assert result.metadata is not None
node_file = next(f for f in result.metadata["files"] if f["is_node"])
cols = {c["name"]: c for c in node_file["columns"]}
assert cols["title"]["is_unique"] is True
assert cols["year"]["is_unique"] is True
assert cols["isbn"]["is_unique"] is False
# Structured constraints metadata
assert "constraints" in node_file
unique_constraints = [
c for c in node_file["constraints"] if c["type"] == "UNIQUENESS"
]
assert len(unique_constraints) == 1
assert unique_constraints[0]["properties"] == ["title", "year"]
key_cs = [c for c in node_file["constraints"] if c["type"] == "KEY"]
assert key_cs == [{"type": "KEY", "properties": [INTERNAL_ID_PROPERTY]}]
@pytest.mark.asyncio
async def test_parquet_writer_existence_constraint_in_metadata() -> None:
"""Node-scoped EXISTENCE appears in structured constraints; synthetic KEY on __id__ remains."""
pytest.importorskip("pyarrow")
with tempfile.TemporaryDirectory() as tmpdir:
out = Path(tmpdir)
dest = _LocalParquetDestination(out)
writer = ParquetWriter(
nodes_dest=dest,
relationships_dest=dest,
collision_handler=FilenameCollisionHandler(),
)
schema_dict: dict[str, Any] = {
"node_types": [
{
"label": "Person",
"properties": [
{"name": "name", "type": "STRING"},
{"name": "email", "type": "STRING"},
],
}
],
"constraints": [
{
"type": "EXISTENCE",
"node_type": "Person",
"property_names": ["name"],
"relationship_type": None,
}
],
}
node = Neo4jNode(
id="n1",
label="Person",
properties={"name": "Alice", "email": "a@b.c"},
)
graph = Neo4jGraph(nodes=[node], relationships=[])
result = await writer.run(
graph=graph, schema=GraphSchema.model_validate(schema_dict)
)
assert result.status == "SUCCESS"
assert result.metadata is not None
node_file = next(f for f in result.metadata["files"] if f["is_node"])
existence_cs = [c for c in node_file["constraints"] if c["type"] == "EXISTENCE"]
assert existence_cs == [{"type": "EXISTENCE", "properties": ["name"]}]
key_cs = [c for c in node_file["constraints"] if c["type"] == "KEY"]
assert key_cs == [{"type": "KEY", "properties": [INTERNAL_ID_PROPERTY]}]
cols = {c["name"]: c for c in node_file["columns"]}
assert cols["name"]["is_primary_key"] is False
assert cols["name"]["is_unique"] is False
@pytest.mark.asyncio
async def test_parquet_writer_key_uniqueness_existence_constraints_metadata() -> None:
"""KEY, UNIQUENESS, and EXISTENCE all appear in constraints; EXISTENCE does not alter column flags."""
pytest.importorskip("pyarrow")
with tempfile.TemporaryDirectory() as tmpdir:
out = Path(tmpdir)
dest = _LocalParquetDestination(out)
writer = ParquetWriter(
nodes_dest=dest,
relationships_dest=dest,
collision_handler=FilenameCollisionHandler(),
)
schema_dict: dict[str, Any] = {
"node_types": [
{
"label": "Person",
"properties": [
{"name": "email", "type": "STRING"},
{"name": "external_id", "type": "STRING"},
{"name": "name", "type": "STRING"},
],
}
],
"constraints": [
{
"type": "KEY",
"node_type": "Person",
"property_names": ["email"],
"relationship_type": None,
},
{
"type": "UNIQUENESS",
"node_type": "Person",
"property_names": ["external_id"],
"relationship_type": None,
},
{
"type": "EXISTENCE",
"node_type": "Person",
"property_names": ["name"],
"relationship_type": None,
},
],
}
node = Neo4jNode(
id="n1",
label="Person",
properties={
"email": "a@b.c",
"external_id": "ext-1",
"name": "Alice",
},
)
graph = Neo4jGraph(nodes=[node], relationships=[])
result = await writer.run(
graph=graph, schema=GraphSchema.model_validate(schema_dict)
)
assert result.status == "SUCCESS"
assert result.metadata is not None
node_file = next(f for f in result.metadata["files"] if f["is_node"])
cs = node_file["constraints"]
assert {"type": "KEY", "properties": ["email"]} in cs
assert {"type": "UNIQUENESS", "properties": ["external_id"]} in cs
assert {"type": "EXISTENCE", "properties": ["name"]} in cs
cols = {c["name"]: c for c in node_file["columns"]}
assert cols["email"]["is_primary_key"] is True
assert cols["email"]["is_unique"] is False
assert cols["external_id"]["is_primary_key"] is False
assert cols["external_id"]["is_unique"] is True
assert cols["name"]["is_primary_key"] is False
assert cols["name"]["is_unique"] is False
@pytest.mark.asyncio
async def test_parquet_writer_run_empty_graph() -> None:
"""ParquetWriter accepts an empty graph and writes no files."""
pytest.importorskip("pyarrow")
with tempfile.TemporaryDirectory() as tmpdir:
out = Path(tmpdir)
dest = _LocalParquetDestination(out)
writer = ParquetWriter(
nodes_dest=dest,
relationships_dest=dest,
collision_handler=FilenameCollisionHandler(),
)
graph = Neo4jGraph(nodes=[], relationships=[])
result = await writer.run(graph=graph)
assert result.status == "SUCCESS"
assert result.metadata is not None
stats = result.metadata.get("statistics") or {}
assert stats["node_count"] == 0
assert stats["relationship_count"] == 0
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"}
# ---------------------------------------------------------------------------
# Regression tests: node embedding column must be present regardless of row order
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"failed_first",
[
pytest.param(True, id="failed_batch_first"),
pytest.param(False, id="succeeded_batch_first"),
],
)
def test_node_embedding_column_present_regardless_of_row_order(
failed_first: bool,
) -> None:
"""Embedding column must exist in the Parquet table regardless of which rows come first.
Regression test for the bug where failed-batch nodes (empty embedding_properties)
appearing before succeeded-batch nodes caused PyArrow to omit the embedding column
entirely from the inferred schema.
"""
pytest.importorskip("pyarrow")
import pyarrow as pa
import pyarrow.parquet as pq
formatter = Neo4jGraphParquetFormatter()
# Rows simulating two batches of the same node label:
# - failed-batch row: no embedding key (as if embedding_properties was empty)
# - succeeded-batch row: embedding key present
failed_row: dict[str, Any] = {
INTERNAL_ID_PROPERTY: "node-1",
"name": "Alice",
"labels": ["Person", "__Entity__"],
}
succeeded_row: dict[str, Any] = {
INTERNAL_ID_PROPERTY: "node-2",
"name": "Bob",
"labels": ["Person", "__Entity__"],
"embedding": [0.1, 0.2, 0.3],
}
rows = [failed_row, succeeded_row] if failed_first else [succeeded_row, failed_row]
parquet_bytes, schema = formatter.format_parquet(rows, "node label 'Person'")
# The embedding column must always be present in the schema
assert "embedding" in schema.names, (
f"'embedding' column missing from schema when failed_first={failed_first}. "
f"Schema columns: {schema.names}"
)
# Read back the table and verify nulls and types
table = pq.read_table(BytesIO(parquet_bytes))
assert "embedding" in table.column_names
# The row without an embedding should have a null value
rows_as_dicts = table.to_pylist()
rows_by_id = {r[INTERNAL_ID_PROPERTY]: r for r in rows_as_dicts}
assert (
rows_by_id["node-1"]["embedding"] is None
), "Row without embedding should have null value in the embedding column"
assert (
rows_by_id["node-2"]["embedding"] is not None
), "Row with embedding should have a non-null value in the embedding column"
# The embedding field type must be a list of floats (variable or fixed-size)
emb_field = schema.field("embedding")
emb_type = emb_field.type
# Because node-1 has null, the formatter must fall back to list_(float32)
assert pa.types.is_list(emb_type) or pa.types.is_fixed_size_list(
emb_type
), f"Unexpected embedding field type: {emb_type}"
# The value type must be float32
assert (
emb_type.value_type == pa.float32()
), f"Embedding value type should be float32, got {emb_type.value_type}"
# ---------------------------------------------------------------------------
# Regression tests: relationship embedding column must be present regardless of row order
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"failed_first",
[
pytest.param(True, id="failed_batch_first"),
pytest.param(False, id="succeeded_batch_first"),
],
)
def test_relationship_embedding_column_present_regardless_of_row_order(
failed_first: bool,
) -> None:
"""Embedding column must exist in the relationship Parquet table regardless of which rows come first.
Regression test for the bug where failed-batch relationships (empty embedding_properties)
appearing before succeeded-batch relationships caused PyArrow to omit the embedding column
entirely from the inferred schema.
"""
pytest.importorskip("pyarrow")
import pyarrow as pa
import pyarrow.parquet as pq
formatter = Neo4jGraphParquetFormatter()
# Rows simulating two batches of the same relationship type:
# - failed-batch row: no embedding key (as if embedding_properties was empty)
# - succeeded-batch row: embedding key present
failed_row: dict[str, Any] = {
"from": "node-1",
"to": "node-2",
"from_label": "Person",
"to_label": "Person",
"type": "KNOWS",
"since": "2020",
}
succeeded_row: dict[str, Any] = {
"from": "node-3",
"to": "node-4",
"from_label": "Person",
"to_label": "Person",
"type": "KNOWS",
"since": "2021",
"embedding": [0.1, 0.2, 0.3],
}
rows = [failed_row, succeeded_row] if failed_first else [succeeded_row, failed_row]
parquet_bytes, schema = formatter.format_parquet(
rows, "relationship 'Person_KNOWS_Person'"
)
# The embedding column must always be present in the schema
assert "embedding" in schema.names, (
f"'embedding' column missing from relationship schema when failed_first={failed_first}. "
f"Schema columns: {schema.names}"
)
# Read back the table and verify nulls and types
table = pq.read_table(BytesIO(parquet_bytes))
assert "embedding" in table.column_names
# The row without an embedding should have a null value
rows_as_dicts = table.to_pylist()
rows_by_from = {r["from"]: r for r in rows_as_dicts}
assert (
rows_by_from["node-1"]["embedding"] is None
), "Relationship row without embedding should have null value in the embedding column"
assert (
rows_by_from["node-3"]["embedding"] is not None
), "Relationship row with embedding should have a non-null value in the embedding column"
# The embedding field type must be a list of floats (variable or fixed-size)
emb_field = schema.field("embedding")
emb_type = emb_field.type
# Because the failed row has null, the formatter must fall back to list_(float32)
assert pa.types.is_list(emb_type) or pa.types.is_fixed_size_list(
emb_type
), f"Unexpected relationship embedding field type: {emb_type}"
# The value type must be float32
assert (
emb_type.value_type == pa.float32()
), f"Relationship embedding value type should be float32, got {emb_type.value_type}"
# ---------------------------------------------------------------------------
# Degenerate case: all rows lack the embedding key (all-null column path)
# ---------------------------------------------------------------------------
def test_format_parquet_all_rows_missing_embedding_does_not_crash() -> None:
"""format_parquet must not raise when no row has an embedding key.
When every row lacks a given key the formatter falls back to pa.null() for
that column's type. This test verifies that path doesn't crash and that
the resulting table contains only the columns that were actually present.
"""
pytest.importorskip("pyarrow")
import pyarrow.parquet as pq
formatter = Neo4jGraphParquetFormatter()
rows: list[dict[str, Any]] = [
{INTERNAL_ID_PROPERTY: "node-1", "name": "Alice", "labels": ["Person"]},
{INTERNAL_ID_PROPERTY: "node-2", "name": "Bob", "labels": ["Person"]},
]
parquet_bytes, schema = formatter.format_parquet(rows, "node label 'Person'")
assert (
"embedding" not in schema.names
), "Embedding column should not appear when no row carries an embedding key"
table = pq.read_table(BytesIO(parquet_bytes))
assert table.num_rows == 2
assert set(table.column_names) == {INTERNAL_ID_PROPERTY, "name", "labels"}
# ---------------------------------------------------------------------------
# Edge case: all rows have an empty list for the embedding key (all-null path)
# ---------------------------------------------------------------------------
def test_format_parquet_all_rows_empty_list_embedding_does_not_crash() -> None:
"""format_parquet must not raise when every row has an empty list for the embedding key.
When the sample dict filters out empty lists (they are falsy but not None, so
they pass the `v is not None` guard), pa.infer_type([[]]) returns list<null>.
This test verifies the resulting table survives a Parquet round-trip and that
all embedding values are empty lists.
"""
pytest.importorskip("pyarrow")
import pyarrow.parquet as pq
formatter = Neo4jGraphParquetFormatter()
rows: list[dict[str, Any]] = [
{
INTERNAL_ID_PROPERTY: "node-1",
"name": "Alice",
"labels": ["Person"],
"embedding": [],
},
{
INTERNAL_ID_PROPERTY: "node-2",
"name": "Bob",
"labels": ["Person"],
"embedding": [],
},
]
parquet_bytes, schema = formatter.format_parquet(rows, "node label 'Person'")
assert (
"embedding" in schema.names
), "Embedding column should be present even when all rows have an empty list"
table = pq.read_table(BytesIO(parquet_bytes))
assert table.num_rows == 2
assert "embedding" in table.column_names
for row in table.to_pylist():
assert (
row["embedding"] == []
), f"Expected empty list for embedding, got {row['embedding']}"
# ---------------------------------------------------------------------------
# Edge case: empty-list row before a float-list row must not crash
# ---------------------------------------------------------------------------
def test_format_parquet_empty_list_before_float_embedding_does_not_crash() -> None:
"""Empty-list row appearing before a float-list row must not raise.
If [] is picked up as the type-inference sample, pa.infer_type([[]]) returns
list<null>, which causes ArrowInvalid when writing the float-list row.
The fix skips both None and [] when collecting samples so the float-list row
always wins as the sample for embedding type inference.
"""
pytest.importorskip("pyarrow")
import pyarrow as pa
import pyarrow.parquet as pq
formatter = Neo4jGraphParquetFormatter()
rows: list[dict[str, Any]] = [
{
INTERNAL_ID_PROPERTY: "node-1",
"name": "Alice",
"labels": ["Person"],
"embedding": [],
},
{
INTERNAL_ID_PROPERTY: "node-2",
"name": "Bob",
"labels": ["Person"],
"embedding": [0.1, 0.2, 0.3],
},
]
parquet_bytes, schema = formatter.format_parquet(rows, "node label 'Person'")
assert "embedding" in schema.names
emb_type = schema.field("embedding").type
assert pa.types.is_list(emb_type), f"Unexpected embedding type: {emb_type}"
assert emb_type.value_type == pa.float32()
table = pq.read_table(BytesIO(parquet_bytes))
rows_by_id = {r[INTERNAL_ID_PROPERTY]: r for r in table.to_pylist()}
assert (
rows_by_id["node-1"]["embedding"] is None
or rows_by_id["node-1"]["embedding"] == []
)
assert rows_by_id["node-2"]["embedding"] is not None
Reference in New Issue Copy Permalink