From 453eeb7ca9a5a8fefb9358d2e09237792a8cb5ee Mon Sep 17 00:00:00 2001 From: KaifAhmad1 Date: Sat, 23 May 2026 00:14:23 +0530 Subject: [PATCH] fix: rename contributing/license pages to avoid Mintlify reserved slug conflict MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mint export fails with 'file does not exist' for pages named 'contributing' and 'license' ā€” these are reserved by Mintlify's GitHub integration layer. Renamed to contributing-guide.md and project-license.md and updated all nav entries and cross-links throughout the docs. Also adds .gitattributes LF rules to prevent CRLF issues from Windows devs. --- .gitattributes | 8 + docs/citation.md | 2 +- docs/community-projects.md | 2 +- docs/community.md | 2 +- ...{contributing.md => contributing-guide.md} | 0 docs/docs.json | 4 +- docs/faq.md | 2 +- docs/governance.md | 2 +- docs/{license.md => project-license.md} | 2 +- examples/parquet_export_example.py | 426 ++--- semantica/export/arango_aql_exporter.py | 1282 +++++++-------- semantica/export/parquet_exporter.py | 1400 ++++++++--------- tests/test_arango_aql_exporter.py | 1076 ++++++------- tests/test_parquet_exporter.py | 1092 ++++++------- 14 files changed, 2654 insertions(+), 2646 deletions(-) rename docs/{contributing.md => contributing-guide.md} (100%) rename docs/{license.md => project-license.md} (99%) diff --git a/.gitattributes b/.gitattributes index 5595f17b..aad0da1d 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,3 +1,11 @@ +# Line endings ā€” force LF so Mintlify/Linux CI parses frontmatter correctly +* text=auto eol=lf +*.md text eol=lf +*.json text eol=lf +*.yml text eol=lf +*.yaml text eol=lf +*.py text eol=lf + # Linguist documentation and generated files # This ensures GitHub language statistics reflect the core Python code diff --git a/docs/citation.md b/docs/citation.md index a937117e..fd4a6728 100644 --- a/docs/citation.md +++ b/docs/citation.md @@ -62,7 +62,7 @@ Published research using Semantica? [Let us know](https://github.com/semantica-a ## See Also - + MIT License details. diff --git a/docs/community-projects.md b/docs/community-projects.md index 9ef10b7f..c09e46f6 100644 --- a/docs/community-projects.md +++ b/docs/community-projects.md @@ -55,7 +55,7 @@ Have a project using Semantica? [Submit it on GitHub](https://github.com/semanti See the [Contributing Guide](contributing) for full details. - + Submit code, docs, or tests. diff --git a/docs/community.md b/docs/community.md index ab624904..17614a34 100644 --- a/docs/community.md +++ b/docs/community.md @@ -61,7 +61,7 @@ See the [Contributing Guide](contributing) for full details. ## See Also - + Step-by-step guide for submitting PRs. diff --git a/docs/contributing.md b/docs/contributing-guide.md similarity index 100% rename from docs/contributing.md rename to docs/contributing-guide.md diff --git a/docs/docs.json b/docs/docs.json index e9324bf4..2be3d248 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -194,10 +194,10 @@ "pages": [ "community", "community-projects", - "contributing", + "contributing-guide", "governance", "citation", - "license" + "project-license" ] } ] diff --git a/docs/faq.md b/docs/faq.md index 43572513..88af3cac 100644 --- a/docs/faq.md +++ b/docs/faq.md @@ -235,7 +235,7 @@ pip install --upgrade semantica Bug reports and feature requests. - + Help improve Semantica. diff --git a/docs/governance.md b/docs/governance.md index f769d67d..791279d8 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -98,7 +98,7 @@ MIT License ā€” see [LICENSE](https://github.com/semantica-agi/semantica/blob/ma ## See Also - + How to submit changes. diff --git a/docs/license.md b/docs/project-license.md similarity index 99% rename from docs/license.md rename to docs/project-license.md index 6e3871f0..bad87bef 100644 --- a/docs/license.md +++ b/docs/project-license.md @@ -83,7 +83,7 @@ By contributing to Semantica, you agree that your contributions will be licensed ## See Also - + How to contribute to the project. diff --git a/examples/parquet_export_example.py b/examples/parquet_export_example.py index 5dcb0db0..0f47e82c 100644 --- a/examples/parquet_export_example.py +++ b/examples/parquet_export_example.py @@ -1,213 +1,213 @@ -""" -Apache Parquet Exporter - Example Usage - -This script demonstrates how to use the ParquetExporter to export -knowledge graphs, entities, and relationships to Apache Parquet format. -""" - -import tempfile -from pathlib import Path - -from semantica.export import ParquetExporter, export_parquet - - -def main(): - print("=" * 70) - print("Apache Parquet Exporter - Example Usage") - print("=" * 70) - - # Create a temporary directory for outputs - temp_dir = Path(tempfile.mkdtemp()) - print(f"\nšŸ“ Output directory: {temp_dir}\n") - - # Sample data - entities = [ - { - "id": "e1", - "text": "Alice", - "type": "Person", - "confidence": 0.95, - "start": 0, - "end": 5, - "metadata": {"age": 30, "city": "New York"}, - }, - { - "id": "e2", - "text": "Acme Corp", - "type": "Organization", - "confidence": 0.88, - "start": 10, - "end": 19, - "metadata": {"location": "NY", "employees": 100}, - }, - { - "id": "e3", - "text": "Bob", - "type": "Person", - "confidence": 0.92, - "metadata": {"age": 35, "department": "Engineering"}, - }, - ] - - relationships = [ - { - "id": "r1", - "source_id": "e1", - "target_id": "e2", - "type": "WORKS_FOR", - "confidence": 0.90, - "metadata": {"role": "Engineer", "since": 2020}, - }, - { - "id": "r2", - "source_id": "e3", - "target_id": "e2", - "type": "WORKS_FOR", - "confidence": 0.85, - "metadata": {"role": "Manager", "since": 2018}, - }, - ] - - knowledge_graph = { - "entities": entities, - "relationships": relationships, - "metadata": {"version": "1.0", "created": "2024-01-01"}, - } - - # Example 1: Export entities using ParquetExporter class - print("Example 1: Export entities to Parquet") - print("-" * 70) - exporter = ParquetExporter(compression="snappy") - entities_path = temp_dir / "entities.parquet" - exporter.export_entities(entities, entities_path) - print(f"āœ“ Entities exported to: {entities_path}") - print(f" File size: {entities_path.stat().st_size} bytes\n") - - # Example 2: Export relationships - print("Example 2: Export relationships to Parquet") - print("-" * 70) - rels_path = temp_dir / "relationships.parquet" - exporter.export_relationships(relationships, rels_path) - print(f"āœ“ Relationships exported to: {rels_path}") - print(f" File size: {rels_path.stat().st_size} bytes\n") - - # Example 3: Export complete knowledge graph - print("Example 3: Export knowledge graph to multiple Parquet files") - print("-" * 70) - kg_base_path = temp_dir / "knowledge_graph" - exporter.export_knowledge_graph(knowledge_graph, kg_base_path) - kg_entities = temp_dir / "knowledge_graph_entities.parquet" - kg_rels = temp_dir / "knowledge_graph_relationships.parquet" - print("āœ“ Knowledge graph exported to:") - print(f" - {kg_entities} ({kg_entities.stat().st_size} bytes)") - print(f" - {kg_rels} ({kg_rels.stat().st_size} bytes)\n") - - # Example 4: Using convenience function - print("Example 4: Using export_parquet convenience function") - print("-" * 70) - conv_path = temp_dir / "convenience_export.parquet" - export_parquet(entities, conv_path, compression="gzip") - print(f"āœ“ Exported using convenience function: {conv_path}") - print(f" File size: {conv_path.stat().st_size} bytes\n") - - # Example 5: Different compression codecs - print("Example 5: Compare compression codecs") - print("-" * 70) - - # Create larger dataset for meaningful comparison - large_entities = entities * 50 - - compression_codecs = ["snappy", "gzip", "brotli", "zstd", "lz4", "none"] - sizes = {} - - for codec in compression_codecs: - codec_exporter = ParquetExporter(compression=codec) - codec_path = temp_dir / f"entities_{codec}.parquet" - codec_exporter.export_entities(large_entities, codec_path) - sizes[codec] = codec_path.stat().st_size - print(f" {codec:8} - {sizes[codec]:,} bytes") - - print() - - # Example 6: Load Parquet with pandas (if available) - print("Example 6: Loading Parquet files with pandas") - print("-" * 70) - try: - import pandas as pd - - df = pd.read_parquet(entities_path) - print("āœ“ Loaded entities as pandas DataFrame") - print(f" Shape: {df.shape}") - print(f" Columns: {list(df.columns)}") - print("\nFirst few rows:") - print(df.head()) - print() - - except ImportError: - print("āš  pandas not installed - skipping pandas example\n") - - # Example 7: Load Parquet with pyarrow - print("Example 7: Loading Parquet files with pyarrow") - print("-" * 70) - try: - import pyarrow.parquet as pq - - table = pq.read_table(entities_path) - print("āœ“ Loaded entities as Arrow Table") - print(f" Rows: {table.num_rows}") - print(f" Columns: {table.num_columns}") - print(" Schema:") - for i, field in enumerate(table.schema): - print(f" - {field.name}: {field.type}") - print() - - except ImportError: - print("āš  pyarrow not installed - skipping pyarrow example\n") - - # Example 8: Schema validation - print("Example 8: Explicit schema validation") - print("-" * 70) - try: - import pyarrow.parquet as pq - - # Read parquet file and verify schema - table = pq.read_table(entities_path) - - print("āœ“ Schema validation:") - print(f" - ID column type: {table.schema.field('id').type}") - print(f" - Text column type: {table.schema.field('text').type}") - print(f" - Confidence column type: {table.schema.field('confidence').type}") - print(f" - Metadata column type: {table.schema.field('metadata').type}") - print() - - # Verify metadata structure - metadata_field = table.schema.field("metadata") - print(" Metadata structure:") - if hasattr(metadata_field.type, "num_fields"): - for i in range(metadata_field.type.num_fields): - subfield = metadata_field.type.field(i) - print(f" - {subfield.name}: {subfield.type}") - print() - - except Exception as e: - print(f"āš  Schema validation error: {e}\n") - - # Summary - print("=" * 70) - print("Summary") - print("=" * 70) - print("āœ“ All examples completed successfully") - print(f"āœ“ Output directory: {temp_dir}") - print(f"āœ“ Files created: {len(list(temp_dir.glob('*.parquet')))}") - print("\nKey Features:") - print(" - Columnar storage optimized for analytics") - print(" - Multiple compression options (snappy, gzip, brotli, zstd, lz4)") - print(" - Compatible with pandas, Spark, Snowflake, BigQuery, Databricks") - print(" - Explicit schemas for type safety") - print(" - Structured metadata handling") - print("\nFor more information, see the Semantica documentation.") - print("=" * 70) - - -if __name__ == "__main__": - main() +""" +Apache Parquet Exporter - Example Usage + +This script demonstrates how to use the ParquetExporter to export +knowledge graphs, entities, and relationships to Apache Parquet format. +""" + +import tempfile +from pathlib import Path + +from semantica.export import ParquetExporter, export_parquet + + +def main(): + print("=" * 70) + print("Apache Parquet Exporter - Example Usage") + print("=" * 70) + + # Create a temporary directory for outputs + temp_dir = Path(tempfile.mkdtemp()) + print(f"\nšŸ“ Output directory: {temp_dir}\n") + + # Sample data + entities = [ + { + "id": "e1", + "text": "Alice", + "type": "Person", + "confidence": 0.95, + "start": 0, + "end": 5, + "metadata": {"age": 30, "city": "New York"}, + }, + { + "id": "e2", + "text": "Acme Corp", + "type": "Organization", + "confidence": 0.88, + "start": 10, + "end": 19, + "metadata": {"location": "NY", "employees": 100}, + }, + { + "id": "e3", + "text": "Bob", + "type": "Person", + "confidence": 0.92, + "metadata": {"age": 35, "department": "Engineering"}, + }, + ] + + relationships = [ + { + "id": "r1", + "source_id": "e1", + "target_id": "e2", + "type": "WORKS_FOR", + "confidence": 0.90, + "metadata": {"role": "Engineer", "since": 2020}, + }, + { + "id": "r2", + "source_id": "e3", + "target_id": "e2", + "type": "WORKS_FOR", + "confidence": 0.85, + "metadata": {"role": "Manager", "since": 2018}, + }, + ] + + knowledge_graph = { + "entities": entities, + "relationships": relationships, + "metadata": {"version": "1.0", "created": "2024-01-01"}, + } + + # Example 1: Export entities using ParquetExporter class + print("Example 1: Export entities to Parquet") + print("-" * 70) + exporter = ParquetExporter(compression="snappy") + entities_path = temp_dir / "entities.parquet" + exporter.export_entities(entities, entities_path) + print(f"āœ“ Entities exported to: {entities_path}") + print(f" File size: {entities_path.stat().st_size} bytes\n") + + # Example 2: Export relationships + print("Example 2: Export relationships to Parquet") + print("-" * 70) + rels_path = temp_dir / "relationships.parquet" + exporter.export_relationships(relationships, rels_path) + print(f"āœ“ Relationships exported to: {rels_path}") + print(f" File size: {rels_path.stat().st_size} bytes\n") + + # Example 3: Export complete knowledge graph + print("Example 3: Export knowledge graph to multiple Parquet files") + print("-" * 70) + kg_base_path = temp_dir / "knowledge_graph" + exporter.export_knowledge_graph(knowledge_graph, kg_base_path) + kg_entities = temp_dir / "knowledge_graph_entities.parquet" + kg_rels = temp_dir / "knowledge_graph_relationships.parquet" + print("āœ“ Knowledge graph exported to:") + print(f" - {kg_entities} ({kg_entities.stat().st_size} bytes)") + print(f" - {kg_rels} ({kg_rels.stat().st_size} bytes)\n") + + # Example 4: Using convenience function + print("Example 4: Using export_parquet convenience function") + print("-" * 70) + conv_path = temp_dir / "convenience_export.parquet" + export_parquet(entities, conv_path, compression="gzip") + print(f"āœ“ Exported using convenience function: {conv_path}") + print(f" File size: {conv_path.stat().st_size} bytes\n") + + # Example 5: Different compression codecs + print("Example 5: Compare compression codecs") + print("-" * 70) + + # Create larger dataset for meaningful comparison + large_entities = entities * 50 + + compression_codecs = ["snappy", "gzip", "brotli", "zstd", "lz4", "none"] + sizes = {} + + for codec in compression_codecs: + codec_exporter = ParquetExporter(compression=codec) + codec_path = temp_dir / f"entities_{codec}.parquet" + codec_exporter.export_entities(large_entities, codec_path) + sizes[codec] = codec_path.stat().st_size + print(f" {codec:8} - {sizes[codec]:,} bytes") + + print() + + # Example 6: Load Parquet with pandas (if available) + print("Example 6: Loading Parquet files with pandas") + print("-" * 70) + try: + import pandas as pd + + df = pd.read_parquet(entities_path) + print("āœ“ Loaded entities as pandas DataFrame") + print(f" Shape: {df.shape}") + print(f" Columns: {list(df.columns)}") + print("\nFirst few rows:") + print(df.head()) + print() + + except ImportError: + print("āš  pandas not installed - skipping pandas example\n") + + # Example 7: Load Parquet with pyarrow + print("Example 7: Loading Parquet files with pyarrow") + print("-" * 70) + try: + import pyarrow.parquet as pq + + table = pq.read_table(entities_path) + print("āœ“ Loaded entities as Arrow Table") + print(f" Rows: {table.num_rows}") + print(f" Columns: {table.num_columns}") + print(" Schema:") + for i, field in enumerate(table.schema): + print(f" - {field.name}: {field.type}") + print() + + except ImportError: + print("āš  pyarrow not installed - skipping pyarrow example\n") + + # Example 8: Schema validation + print("Example 8: Explicit schema validation") + print("-" * 70) + try: + import pyarrow.parquet as pq + + # Read parquet file and verify schema + table = pq.read_table(entities_path) + + print("āœ“ Schema validation:") + print(f" - ID column type: {table.schema.field('id').type}") + print(f" - Text column type: {table.schema.field('text').type}") + print(f" - Confidence column type: {table.schema.field('confidence').type}") + print(f" - Metadata column type: {table.schema.field('metadata').type}") + print() + + # Verify metadata structure + metadata_field = table.schema.field("metadata") + print(" Metadata structure:") + if hasattr(metadata_field.type, "num_fields"): + for i in range(metadata_field.type.num_fields): + subfield = metadata_field.type.field(i) + print(f" - {subfield.name}: {subfield.type}") + print() + + except Exception as e: + print(f"āš  Schema validation error: {e}\n") + + # Summary + print("=" * 70) + print("Summary") + print("=" * 70) + print("āœ“ All examples completed successfully") + print(f"āœ“ Output directory: {temp_dir}") + print(f"āœ“ Files created: {len(list(temp_dir.glob('*.parquet')))}") + print("\nKey Features:") + print(" - Columnar storage optimized for analytics") + print(" - Multiple compression options (snappy, gzip, brotli, zstd, lz4)") + print(" - Compatible with pandas, Spark, Snowflake, BigQuery, Databricks") + print(" - Explicit schemas for type safety") + print(" - Structured metadata handling") + print("\nFor more information, see the Semantica documentation.") + print("=" * 70) + + +if __name__ == "__main__": + main() diff --git a/semantica/export/arango_aql_exporter.py b/semantica/export/arango_aql_exporter.py index 712cb1d3..69140481 100644 --- a/semantica/export/arango_aql_exporter.py +++ b/semantica/export/arango_aql_exporter.py @@ -1,641 +1,641 @@ -""" -ArangoDB AQL Export Module - -This module provides comprehensive AQL export capabilities for the Semantica framework, -enabling export to ArangoDB multi-model graph databases. - -Key Features: - - AQL format export for ArangoDB - - INSERT statement generation for vertex and edge collections - - Configurable collection names - - Entity and relationship identifier preservation - - Batch insert support for performance - - Proper string escaping and special character handling - -Example Usage: - >>> from semantica.export import ArangoAQLExporter - >>> exporter = ArangoAQLExporter() - >>> exporter.export_knowledge_graph(kg, "output.aql") - >>> # With custom collection names - >>> exporter = ArangoAQLExporter( - ... vertex_collection="nodes", - ... edge_collection="links" - ... ) - >>> exporter.export(kg, "graph.aql") -""" - -import json -from pathlib import Path -from typing import Any, Dict, List, Optional, Union - -from ..utils.helpers import ensure_directory -from ..utils.logging import get_logger -from ..utils.progress_tracker import get_progress_tracker - - -class ArangoAQLExporter: - """ - ArangoDB AQL exporter for knowledge graphs. - - This class provides comprehensive AQL export functionality for knowledge graphs, - supporting export to ArangoDB multi-model databases via AQL INSERT statements. - - Features: - - AQL INSERT statement generation for vertices and edges - - Configurable collection names - - Entity and relationship identifier preservation - - Batch insert support for performance - - Proper string escaping and special character handling - - Support for nested properties via JSON serialization - - Example Usage: - >>> exporter = ArangoAQLExporter( - ... vertex_collection="entities", - ... edge_collection="relationships" - ... ) - >>> exporter.export_knowledge_graph(kg, "output.aql") - """ - - def __init__( - self, - vertex_collection: str = "vertices", - edge_collection: str = "edges", - batch_size: int = 1000, - include_collection_creation: bool = True, - config: Optional[Dict[str, Any]] = None, - **kwargs, - ): - """ - Initialize ArangoDB AQL exporter. - - Sets up the exporter with collection names and batch processing options. - - Args: - vertex_collection: Name of the vertex collection - (default: "vertices") - edge_collection: Name of the edge collection (default: "edges") - batch_size: Batch size for INSERT operations (default: 1000) - include_collection_creation: Whether to include collection - creation statements (default: True) - config: Optional configuration dictionary (merged with kwargs) - **kwargs: Additional configuration options - """ - self.logger = get_logger("arango_aql_exporter") - self.config = config or {} - self.config.update(kwargs) - - # Validate collection names - self._validate_collection_name(vertex_collection, "vertex_collection") - self._validate_collection_name(edge_collection, "edge_collection") - - # AQL export configuration - self.vertex_collection = vertex_collection - self.edge_collection = edge_collection - self.batch_size = batch_size - self.include_collection_creation = include_collection_creation - - # Initialize progress tracker - self.progress_tracker = get_progress_tracker() - - self.logger.debug( - f"ArangoDB AQL exporter initialized: " - f"vertex_collection={vertex_collection}, " - f"edge_collection={edge_collection}, batch_size={batch_size}" - ) - - def export( - self, knowledge_graph: Dict[str, Any], file_path: Union[str, Path], **options - ) -> None: - """ - Export knowledge graph to ArangoDB AQL format. - - This method exports a knowledge graph to AQL INSERT statements that can - be imported into ArangoDB multi-model databases. - - Args: - knowledge_graph: Knowledge graph dictionary containing: - - entities: List of entity dictionaries - - relationships: List of relationship dictionaries - - nodes: List of node dictionaries (optional, alternative to - entities) - - edges: List of edge dictionaries (optional, alternative to - relationships) - file_path: Output AQL file path - **options: Additional export options: - - vertex_collection: Override default vertex collection name - - edge_collection: Override default edge collection name - - Example: - >>> kg = { - ... "entities": [...], - ... "relationships": [...] - ... } - >>> exporter.export(kg, "graph.aql") - """ - file_path = Path(file_path) - ensure_directory(file_path.parent) - - tracking_id = self.progress_tracker.start_tracking( - file=str(file_path), - module="export", - submodule="ArangoAQLExporter", - message=f"Exporting knowledge graph to ArangoDB AQL format: {file_path}", - ) - - try: - # Override collection names if provided in options - vertex_collection = options.pop("vertex_collection", self.vertex_collection) - edge_collection = options.pop("edge_collection", self.edge_collection) - - # Validate overridden collection names - if vertex_collection != self.vertex_collection: - self._validate_collection_name(vertex_collection, "vertex_collection") - if edge_collection != self.edge_collection: - self._validate_collection_name(edge_collection, "edge_collection") - - # Generate AQL statements - aql_statements = self._generate_aql_statements( - knowledge_graph, vertex_collection, edge_collection, **options - ) - - # Write to file - with open(file_path, "w", encoding="utf-8") as f: - f.write("\n".join(aql_statements)) - - self.progress_tracker.stop_tracking( - tracking_id, - status="completed", - message=f"Exported {len(aql_statements)} AQL statements", - ) - self.logger.info( - f"Exported knowledge graph to ArangoDB AQL format: " f"{file_path}" - ) - - except Exception as e: - self.progress_tracker.stop_tracking( - tracking_id, status="failed", message=str(e) - ) - raise - - def _generate_aql_statements( - self, - knowledge_graph: Dict[str, Any], - vertex_collection: str, - edge_collection: str, - **options, - ) -> List[str]: - """ - Generate AQL INSERT statements from knowledge graph. - - Args: - knowledge_graph: Knowledge graph dictionary - vertex_collection: Vertex collection name - edge_collection: Edge collection name - **options: Additional options - - Returns: - List of AQL statement strings - """ - statements = [] - - # Add collection creation statements if requested - if self.include_collection_creation: - statements.extend( - self._generate_collection_creation(vertex_collection, edge_collection) - ) - - # Extract entities and relationships - entities = knowledge_graph.get("entities", []) - relationships = knowledge_graph.get("relationships", []) - nodes = knowledge_graph.get("nodes", entities) - edges = knowledge_graph.get("edges", relationships) - - # Use nodes/edges if entities/relationships are empty - if not entities and nodes: - entities = nodes - if not relationships and edges: - relationships = edges - - # Generate vertex INSERT statements - vertex_statements = self._generate_vertex_inserts(entities, vertex_collection) - statements.extend(vertex_statements) - - # Generate edge INSERT statements - edge_statements = self._generate_edge_inserts( - relationships, edge_collection, vertex_collection - ) - statements.extend(edge_statements) - - return statements - - def _generate_collection_creation( - self, vertex_collection: str, edge_collection: str - ) -> List[str]: - """ - Generate AQL collection creation statements. - - Args: - vertex_collection: Vertex collection name - edge_collection: Edge collection name - - Returns: - List of collection creation statements - """ - statements = [ - "// Create vertex collection if it doesn't exist", - f"// db._createDocumentCollection('{vertex_collection}');", - "", - "// Create edge collection if it doesn't exist", - f"// db._createEdgeCollection('{edge_collection}');", - "", - ] - return statements - - def _generate_vertex_inserts( - self, vertices: List[Dict[str, Any]], collection: str - ) -> List[str]: - """ - Generate AQL INSERT statements for vertices. - - Args: - vertices: List of vertex/entity dictionaries - collection: Vertex collection name - - Returns: - List of AQL INSERT statements - """ - statements = [] - - # Add header comment - statements.append(f"// Inserting {len(vertices)} vertices into {collection}") - statements.append("") - - # Process vertices in batches - for i in range(0, len(vertices), self.batch_size): - batch = vertices[i : i + self.batch_size] - batch_statement = self._create_vertex_batch_insert(batch, collection) - statements.append(batch_statement) - statements.append("") - - return statements - - def _create_vertex_batch_insert( - self, vertices: List[Dict[str, Any]], collection: str - ) -> str: - """ - Create a batch INSERT statement for vertices. - - Args: - vertices: Batch of vertex dictionaries - collection: Vertex collection name - - Returns: - AQL INSERT statement - """ - if not vertices: - return "" - - # Build document list - documents = [] - for idx, vertex in enumerate(vertices): - doc = self._convert_vertex_to_document(vertex, idx) - documents.append(doc) - - # Format as AQL - docs_json = json.dumps(documents, indent=2, ensure_ascii=False) - statement = f"FOR doc IN {docs_json}\n INSERT doc INTO {collection}" - - return statement - - def _convert_vertex_to_document( - self, vertex: Dict[str, Any], idx: int - ) -> Dict[str, Any]: - """ - Convert a vertex/entity to an ArangoDB document. - - Additional properties from the input entity are preserved as-is in the - document root (not flattened or merged). Nested dictionaries and lists - are preserved as JSON-serializable structures. The 'properties' field, - if present, is also preserved as-is rather than being flattened into - the document root. - - Args: - vertex: Vertex/entity dictionary - idx: Index for generating fallback IDs - - Returns: - ArangoDB document dictionary - """ - document = {} - - # Set _key from id or generate one - vertex_id = vertex.get("id") or vertex.get("entity_id") or f"vertex_{idx}" - document["_key"] = self._sanitize_key(str(vertex_id)) - - # Add original ID if different from _key - if str(vertex_id) != document["_key"]: - document["original_id"] = str(vertex_id) - - # Add type/label information - vertex_type = vertex.get("type") or vertex.get("entity_type", "Entity") - document["type"] = vertex_type - - # Add name/label - label = ( - vertex.get("label") - or vertex.get("name") - or vertex.get("text") - or document["_key"] - ) - document["name"] = label - - # Add all other properties - for key, value in vertex.items(): - if key not in [ - "_key", - "_id", - "_rev", - "id", - "entity_id", - "type", - "entity_type", - "label", - "name", - "text", - ]: - # Handle nested dictionaries and lists - if isinstance(value, (dict, list)): - document[key] = value - elif value is not None: - document[key] = value - - return document - - def _generate_edge_inserts( - self, edges: List[Dict[str, Any]], collection: str, vertex_collection: str - ) -> List[str]: - """ - Generate AQL INSERT statements for edges. - - Args: - edges: List of edge/relationship dictionaries - collection: Edge collection name - vertex_collection: Vertex collection name for _from/_to references - - Returns: - List of AQL INSERT statements - """ - statements = [] - - # Add header comment - statements.append( - f"// Attempting to insert {len(edges)} edges into {collection}" - ) - statements.append("") - - # Process edges in batches - for i in range(0, len(edges), self.batch_size): - batch = edges[i : i + self.batch_size] - batch_statement = self._create_edge_batch_insert( - batch, collection, vertex_collection - ) - if batch_statement: # Only add non-empty statements - statements.append(batch_statement) - statements.append("") - - return statements - - def _create_edge_batch_insert( - self, edges: List[Dict[str, Any]], collection: str, vertex_collection: str - ) -> str: - """ - Create a batch INSERT statement for edges. - - Args: - edges: Batch of edge dictionaries - collection: Edge collection name - vertex_collection: Vertex collection name for _from/_to references - - Returns: - AQL INSERT statement - """ - # Build document list - documents = [] - for idx, edge in enumerate(edges): - doc = self._convert_edge_to_document(edge, idx, vertex_collection) - if doc: # Only add valid edges - documents.append(doc) - - if not documents: - return "" - - # Format as AQL - docs_json = json.dumps(documents, indent=2, ensure_ascii=False) - statement = f"FOR doc IN {docs_json}\n INSERT doc INTO {collection}" - - return statement - - def _convert_edge_to_document( - self, edge: Dict[str, Any], idx: int, vertex_collection: str - ) -> Optional[Dict[str, Any]]: - """ - Convert an edge/relationship to an ArangoDB edge document. - - Args: - edge: Edge/relationship dictionary - idx: Index for generating fallback IDs - vertex_collection: Vertex collection name for _from/_to references - - Returns: - ArangoDB edge document dictionary, or None if source/target missing - """ - # Extract source and target - source_id = edge.get("source") or edge.get("source_id") - target_id = edge.get("target") or edge.get("target_id") - - if not source_id or not target_id: - self.logger.warning( - f"Skipping edge {idx}: missing source or target " - f"(source={source_id}, target={target_id})" - ) - return None - - document = {} - - # Set _key from id or generate one - edge_id = edge.get("id") or edge.get("relationship_id") or f"edge_{idx}" - document["_key"] = self._sanitize_key(str(edge_id)) - - # Set _from and _to (required for edges in ArangoDB) - document["_from"] = f"{vertex_collection}/{self._sanitize_key(str(source_id))}" - document["_to"] = f"{vertex_collection}/{self._sanitize_key(str(target_id))}" - - # Add original ID if different from _key - if str(edge_id) != document["_key"]: - document["original_id"] = str(edge_id) - - # Add relationship type - rel_type = edge.get("type") or edge.get("relationship_type", "RELATED_TO") - document["type"] = rel_type - - # Add all other properties - for key, value in edge.items(): - if key not in [ - "_key", - "_id", - "_rev", - "_from", - "_to", - "id", - "relationship_id", - "source", - "source_id", - "target", - "target_id", - "type", - "relationship_type", - ]: - # Handle nested dictionaries and lists - if isinstance(value, (dict, list)): - document[key] = value - elif value is not None: - document[key] = value - - return document - - def _validate_collection_name(self, name: str, param_name: str) -> None: - """ - Validate an ArangoDB collection name. - - ArangoDB collection names must: - - Start with a letter or underscore - - Contain only alphanumeric characters, hyphens, and underscores - - Not exceed 256 characters - - Args: - name: Collection name to validate - param_name: Parameter name for error messages - - Raises: - ValueError: If the collection name is invalid - """ - if not name: - raise ValueError(f"{param_name} cannot be empty") - - if len(name) > 256: - raise ValueError( - f"{param_name} '{name}' exceeds maximum length of 256 characters" - ) - - # Check first character - if not (name[0].isalpha() or name[0] == "_"): - raise ValueError( - f"{param_name} '{name}' must start with a letter or underscore" - ) - - # Check remaining characters - for char in name: - if not (char.isalnum() or char in ("-", "_")): - raise ValueError( - f"{param_name} '{name}' contains invalid character " - f"'{char}'. Only alphanumeric characters, hyphens, and " - "underscores are allowed." - ) - - def _sanitize_key(self, key: str) -> str: - """ - Sanitize a key for use as ArangoDB _key. - - ArangoDB _key must contain only alphanumeric characters, hyphens, - and underscores. It cannot start with an underscore (unless it's - a system collection). - - Args: - key: Original key string - - Returns: - Sanitized key string - """ - # Replace invalid characters with underscores - sanitized = "" - for char in key: - if char.isalnum() or char in ("-", "_"): - sanitized += char - else: - sanitized += "_" - - # Ensure key doesn't start with underscore - if sanitized.startswith("_"): - sanitized = "k" + sanitized - - # Ensure key is not empty - if not sanitized: - sanitized = "key" - - return sanitized - - def export_knowledge_graph( - self, knowledge_graph: Dict[str, Any], file_path: Union[str, Path], **options - ) -> None: - """ - Export knowledge graph to ArangoDB AQL format. - - Convenience method that calls export(). - - Args: - knowledge_graph: Knowledge graph dictionary - file_path: Output AQL file path - **options: Additional export options - - Example: - >>> kg = { - ... "entities": [...], - ... "relationships": [...] - ... } - >>> exporter.export_knowledge_graph(kg, "output.aql") - """ - self.export(knowledge_graph, file_path, **options) - - def export_entities( - self, entities: List[Dict[str, Any]], file_path: Union[str, Path], **options - ) -> None: - """ - Export entities to ArangoDB AQL format (vertices only). - - Args: - entities: List of entity dictionaries - file_path: Output AQL file path - **options: Additional export options - - Example: - >>> entities = [ - ... {"id": "e1", "type": "Person", "name": "Alice"}, - ... {"id": "e2", "type": "Organization", "name": "Acme Corp"} - ... ] - >>> exporter.export_entities(entities, "entities.aql") - """ - knowledge_graph = {"entities": entities, "relationships": []} - self.export(knowledge_graph, file_path, **options) - - def export_relationships( - self, - relationships: List[Dict[str, Any]], - file_path: Union[str, Path], - **options, - ) -> None: - """ - Export relationships to ArangoDB AQL format (edges only). - - Args: - relationships: List of relationship dictionaries - file_path: Output AQL file path - **options: Additional export options - - Example: - >>> relationships = [ - ... {"id": "r1", "source": "e1", "target": "e2", "type": "WORKS_FOR"} - ... ] - >>> exporter.export_relationships(relationships, "relationships.aql") - """ - knowledge_graph = {"entities": [], "relationships": relationships} - self.export(knowledge_graph, file_path, **options) +""" +ArangoDB AQL Export Module + +This module provides comprehensive AQL export capabilities for the Semantica framework, +enabling export to ArangoDB multi-model graph databases. + +Key Features: + - AQL format export for ArangoDB + - INSERT statement generation for vertex and edge collections + - Configurable collection names + - Entity and relationship identifier preservation + - Batch insert support for performance + - Proper string escaping and special character handling + +Example Usage: + >>> from semantica.export import ArangoAQLExporter + >>> exporter = ArangoAQLExporter() + >>> exporter.export_knowledge_graph(kg, "output.aql") + >>> # With custom collection names + >>> exporter = ArangoAQLExporter( + ... vertex_collection="nodes", + ... edge_collection="links" + ... ) + >>> exporter.export(kg, "graph.aql") +""" + +import json +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from ..utils.helpers import ensure_directory +from ..utils.logging import get_logger +from ..utils.progress_tracker import get_progress_tracker + + +class ArangoAQLExporter: + """ + ArangoDB AQL exporter for knowledge graphs. + + This class provides comprehensive AQL export functionality for knowledge graphs, + supporting export to ArangoDB multi-model databases via AQL INSERT statements. + + Features: + - AQL INSERT statement generation for vertices and edges + - Configurable collection names + - Entity and relationship identifier preservation + - Batch insert support for performance + - Proper string escaping and special character handling + - Support for nested properties via JSON serialization + + Example Usage: + >>> exporter = ArangoAQLExporter( + ... vertex_collection="entities", + ... edge_collection="relationships" + ... ) + >>> exporter.export_knowledge_graph(kg, "output.aql") + """ + + def __init__( + self, + vertex_collection: str = "vertices", + edge_collection: str = "edges", + batch_size: int = 1000, + include_collection_creation: bool = True, + config: Optional[Dict[str, Any]] = None, + **kwargs, + ): + """ + Initialize ArangoDB AQL exporter. + + Sets up the exporter with collection names and batch processing options. + + Args: + vertex_collection: Name of the vertex collection + (default: "vertices") + edge_collection: Name of the edge collection (default: "edges") + batch_size: Batch size for INSERT operations (default: 1000) + include_collection_creation: Whether to include collection + creation statements (default: True) + config: Optional configuration dictionary (merged with kwargs) + **kwargs: Additional configuration options + """ + self.logger = get_logger("arango_aql_exporter") + self.config = config or {} + self.config.update(kwargs) + + # Validate collection names + self._validate_collection_name(vertex_collection, "vertex_collection") + self._validate_collection_name(edge_collection, "edge_collection") + + # AQL export configuration + self.vertex_collection = vertex_collection + self.edge_collection = edge_collection + self.batch_size = batch_size + self.include_collection_creation = include_collection_creation + + # Initialize progress tracker + self.progress_tracker = get_progress_tracker() + + self.logger.debug( + f"ArangoDB AQL exporter initialized: " + f"vertex_collection={vertex_collection}, " + f"edge_collection={edge_collection}, batch_size={batch_size}" + ) + + def export( + self, knowledge_graph: Dict[str, Any], file_path: Union[str, Path], **options + ) -> None: + """ + Export knowledge graph to ArangoDB AQL format. + + This method exports a knowledge graph to AQL INSERT statements that can + be imported into ArangoDB multi-model databases. + + Args: + knowledge_graph: Knowledge graph dictionary containing: + - entities: List of entity dictionaries + - relationships: List of relationship dictionaries + - nodes: List of node dictionaries (optional, alternative to + entities) + - edges: List of edge dictionaries (optional, alternative to + relationships) + file_path: Output AQL file path + **options: Additional export options: + - vertex_collection: Override default vertex collection name + - edge_collection: Override default edge collection name + + Example: + >>> kg = { + ... "entities": [...], + ... "relationships": [...] + ... } + >>> exporter.export(kg, "graph.aql") + """ + file_path = Path(file_path) + ensure_directory(file_path.parent) + + tracking_id = self.progress_tracker.start_tracking( + file=str(file_path), + module="export", + submodule="ArangoAQLExporter", + message=f"Exporting knowledge graph to ArangoDB AQL format: {file_path}", + ) + + try: + # Override collection names if provided in options + vertex_collection = options.pop("vertex_collection", self.vertex_collection) + edge_collection = options.pop("edge_collection", self.edge_collection) + + # Validate overridden collection names + if vertex_collection != self.vertex_collection: + self._validate_collection_name(vertex_collection, "vertex_collection") + if edge_collection != self.edge_collection: + self._validate_collection_name(edge_collection, "edge_collection") + + # Generate AQL statements + aql_statements = self._generate_aql_statements( + knowledge_graph, vertex_collection, edge_collection, **options + ) + + # Write to file + with open(file_path, "w", encoding="utf-8") as f: + f.write("\n".join(aql_statements)) + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Exported {len(aql_statements)} AQL statements", + ) + self.logger.info( + f"Exported knowledge graph to ArangoDB AQL format: " f"{file_path}" + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + raise + + def _generate_aql_statements( + self, + knowledge_graph: Dict[str, Any], + vertex_collection: str, + edge_collection: str, + **options, + ) -> List[str]: + """ + Generate AQL INSERT statements from knowledge graph. + + Args: + knowledge_graph: Knowledge graph dictionary + vertex_collection: Vertex collection name + edge_collection: Edge collection name + **options: Additional options + + Returns: + List of AQL statement strings + """ + statements = [] + + # Add collection creation statements if requested + if self.include_collection_creation: + statements.extend( + self._generate_collection_creation(vertex_collection, edge_collection) + ) + + # Extract entities and relationships + entities = knowledge_graph.get("entities", []) + relationships = knowledge_graph.get("relationships", []) + nodes = knowledge_graph.get("nodes", entities) + edges = knowledge_graph.get("edges", relationships) + + # Use nodes/edges if entities/relationships are empty + if not entities and nodes: + entities = nodes + if not relationships and edges: + relationships = edges + + # Generate vertex INSERT statements + vertex_statements = self._generate_vertex_inserts(entities, vertex_collection) + statements.extend(vertex_statements) + + # Generate edge INSERT statements + edge_statements = self._generate_edge_inserts( + relationships, edge_collection, vertex_collection + ) + statements.extend(edge_statements) + + return statements + + def _generate_collection_creation( + self, vertex_collection: str, edge_collection: str + ) -> List[str]: + """ + Generate AQL collection creation statements. + + Args: + vertex_collection: Vertex collection name + edge_collection: Edge collection name + + Returns: + List of collection creation statements + """ + statements = [ + "// Create vertex collection if it doesn't exist", + f"// db._createDocumentCollection('{vertex_collection}');", + "", + "// Create edge collection if it doesn't exist", + f"// db._createEdgeCollection('{edge_collection}');", + "", + ] + return statements + + def _generate_vertex_inserts( + self, vertices: List[Dict[str, Any]], collection: str + ) -> List[str]: + """ + Generate AQL INSERT statements for vertices. + + Args: + vertices: List of vertex/entity dictionaries + collection: Vertex collection name + + Returns: + List of AQL INSERT statements + """ + statements = [] + + # Add header comment + statements.append(f"// Inserting {len(vertices)} vertices into {collection}") + statements.append("") + + # Process vertices in batches + for i in range(0, len(vertices), self.batch_size): + batch = vertices[i : i + self.batch_size] + batch_statement = self._create_vertex_batch_insert(batch, collection) + statements.append(batch_statement) + statements.append("") + + return statements + + def _create_vertex_batch_insert( + self, vertices: List[Dict[str, Any]], collection: str + ) -> str: + """ + Create a batch INSERT statement for vertices. + + Args: + vertices: Batch of vertex dictionaries + collection: Vertex collection name + + Returns: + AQL INSERT statement + """ + if not vertices: + return "" + + # Build document list + documents = [] + for idx, vertex in enumerate(vertices): + doc = self._convert_vertex_to_document(vertex, idx) + documents.append(doc) + + # Format as AQL + docs_json = json.dumps(documents, indent=2, ensure_ascii=False) + statement = f"FOR doc IN {docs_json}\n INSERT doc INTO {collection}" + + return statement + + def _convert_vertex_to_document( + self, vertex: Dict[str, Any], idx: int + ) -> Dict[str, Any]: + """ + Convert a vertex/entity to an ArangoDB document. + + Additional properties from the input entity are preserved as-is in the + document root (not flattened or merged). Nested dictionaries and lists + are preserved as JSON-serializable structures. The 'properties' field, + if present, is also preserved as-is rather than being flattened into + the document root. + + Args: + vertex: Vertex/entity dictionary + idx: Index for generating fallback IDs + + Returns: + ArangoDB document dictionary + """ + document = {} + + # Set _key from id or generate one + vertex_id = vertex.get("id") or vertex.get("entity_id") or f"vertex_{idx}" + document["_key"] = self._sanitize_key(str(vertex_id)) + + # Add original ID if different from _key + if str(vertex_id) != document["_key"]: + document["original_id"] = str(vertex_id) + + # Add type/label information + vertex_type = vertex.get("type") or vertex.get("entity_type", "Entity") + document["type"] = vertex_type + + # Add name/label + label = ( + vertex.get("label") + or vertex.get("name") + or vertex.get("text") + or document["_key"] + ) + document["name"] = label + + # Add all other properties + for key, value in vertex.items(): + if key not in [ + "_key", + "_id", + "_rev", + "id", + "entity_id", + "type", + "entity_type", + "label", + "name", + "text", + ]: + # Handle nested dictionaries and lists + if isinstance(value, (dict, list)): + document[key] = value + elif value is not None: + document[key] = value + + return document + + def _generate_edge_inserts( + self, edges: List[Dict[str, Any]], collection: str, vertex_collection: str + ) -> List[str]: + """ + Generate AQL INSERT statements for edges. + + Args: + edges: List of edge/relationship dictionaries + collection: Edge collection name + vertex_collection: Vertex collection name for _from/_to references + + Returns: + List of AQL INSERT statements + """ + statements = [] + + # Add header comment + statements.append( + f"// Attempting to insert {len(edges)} edges into {collection}" + ) + statements.append("") + + # Process edges in batches + for i in range(0, len(edges), self.batch_size): + batch = edges[i : i + self.batch_size] + batch_statement = self._create_edge_batch_insert( + batch, collection, vertex_collection + ) + if batch_statement: # Only add non-empty statements + statements.append(batch_statement) + statements.append("") + + return statements + + def _create_edge_batch_insert( + self, edges: List[Dict[str, Any]], collection: str, vertex_collection: str + ) -> str: + """ + Create a batch INSERT statement for edges. + + Args: + edges: Batch of edge dictionaries + collection: Edge collection name + vertex_collection: Vertex collection name for _from/_to references + + Returns: + AQL INSERT statement + """ + # Build document list + documents = [] + for idx, edge in enumerate(edges): + doc = self._convert_edge_to_document(edge, idx, vertex_collection) + if doc: # Only add valid edges + documents.append(doc) + + if not documents: + return "" + + # Format as AQL + docs_json = json.dumps(documents, indent=2, ensure_ascii=False) + statement = f"FOR doc IN {docs_json}\n INSERT doc INTO {collection}" + + return statement + + def _convert_edge_to_document( + self, edge: Dict[str, Any], idx: int, vertex_collection: str + ) -> Optional[Dict[str, Any]]: + """ + Convert an edge/relationship to an ArangoDB edge document. + + Args: + edge: Edge/relationship dictionary + idx: Index for generating fallback IDs + vertex_collection: Vertex collection name for _from/_to references + + Returns: + ArangoDB edge document dictionary, or None if source/target missing + """ + # Extract source and target + source_id = edge.get("source") or edge.get("source_id") + target_id = edge.get("target") or edge.get("target_id") + + if not source_id or not target_id: + self.logger.warning( + f"Skipping edge {idx}: missing source or target " + f"(source={source_id}, target={target_id})" + ) + return None + + document = {} + + # Set _key from id or generate one + edge_id = edge.get("id") or edge.get("relationship_id") or f"edge_{idx}" + document["_key"] = self._sanitize_key(str(edge_id)) + + # Set _from and _to (required for edges in ArangoDB) + document["_from"] = f"{vertex_collection}/{self._sanitize_key(str(source_id))}" + document["_to"] = f"{vertex_collection}/{self._sanitize_key(str(target_id))}" + + # Add original ID if different from _key + if str(edge_id) != document["_key"]: + document["original_id"] = str(edge_id) + + # Add relationship type + rel_type = edge.get("type") or edge.get("relationship_type", "RELATED_TO") + document["type"] = rel_type + + # Add all other properties + for key, value in edge.items(): + if key not in [ + "_key", + "_id", + "_rev", + "_from", + "_to", + "id", + "relationship_id", + "source", + "source_id", + "target", + "target_id", + "type", + "relationship_type", + ]: + # Handle nested dictionaries and lists + if isinstance(value, (dict, list)): + document[key] = value + elif value is not None: + document[key] = value + + return document + + def _validate_collection_name(self, name: str, param_name: str) -> None: + """ + Validate an ArangoDB collection name. + + ArangoDB collection names must: + - Start with a letter or underscore + - Contain only alphanumeric characters, hyphens, and underscores + - Not exceed 256 characters + + Args: + name: Collection name to validate + param_name: Parameter name for error messages + + Raises: + ValueError: If the collection name is invalid + """ + if not name: + raise ValueError(f"{param_name} cannot be empty") + + if len(name) > 256: + raise ValueError( + f"{param_name} '{name}' exceeds maximum length of 256 characters" + ) + + # Check first character + if not (name[0].isalpha() or name[0] == "_"): + raise ValueError( + f"{param_name} '{name}' must start with a letter or underscore" + ) + + # Check remaining characters + for char in name: + if not (char.isalnum() or char in ("-", "_")): + raise ValueError( + f"{param_name} '{name}' contains invalid character " + f"'{char}'. Only alphanumeric characters, hyphens, and " + "underscores are allowed." + ) + + def _sanitize_key(self, key: str) -> str: + """ + Sanitize a key for use as ArangoDB _key. + + ArangoDB _key must contain only alphanumeric characters, hyphens, + and underscores. It cannot start with an underscore (unless it's + a system collection). + + Args: + key: Original key string + + Returns: + Sanitized key string + """ + # Replace invalid characters with underscores + sanitized = "" + for char in key: + if char.isalnum() or char in ("-", "_"): + sanitized += char + else: + sanitized += "_" + + # Ensure key doesn't start with underscore + if sanitized.startswith("_"): + sanitized = "k" + sanitized + + # Ensure key is not empty + if not sanitized: + sanitized = "key" + + return sanitized + + def export_knowledge_graph( + self, knowledge_graph: Dict[str, Any], file_path: Union[str, Path], **options + ) -> None: + """ + Export knowledge graph to ArangoDB AQL format. + + Convenience method that calls export(). + + Args: + knowledge_graph: Knowledge graph dictionary + file_path: Output AQL file path + **options: Additional export options + + Example: + >>> kg = { + ... "entities": [...], + ... "relationships": [...] + ... } + >>> exporter.export_knowledge_graph(kg, "output.aql") + """ + self.export(knowledge_graph, file_path, **options) + + def export_entities( + self, entities: List[Dict[str, Any]], file_path: Union[str, Path], **options + ) -> None: + """ + Export entities to ArangoDB AQL format (vertices only). + + Args: + entities: List of entity dictionaries + file_path: Output AQL file path + **options: Additional export options + + Example: + >>> entities = [ + ... {"id": "e1", "type": "Person", "name": "Alice"}, + ... {"id": "e2", "type": "Organization", "name": "Acme Corp"} + ... ] + >>> exporter.export_entities(entities, "entities.aql") + """ + knowledge_graph = {"entities": entities, "relationships": []} + self.export(knowledge_graph, file_path, **options) + + def export_relationships( + self, + relationships: List[Dict[str, Any]], + file_path: Union[str, Path], + **options, + ) -> None: + """ + Export relationships to ArangoDB AQL format (edges only). + + Args: + relationships: List of relationship dictionaries + file_path: Output AQL file path + **options: Additional export options + + Example: + >>> relationships = [ + ... {"id": "r1", "source": "e1", "target": "e2", "type": "WORKS_FOR"} + ... ] + >>> exporter.export_relationships(relationships, "relationships.aql") + """ + knowledge_graph = {"entities": [], "relationships": relationships} + self.export(knowledge_graph, file_path, **options) diff --git a/semantica/export/parquet_exporter.py b/semantica/export/parquet_exporter.py index b2d6911e..8934f2e0 100644 --- a/semantica/export/parquet_exporter.py +++ b/semantica/export/parquet_exporter.py @@ -1,700 +1,700 @@ -""" -Apache Parquet Exporter Module - -This module provides comprehensive Apache Parquet export capabilities for the -Semantica framework, enabling efficient columnar data export for entities, -relationships, and knowledge graphs optimized for analytics and data warehousing. - -Key Features: - - Parquet file export (.parquet) - - Explicit schema definition (no inference) - - Entity and relationship export with metadata - - Knowledge graph export to multiple Parquet files - - Compatible with pandas, Spark, Snowflake, BigQuery, and Databricks - - Configurable compression (snappy, gzip, brotli, zstd, lz4) - - Batch export processing - - Structured metadata handling - -Example Usage: - >>> from semantica.export import ParquetExporter - >>> exporter = ParquetExporter(compression="snappy") - >>> exporter.export_entities(entities, "entities.parquet") - >>> exporter.export_knowledge_graph(kg, "kg_base") - -Author: Semantica Contributors -License: MIT -""" - -import json -from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union - -if TYPE_CHECKING: - import pyarrow as pa # noqa: F401 - -try: - import pyarrow as pa # noqa: F811 - import pyarrow.parquet as pq - - PARQUET_AVAILABLE = True - - # Explicit Parquet Schemas (no inference) - ENTITY_SCHEMA = pa.schema( - [ - pa.field("id", pa.string(), nullable=False), - pa.field("text", pa.string(), nullable=True), - pa.field("type", pa.string(), nullable=True), - pa.field("confidence", pa.float64(), nullable=True), - pa.field("start", pa.int64(), nullable=True), - pa.field("end", pa.int64(), nullable=True), - pa.field( - "metadata", - pa.struct( - [ - pa.field("keys", pa.list_(pa.string())), - pa.field("values", pa.list_(pa.string())), - ] - ), - nullable=True, - ), - ] - ) - - RELATIONSHIP_SCHEMA = pa.schema( - [ - pa.field("id", pa.string(), nullable=False), - pa.field("source_id", pa.string(), nullable=False), - pa.field("target_id", pa.string(), nullable=False), - pa.field("type", pa.string(), nullable=True), - pa.field("confidence", pa.float64(), nullable=True), - pa.field( - "metadata", - pa.struct( - [ - pa.field("keys", pa.list_(pa.string())), - pa.field("values", pa.list_(pa.string())), - ] - ), - nullable=True, - ), - ] - ) -except ImportError: - PARQUET_AVAILABLE = False - ENTITY_SCHEMA = None - RELATIONSHIP_SCHEMA = None - -from ..utils.exceptions import ProcessingError, ValidationError -from ..utils.helpers import ensure_directory -from ..utils.logging import get_logger -from ..utils.progress_tracker import get_progress_tracker - - -class ParquetExporter: - """ - Apache Parquet exporter for knowledge graphs and structured data. - - This class provides comprehensive Parquet export functionality for entities, - relationships, and knowledge graphs. Uses explicit schemas for type safety - and compatibility with analytics platforms like pandas, Spark, Snowflake, - BigQuery, and Databricks. - - Features: - - Entity and relationship export - - Knowledge graph export to multiple Parquet files - - Explicit schema definition (no inference) - - Metadata serialization as Parquet struct fields - - Configurable compression (snappy, gzip, brotli, zstd, lz4) - - Compatible with major analytics platforms - - Progress tracking and error handling - - Example Usage: - >>> exporter = ParquetExporter(compression="snappy") - >>> exporter.export_entities(entities, "entities.parquet") - >>> exporter.export_knowledge_graph(kg, "output_base") - """ - - def __init__( - self, - compression: str = "snappy", - config: Optional[Dict[str, Any]] = None, - **kwargs, - ): - """ - Initialize Parquet exporter. - - Sets up the exporter with specified Parquet formatting options. - - Args: - compression: Compression codec (default: "snappy") - - "snappy": Snappy compression (fast, good compression) - - "gzip": GZIP compression (slower, better compression) - - "brotli": Brotli compression (slow, best compression) - - "zstd": Zstandard compression (balanced) - - "lz4": LZ4 compression (very fast, moderate compression) - - "none" or None: No compression - config: Optional configuration dictionary (merged with kwargs) - **kwargs: Additional configuration options - - Raises: - ImportError: If pyarrow is not installed - """ - if not PARQUET_AVAILABLE: - raise ImportError( - "pyarrow is not installed. Please install it with: " - "pip install pyarrow" - ) - - self.logger = get_logger("parquet_exporter") - self.config = config or {} - self.config.update(kwargs) - - # Parquet configuration - self.compression = compression if compression != "none" else None - - # Initialize progress tracker - self.progress_tracker = get_progress_tracker() - # Ensure progress tracker is enabled - if not self.progress_tracker.enabled: - self.progress_tracker.enabled = True - - self.logger.debug(f"Parquet exporter initialized: compression={compression}") - - def export( - self, - data: Union[List[Dict[str, Any]], Dict[str, Any]], - file_path: Union[str, Path], - schema: Optional["pa.Schema"] = None, - **options, - ) -> None: - """ - Export data to Parquet file(s). - - This method handles both single Parquet file export (from list) and multiple - Parquet file export (from dictionary with multiple keys). - - Args: - data: Data to export: - - List of dicts: Exports to single Parquet file - - Dict with list values: Exports each key as separate Parquet file - file_path: Output file path (base path for dict exports) - schema: Parquet schema to use (default: auto-select based on data) - **options: Additional options - - Raises: - ValidationError: If data type is unsupported - - Example: - >>> # Single Parquet file - >>> exporter.export([{"id": "1", "name": "A"}], "data.parquet") - >>> # Multiple Parquet files - >>> exporter.export( - ... {"entities": [...], "relationships": [...]}, - ... "output_base" - ... ) - """ - # Track Parquet export - tracking_id = self.progress_tracker.start_tracking( - file=str(file_path), - module="export", - submodule="ParquetExporter", - message=f"Exporting data to Parquet: {file_path}", - ) - - try: - file_path = Path(file_path) - ensure_directory(file_path.parent) - - self.logger.debug(f"Exporting data to Parquet: {file_path}") - - # Handle different data structures - if isinstance(data, dict): - # Export each key as separate Parquet file - exported_files = [] - self.progress_tracker.update_tracking( - tracking_id, message=f"Exporting {len(data)} data groups..." - ) - for key, value in data.items(): - if isinstance(value, list): - output_path = ( - file_path.parent / f"{file_path.stem}_{key}.parquet" - ) - - # Use dedicated export methods for entities and relationships - # to ensure proper normalization - if key == "entities" and schema is None: - self.export_entities(value, output_path, **options) - elif key == "relationships" and schema is None: - self.export_relationships(value, output_path, **options) - else: - # For other keys, write directly with provided schema - self._write_parquet( - value, output_path, schema=schema, **options - ) - - exported_files.append(output_path) - else: - self.logger.warning( - f"Skipping key '{key}': value is not a list " - f"(type: {type(value)})" - ) - - self.logger.info( - f"Exported {len(exported_files)} Parquet file(s) from dictionary: " - f"{', '.join(str(f) for f in exported_files)}" - ) - self.progress_tracker.stop_tracking( - tracking_id, - status="completed", - message=f"Exported {len(exported_files)} Parquet files", - ) - elif isinstance(data, list): - # Single Parquet file - auto-detect if entities or relationships - self.progress_tracker.update_tracking( - tracking_id, message=f"Exporting {len(data)} records..." - ) - - # If no schema provided, try to auto-detect from data structure - if schema is None: - if not data: - raise ValidationError( - "Cannot export empty list without explicit schema. " - "Provide a schema or use " - "export_entities/export_relationships." - ) - sample = data[0] - has_source = any( - k in sample for k in ["source_id", "source", "from_id", "from"] - ) - has_target = any( - k in sample for k in ["target_id", "target", "to_id", "to"] - ) - - if has_source and has_target: - # Use dedicated method for relationship normalization - self.export_relationships(data, file_path, **options) - else: - # Use dedicated method for entity normalization - self.export_entities(data, file_path, **options) - else: - # Schema provided - write directly - self._write_parquet(data, file_path, schema=schema, **options) - - self.logger.info(f"Exported Parquet to: {file_path}") - self.progress_tracker.stop_tracking( - tracking_id, - status="completed", - message=f"Exported Parquet to: {file_path}", - ) - else: - raise ValidationError( - f"Unsupported data type: {type(data)}. " - "Expected list of dicts or dict with list values." - ) - - except Exception as e: - self.progress_tracker.stop_tracking( - tracking_id, status="failed", message=str(e) - ) - raise - - def export_entities( - self, entities: List[Dict[str, Any]], file_path: Union[str, Path], **options - ) -> None: - """ - Export entities to Parquet file. - - This method normalizes entity data to a consistent format and exports - to Parquet using the explicit ENTITY_SCHEMA. Handles various entity field - name variations and serializes metadata as Parquet structs. - - Normalized Fields: - - id: Entity identifier (required) - - text: Entity text/label/name - - type: Entity type - - confidence: Confidence score - - start: Start offset/position - - end: End offset/position - - metadata: Metadata as struct (keys and values lists) - - Args: - entities: List of entity dictionaries with various field names - file_path: Output Parquet file path - **options: Additional options passed to _write_parquet() - - Raises: - ValidationError: If entities list is empty - - Example: - >>> entities = [ - ... {"id": "e1", "text": "Entity 1", "type": "PERSON"}, - ... {"id": "e2", "label": "Entity 2", "entity_type": "ORG"} - ... ] - >>> exporter.export_entities(entities, "entities.parquet") - """ - if not entities: - raise ValidationError("No entities to export. Entities list is empty.") - - self.logger.debug(f"Exporting {len(entities)} entity(ies) to Parquet") - - # Normalize entity data to consistent format - normalized_entities = [] - for i, entity in enumerate(entities): - if not isinstance(entity, dict): - self.logger.warning(f"Entity {i} is not a dictionary, skipping") - continue - - # Extract and normalize fields - entity_id = entity.get("id") or entity.get("entity_id") - if not entity_id: - self.logger.warning(f"Entity {i} missing ID, skipping") - continue - - # Normalize confidence to float, with validation - raw_confidence = entity.get("confidence") - confidence_value = None - if raw_confidence is not None: - try: - confidence_value = float(raw_confidence) - except (TypeError, ValueError): - self.logger.warning( - f"Entity {i} has non-numeric confidence {raw_confidence!r}; " - "setting to None" - ) - confidence_value = None - - # Normalize start/end to int, with validation - start_value = entity.get("start") - if start_value is None: - start_value = entity.get("start_offset") - if start_value is not None: - try: - start_value = int(start_value) - except (TypeError, ValueError): - self.logger.warning( - f"Entity {i} has non-integer start {start_value!r}; " - "setting to None" - ) - start_value = None - - end_value = entity.get("end") - if end_value is None: - end_value = entity.get("end_offset") - if end_value is not None: - try: - end_value = int(end_value) - except (TypeError, ValueError): - self.logger.warning( - f"Entity {i} has non-integer end {end_value!r}; " - "setting to None" - ) - end_value = None - - normalized = { - "id": str(entity_id), - "text": ( - entity.get("text") or entity.get("label") or entity.get("name") - ), - "type": entity.get("type") or entity.get("entity_type"), - "confidence": confidence_value, - "start": start_value, - "end": end_value, - } - - # Convert metadata to struct format (keys and values lists) - if "metadata" in entity and isinstance(entity["metadata"], dict): - metadata_dict = entity["metadata"] - normalized["metadata"] = { - "keys": list(metadata_dict.keys()), - "values": [ - json.dumps(v) if not isinstance(v, str) else v - for v in metadata_dict.values() - ], - } - else: - normalized["metadata"] = None - - normalized_entities.append(normalized) - - self.logger.debug( - f"Normalized {len(normalized_entities)} entity(ies) for Parquet export" - ) - - if not normalized_entities: - raise ValidationError( - "No valid entities to export after normalization. " - "All entities were skipped due to missing IDs or invalid format." - ) - - self._write_parquet( - normalized_entities, file_path, schema=ENTITY_SCHEMA, **options - ) - - def export_relationships( - self, - relationships: List[Dict[str, Any]], - file_path: Union[str, Path], - **options, - ) -> None: - """ - Export relationships to Parquet file. - - This method normalizes relationship data to a consistent format and exports - to Parquet using the explicit RELATIONSHIP_SCHEMA. Handles various relationship - field name variations and serializes metadata as Parquet structs. - - Normalized Fields: - - id: Relationship identifier (generated if missing) - - source_id: Source entity ID (required) - - target_id: Target entity ID (required) - - type: Relationship type - - confidence: Confidence score - - metadata: Metadata as struct (keys and values lists) - - Args: - relationships: List of relationship dictionaries with various field names - file_path: Output Parquet file path - **options: Additional options passed to _write_parquet() - - Raises: - ValidationError: If relationships list is empty - - Example: - >>> relationships = [ - ... {"id": "r1", "source": "e1", "target": "e2", "type": "KNOWS"}, - ... {"source_id": "e2", "target_id": "e3", "relationship_type": "LIKES"} - ... ] - >>> exporter.export_relationships(relationships, "relationships.parquet") - """ - if not relationships: - raise ValidationError( - "No relationships to export. Relationships list is empty." - ) - - self.logger.debug(f"Exporting {len(relationships)} relationship(s) to Parquet") - - # Normalize relationship data to consistent format - normalized_rels = [] - for i, rel in enumerate(relationships): - if not isinstance(rel, dict): - self.logger.warning(f"Relationship {i} is not a dictionary, skipping") - continue - - # Extract source and target IDs - source_id = ( - rel.get("source_id") - or rel.get("source") - or rel.get("from_id") - or rel.get("from") - ) - target_id = ( - rel.get("target_id") - or rel.get("target") - or rel.get("to_id") - or rel.get("to") - ) - - if not source_id or not target_id: - self.logger.warning( - f"Relationship {i} missing source or target ID, skipping" - ) - continue - - # Generate ID if missing - rel_id = rel.get("id") or rel.get("relationship_id") or f"rel_{i}" - - # Normalize confidence to float, with validation - raw_confidence = rel.get("confidence") - confidence_value = None - if raw_confidence is not None: - try: - confidence_value = float(raw_confidence) - except (TypeError, ValueError): - self.logger.warning( - f"Relationship {i} has non-numeric confidence " - f"{raw_confidence!r}; setting to None" - ) - confidence_value = None - - normalized = { - "id": str(rel_id), - "source_id": str(source_id), - "target_id": str(target_id), - "type": ( - rel.get("type") - or rel.get("relationship_type") - or rel.get("relation_type") - ), - "confidence": confidence_value, - } - - # Convert metadata to struct format (keys and values lists) - if "metadata" in rel and isinstance(rel["metadata"], dict): - metadata_dict = rel["metadata"] - normalized["metadata"] = { - "keys": list(metadata_dict.keys()), - "values": [ - json.dumps(v) if not isinstance(v, str) else v - for v in metadata_dict.values() - ], - } - else: - normalized["metadata"] = None - - normalized_rels.append(normalized) - - self.logger.debug( - f"Normalized {len(normalized_rels)} relationship(s) for Parquet export" - ) - - if not normalized_rels: - raise ValidationError( - "No valid relationships to export after normalization. " - "Ensure each relationship is a dictionary and includes valid " - "'source'/'source_id' and 'target'/'target_id' fields." - ) - - self._write_parquet( - normalized_rels, file_path, schema=RELATIONSHIP_SCHEMA, **options - ) - - def export_knowledge_graph( - self, kg: Dict[str, Any], base_path: Union[str, Path], **options - ) -> None: - """ - Export knowledge graph to multiple Parquet files. - - This method exports a knowledge graph to separate Parquet files for - entities and relationships. Files are named using the base_path with - suffixes: _entities.parquet and _relationships.parquet. - - Args: - kg: Knowledge graph dictionary with 'entities' and 'relationships' keys - base_path: Base path for output files (without extension) - **options: Additional options passed to export methods - - Raises: - ValidationError: If knowledge graph is missing required keys - - Example: - >>> kg = { - ... "entities": [...], - ... "relationships": [...] - ... } - >>> exporter.export_knowledge_graph(kg, "output/kg_base") - # Creates: output/kg_base_entities.parquet and - # output/kg_base_relationships.parquet - """ - if not isinstance(kg, dict): - raise ValidationError( - f"Knowledge graph must be a dictionary, got {type(kg)}" - ) - - if "entities" not in kg and "relationships" not in kg: - raise ValidationError( - "Knowledge graph must contain 'entities' or 'relationships' key" - ) - - base_path = Path(base_path) - ensure_directory(base_path.parent) - - self.logger.debug(f"Exporting knowledge graph to Parquet: {base_path}") - - # Track KG export - tracking_id = self.progress_tracker.start_tracking( - file=str(base_path), - module="export", - submodule="ParquetExporter", - message=f"Exporting knowledge graph to Parquet: {base_path}", - ) - - try: - exported_files = [] - - # Export entities - if "entities" in kg and kg["entities"]: - entities_path = base_path.parent / f"{base_path.stem}_entities.parquet" - self.export_entities(kg["entities"], entities_path, **options) - exported_files.append(entities_path) - self.logger.info(f"Exported entities to: {entities_path}") - - # Export relationships - if "relationships" in kg and kg["relationships"]: - rels_path = base_path.parent / f"{base_path.stem}_relationships.parquet" - self.export_relationships(kg["relationships"], rels_path, **options) - exported_files.append(rels_path) - self.logger.info(f"Exported relationships to: {rels_path}") - - self.logger.info( - f"Exported knowledge graph to {len(exported_files)} Parquet file(s)" - ) - self.progress_tracker.stop_tracking( - tracking_id, - status="completed", - message=f"Exported {len(exported_files)} Parquet files", - ) - - except Exception as e: - self.progress_tracker.stop_tracking( - tracking_id, status="failed", message=str(e) - ) - raise - - def _write_parquet( - self, - data: List[Dict[str, Any]], - file_path: Union[str, Path], - schema: Optional["pa.Schema"] = None, - **options, - ) -> None: - """ - Write data to Parquet file. - - Internal method that handles the actual Parquet file writing using pyarrow. - - Args: - data: List of dictionaries to write - file_path: Output Parquet file path - schema: Parquet schema to use (required for type safety) - **options: Additional parquet write options - - Raises: - ProcessingError: If Parquet writing fails - ValidationError: If schema is not provided - """ - if not schema: - raise ValidationError("Schema is required for Parquet export") - - # Ensure file_path is a Path object - file_path = Path(file_path) - - if not data: - self.logger.warning(f"No data to write to {file_path}") - # Write empty Parquet file with schema - empty_table = pa.table({field.name: [] for field in schema}, schema=schema) - pq.write_table( - empty_table, str(file_path), compression=self.compression, **options - ) - return - - try: - # Create PyArrow table from data using explicit schema - table = pa.Table.from_pylist(data, schema=schema) - - # Write to Parquet file - pq.write_table( - table, str(file_path), compression=self.compression, **options - ) - - file_size = file_path.stat().st_size - self.logger.debug( - f"Wrote {len(data)} row(s) to {file_path} ({file_size} bytes)" - ) - - except pa.ArrowInvalid as e: - raise ProcessingError( - f"Failed to create Parquet table: {e}. " - "Check that data matches schema." - ) - except Exception as e: - raise ProcessingError(f"Failed to write Parquet file: {e}") +""" +Apache Parquet Exporter Module + +This module provides comprehensive Apache Parquet export capabilities for the +Semantica framework, enabling efficient columnar data export for entities, +relationships, and knowledge graphs optimized for analytics and data warehousing. + +Key Features: + - Parquet file export (.parquet) + - Explicit schema definition (no inference) + - Entity and relationship export with metadata + - Knowledge graph export to multiple Parquet files + - Compatible with pandas, Spark, Snowflake, BigQuery, and Databricks + - Configurable compression (snappy, gzip, brotli, zstd, lz4) + - Batch export processing + - Structured metadata handling + +Example Usage: + >>> from semantica.export import ParquetExporter + >>> exporter = ParquetExporter(compression="snappy") + >>> exporter.export_entities(entities, "entities.parquet") + >>> exporter.export_knowledge_graph(kg, "kg_base") + +Author: Semantica Contributors +License: MIT +""" + +import json +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +if TYPE_CHECKING: + import pyarrow as pa # noqa: F401 + +try: + import pyarrow as pa # noqa: F811 + import pyarrow.parquet as pq + + PARQUET_AVAILABLE = True + + # Explicit Parquet Schemas (no inference) + ENTITY_SCHEMA = pa.schema( + [ + pa.field("id", pa.string(), nullable=False), + pa.field("text", pa.string(), nullable=True), + pa.field("type", pa.string(), nullable=True), + pa.field("confidence", pa.float64(), nullable=True), + pa.field("start", pa.int64(), nullable=True), + pa.field("end", pa.int64(), nullable=True), + pa.field( + "metadata", + pa.struct( + [ + pa.field("keys", pa.list_(pa.string())), + pa.field("values", pa.list_(pa.string())), + ] + ), + nullable=True, + ), + ] + ) + + RELATIONSHIP_SCHEMA = pa.schema( + [ + pa.field("id", pa.string(), nullable=False), + pa.field("source_id", pa.string(), nullable=False), + pa.field("target_id", pa.string(), nullable=False), + pa.field("type", pa.string(), nullable=True), + pa.field("confidence", pa.float64(), nullable=True), + pa.field( + "metadata", + pa.struct( + [ + pa.field("keys", pa.list_(pa.string())), + pa.field("values", pa.list_(pa.string())), + ] + ), + nullable=True, + ), + ] + ) +except ImportError: + PARQUET_AVAILABLE = False + ENTITY_SCHEMA = None + RELATIONSHIP_SCHEMA = None + +from ..utils.exceptions import ProcessingError, ValidationError +from ..utils.helpers import ensure_directory +from ..utils.logging import get_logger +from ..utils.progress_tracker import get_progress_tracker + + +class ParquetExporter: + """ + Apache Parquet exporter for knowledge graphs and structured data. + + This class provides comprehensive Parquet export functionality for entities, + relationships, and knowledge graphs. Uses explicit schemas for type safety + and compatibility with analytics platforms like pandas, Spark, Snowflake, + BigQuery, and Databricks. + + Features: + - Entity and relationship export + - Knowledge graph export to multiple Parquet files + - Explicit schema definition (no inference) + - Metadata serialization as Parquet struct fields + - Configurable compression (snappy, gzip, brotli, zstd, lz4) + - Compatible with major analytics platforms + - Progress tracking and error handling + + Example Usage: + >>> exporter = ParquetExporter(compression="snappy") + >>> exporter.export_entities(entities, "entities.parquet") + >>> exporter.export_knowledge_graph(kg, "output_base") + """ + + def __init__( + self, + compression: str = "snappy", + config: Optional[Dict[str, Any]] = None, + **kwargs, + ): + """ + Initialize Parquet exporter. + + Sets up the exporter with specified Parquet formatting options. + + Args: + compression: Compression codec (default: "snappy") + - "snappy": Snappy compression (fast, good compression) + - "gzip": GZIP compression (slower, better compression) + - "brotli": Brotli compression (slow, best compression) + - "zstd": Zstandard compression (balanced) + - "lz4": LZ4 compression (very fast, moderate compression) + - "none" or None: No compression + config: Optional configuration dictionary (merged with kwargs) + **kwargs: Additional configuration options + + Raises: + ImportError: If pyarrow is not installed + """ + if not PARQUET_AVAILABLE: + raise ImportError( + "pyarrow is not installed. Please install it with: " + "pip install pyarrow" + ) + + self.logger = get_logger("parquet_exporter") + self.config = config or {} + self.config.update(kwargs) + + # Parquet configuration + self.compression = compression if compression != "none" else None + + # Initialize progress tracker + self.progress_tracker = get_progress_tracker() + # Ensure progress tracker is enabled + if not self.progress_tracker.enabled: + self.progress_tracker.enabled = True + + self.logger.debug(f"Parquet exporter initialized: compression={compression}") + + def export( + self, + data: Union[List[Dict[str, Any]], Dict[str, Any]], + file_path: Union[str, Path], + schema: Optional["pa.Schema"] = None, + **options, + ) -> None: + """ + Export data to Parquet file(s). + + This method handles both single Parquet file export (from list) and multiple + Parquet file export (from dictionary with multiple keys). + + Args: + data: Data to export: + - List of dicts: Exports to single Parquet file + - Dict with list values: Exports each key as separate Parquet file + file_path: Output file path (base path for dict exports) + schema: Parquet schema to use (default: auto-select based on data) + **options: Additional options + + Raises: + ValidationError: If data type is unsupported + + Example: + >>> # Single Parquet file + >>> exporter.export([{"id": "1", "name": "A"}], "data.parquet") + >>> # Multiple Parquet files + >>> exporter.export( + ... {"entities": [...], "relationships": [...]}, + ... "output_base" + ... ) + """ + # Track Parquet export + tracking_id = self.progress_tracker.start_tracking( + file=str(file_path), + module="export", + submodule="ParquetExporter", + message=f"Exporting data to Parquet: {file_path}", + ) + + try: + file_path = Path(file_path) + ensure_directory(file_path.parent) + + self.logger.debug(f"Exporting data to Parquet: {file_path}") + + # Handle different data structures + if isinstance(data, dict): + # Export each key as separate Parquet file + exported_files = [] + self.progress_tracker.update_tracking( + tracking_id, message=f"Exporting {len(data)} data groups..." + ) + for key, value in data.items(): + if isinstance(value, list): + output_path = ( + file_path.parent / f"{file_path.stem}_{key}.parquet" + ) + + # Use dedicated export methods for entities and relationships + # to ensure proper normalization + if key == "entities" and schema is None: + self.export_entities(value, output_path, **options) + elif key == "relationships" and schema is None: + self.export_relationships(value, output_path, **options) + else: + # For other keys, write directly with provided schema + self._write_parquet( + value, output_path, schema=schema, **options + ) + + exported_files.append(output_path) + else: + self.logger.warning( + f"Skipping key '{key}': value is not a list " + f"(type: {type(value)})" + ) + + self.logger.info( + f"Exported {len(exported_files)} Parquet file(s) from dictionary: " + f"{', '.join(str(f) for f in exported_files)}" + ) + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Exported {len(exported_files)} Parquet files", + ) + elif isinstance(data, list): + # Single Parquet file - auto-detect if entities or relationships + self.progress_tracker.update_tracking( + tracking_id, message=f"Exporting {len(data)} records..." + ) + + # If no schema provided, try to auto-detect from data structure + if schema is None: + if not data: + raise ValidationError( + "Cannot export empty list without explicit schema. " + "Provide a schema or use " + "export_entities/export_relationships." + ) + sample = data[0] + has_source = any( + k in sample for k in ["source_id", "source", "from_id", "from"] + ) + has_target = any( + k in sample for k in ["target_id", "target", "to_id", "to"] + ) + + if has_source and has_target: + # Use dedicated method for relationship normalization + self.export_relationships(data, file_path, **options) + else: + # Use dedicated method for entity normalization + self.export_entities(data, file_path, **options) + else: + # Schema provided - write directly + self._write_parquet(data, file_path, schema=schema, **options) + + self.logger.info(f"Exported Parquet to: {file_path}") + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Exported Parquet to: {file_path}", + ) + else: + raise ValidationError( + f"Unsupported data type: {type(data)}. " + "Expected list of dicts or dict with list values." + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + raise + + def export_entities( + self, entities: List[Dict[str, Any]], file_path: Union[str, Path], **options + ) -> None: + """ + Export entities to Parquet file. + + This method normalizes entity data to a consistent format and exports + to Parquet using the explicit ENTITY_SCHEMA. Handles various entity field + name variations and serializes metadata as Parquet structs. + + Normalized Fields: + - id: Entity identifier (required) + - text: Entity text/label/name + - type: Entity type + - confidence: Confidence score + - start: Start offset/position + - end: End offset/position + - metadata: Metadata as struct (keys and values lists) + + Args: + entities: List of entity dictionaries with various field names + file_path: Output Parquet file path + **options: Additional options passed to _write_parquet() + + Raises: + ValidationError: If entities list is empty + + Example: + >>> entities = [ + ... {"id": "e1", "text": "Entity 1", "type": "PERSON"}, + ... {"id": "e2", "label": "Entity 2", "entity_type": "ORG"} + ... ] + >>> exporter.export_entities(entities, "entities.parquet") + """ + if not entities: + raise ValidationError("No entities to export. Entities list is empty.") + + self.logger.debug(f"Exporting {len(entities)} entity(ies) to Parquet") + + # Normalize entity data to consistent format + normalized_entities = [] + for i, entity in enumerate(entities): + if not isinstance(entity, dict): + self.logger.warning(f"Entity {i} is not a dictionary, skipping") + continue + + # Extract and normalize fields + entity_id = entity.get("id") or entity.get("entity_id") + if not entity_id: + self.logger.warning(f"Entity {i} missing ID, skipping") + continue + + # Normalize confidence to float, with validation + raw_confidence = entity.get("confidence") + confidence_value = None + if raw_confidence is not None: + try: + confidence_value = float(raw_confidence) + except (TypeError, ValueError): + self.logger.warning( + f"Entity {i} has non-numeric confidence {raw_confidence!r}; " + "setting to None" + ) + confidence_value = None + + # Normalize start/end to int, with validation + start_value = entity.get("start") + if start_value is None: + start_value = entity.get("start_offset") + if start_value is not None: + try: + start_value = int(start_value) + except (TypeError, ValueError): + self.logger.warning( + f"Entity {i} has non-integer start {start_value!r}; " + "setting to None" + ) + start_value = None + + end_value = entity.get("end") + if end_value is None: + end_value = entity.get("end_offset") + if end_value is not None: + try: + end_value = int(end_value) + except (TypeError, ValueError): + self.logger.warning( + f"Entity {i} has non-integer end {end_value!r}; " + "setting to None" + ) + end_value = None + + normalized = { + "id": str(entity_id), + "text": ( + entity.get("text") or entity.get("label") or entity.get("name") + ), + "type": entity.get("type") or entity.get("entity_type"), + "confidence": confidence_value, + "start": start_value, + "end": end_value, + } + + # Convert metadata to struct format (keys and values lists) + if "metadata" in entity and isinstance(entity["metadata"], dict): + metadata_dict = entity["metadata"] + normalized["metadata"] = { + "keys": list(metadata_dict.keys()), + "values": [ + json.dumps(v) if not isinstance(v, str) else v + for v in metadata_dict.values() + ], + } + else: + normalized["metadata"] = None + + normalized_entities.append(normalized) + + self.logger.debug( + f"Normalized {len(normalized_entities)} entity(ies) for Parquet export" + ) + + if not normalized_entities: + raise ValidationError( + "No valid entities to export after normalization. " + "All entities were skipped due to missing IDs or invalid format." + ) + + self._write_parquet( + normalized_entities, file_path, schema=ENTITY_SCHEMA, **options + ) + + def export_relationships( + self, + relationships: List[Dict[str, Any]], + file_path: Union[str, Path], + **options, + ) -> None: + """ + Export relationships to Parquet file. + + This method normalizes relationship data to a consistent format and exports + to Parquet using the explicit RELATIONSHIP_SCHEMA. Handles various relationship + field name variations and serializes metadata as Parquet structs. + + Normalized Fields: + - id: Relationship identifier (generated if missing) + - source_id: Source entity ID (required) + - target_id: Target entity ID (required) + - type: Relationship type + - confidence: Confidence score + - metadata: Metadata as struct (keys and values lists) + + Args: + relationships: List of relationship dictionaries with various field names + file_path: Output Parquet file path + **options: Additional options passed to _write_parquet() + + Raises: + ValidationError: If relationships list is empty + + Example: + >>> relationships = [ + ... {"id": "r1", "source": "e1", "target": "e2", "type": "KNOWS"}, + ... {"source_id": "e2", "target_id": "e3", "relationship_type": "LIKES"} + ... ] + >>> exporter.export_relationships(relationships, "relationships.parquet") + """ + if not relationships: + raise ValidationError( + "No relationships to export. Relationships list is empty." + ) + + self.logger.debug(f"Exporting {len(relationships)} relationship(s) to Parquet") + + # Normalize relationship data to consistent format + normalized_rels = [] + for i, rel in enumerate(relationships): + if not isinstance(rel, dict): + self.logger.warning(f"Relationship {i} is not a dictionary, skipping") + continue + + # Extract source and target IDs + source_id = ( + rel.get("source_id") + or rel.get("source") + or rel.get("from_id") + or rel.get("from") + ) + target_id = ( + rel.get("target_id") + or rel.get("target") + or rel.get("to_id") + or rel.get("to") + ) + + if not source_id or not target_id: + self.logger.warning( + f"Relationship {i} missing source or target ID, skipping" + ) + continue + + # Generate ID if missing + rel_id = rel.get("id") or rel.get("relationship_id") or f"rel_{i}" + + # Normalize confidence to float, with validation + raw_confidence = rel.get("confidence") + confidence_value = None + if raw_confidence is not None: + try: + confidence_value = float(raw_confidence) + except (TypeError, ValueError): + self.logger.warning( + f"Relationship {i} has non-numeric confidence " + f"{raw_confidence!r}; setting to None" + ) + confidence_value = None + + normalized = { + "id": str(rel_id), + "source_id": str(source_id), + "target_id": str(target_id), + "type": ( + rel.get("type") + or rel.get("relationship_type") + or rel.get("relation_type") + ), + "confidence": confidence_value, + } + + # Convert metadata to struct format (keys and values lists) + if "metadata" in rel and isinstance(rel["metadata"], dict): + metadata_dict = rel["metadata"] + normalized["metadata"] = { + "keys": list(metadata_dict.keys()), + "values": [ + json.dumps(v) if not isinstance(v, str) else v + for v in metadata_dict.values() + ], + } + else: + normalized["metadata"] = None + + normalized_rels.append(normalized) + + self.logger.debug( + f"Normalized {len(normalized_rels)} relationship(s) for Parquet export" + ) + + if not normalized_rels: + raise ValidationError( + "No valid relationships to export after normalization. " + "Ensure each relationship is a dictionary and includes valid " + "'source'/'source_id' and 'target'/'target_id' fields." + ) + + self._write_parquet( + normalized_rels, file_path, schema=RELATIONSHIP_SCHEMA, **options + ) + + def export_knowledge_graph( + self, kg: Dict[str, Any], base_path: Union[str, Path], **options + ) -> None: + """ + Export knowledge graph to multiple Parquet files. + + This method exports a knowledge graph to separate Parquet files for + entities and relationships. Files are named using the base_path with + suffixes: _entities.parquet and _relationships.parquet. + + Args: + kg: Knowledge graph dictionary with 'entities' and 'relationships' keys + base_path: Base path for output files (without extension) + **options: Additional options passed to export methods + + Raises: + ValidationError: If knowledge graph is missing required keys + + Example: + >>> kg = { + ... "entities": [...], + ... "relationships": [...] + ... } + >>> exporter.export_knowledge_graph(kg, "output/kg_base") + # Creates: output/kg_base_entities.parquet and + # output/kg_base_relationships.parquet + """ + if not isinstance(kg, dict): + raise ValidationError( + f"Knowledge graph must be a dictionary, got {type(kg)}" + ) + + if "entities" not in kg and "relationships" not in kg: + raise ValidationError( + "Knowledge graph must contain 'entities' or 'relationships' key" + ) + + base_path = Path(base_path) + ensure_directory(base_path.parent) + + self.logger.debug(f"Exporting knowledge graph to Parquet: {base_path}") + + # Track KG export + tracking_id = self.progress_tracker.start_tracking( + file=str(base_path), + module="export", + submodule="ParquetExporter", + message=f"Exporting knowledge graph to Parquet: {base_path}", + ) + + try: + exported_files = [] + + # Export entities + if "entities" in kg and kg["entities"]: + entities_path = base_path.parent / f"{base_path.stem}_entities.parquet" + self.export_entities(kg["entities"], entities_path, **options) + exported_files.append(entities_path) + self.logger.info(f"Exported entities to: {entities_path}") + + # Export relationships + if "relationships" in kg and kg["relationships"]: + rels_path = base_path.parent / f"{base_path.stem}_relationships.parquet" + self.export_relationships(kg["relationships"], rels_path, **options) + exported_files.append(rels_path) + self.logger.info(f"Exported relationships to: {rels_path}") + + self.logger.info( + f"Exported knowledge graph to {len(exported_files)} Parquet file(s)" + ) + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Exported {len(exported_files)} Parquet files", + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + raise + + def _write_parquet( + self, + data: List[Dict[str, Any]], + file_path: Union[str, Path], + schema: Optional["pa.Schema"] = None, + **options, + ) -> None: + """ + Write data to Parquet file. + + Internal method that handles the actual Parquet file writing using pyarrow. + + Args: + data: List of dictionaries to write + file_path: Output Parquet file path + schema: Parquet schema to use (required for type safety) + **options: Additional parquet write options + + Raises: + ProcessingError: If Parquet writing fails + ValidationError: If schema is not provided + """ + if not schema: + raise ValidationError("Schema is required for Parquet export") + + # Ensure file_path is a Path object + file_path = Path(file_path) + + if not data: + self.logger.warning(f"No data to write to {file_path}") + # Write empty Parquet file with schema + empty_table = pa.table({field.name: [] for field in schema}, schema=schema) + pq.write_table( + empty_table, str(file_path), compression=self.compression, **options + ) + return + + try: + # Create PyArrow table from data using explicit schema + table = pa.Table.from_pylist(data, schema=schema) + + # Write to Parquet file + pq.write_table( + table, str(file_path), compression=self.compression, **options + ) + + file_size = file_path.stat().st_size + self.logger.debug( + f"Wrote {len(data)} row(s) to {file_path} ({file_size} bytes)" + ) + + except pa.ArrowInvalid as e: + raise ProcessingError( + f"Failed to create Parquet table: {e}. " + "Check that data matches schema." + ) + except Exception as e: + raise ProcessingError(f"Failed to write Parquet file: {e}") diff --git a/tests/test_arango_aql_exporter.py b/tests/test_arango_aql_exporter.py index 36e2e000..9a5bb570 100644 --- a/tests/test_arango_aql_exporter.py +++ b/tests/test_arango_aql_exporter.py @@ -1,538 +1,538 @@ -""" -Tests for ArangoDB AQL Exporter Module - -This module contains comprehensive tests for the ArangoDB AQL exporter, -validating AQL syntax generation, node and edge handling, and edge cases. -""" - -import json -import re -import shutil -import tempfile -import unittest -from pathlib import Path - -from semantica.export import ArangoAQLExporter - - -class TestArangoAQLExporter(unittest.TestCase): - """Test cases for ArangoDB AQL Exporter.""" - - def setUp(self): - """Set up test fixtures.""" - self.test_dir = tempfile.mkdtemp() - - # Sample entities for testing - self.entities = [ - { - "id": "e1", - "type": "Person", - "name": "Alice", - "label": "Alice", - "properties": {"age": 30, "email": "alice@example.com"}, - }, - { - "id": "e2", - "type": "Organization", - "name": "Acme Corp", - "label": "Acme Corp", - "properties": {"location": "New York", "founded": 2010}, - }, - { - "id": "e3", - "type": "Person", - "name": "Bob", - "label": "Bob", - "properties": {"age": 25}, - }, - ] - - # Sample relationships for testing - self.relationships = [ - { - "id": "r1", - "source": "e1", - "target": "e2", - "type": "WORKS_FOR", - "properties": {"role": "Engineer", "since": 2020}, - }, - { - "id": "r2", - "source": "e3", - "target": "e2", - "type": "WORKS_FOR", - "properties": {"role": "Manager"}, - }, - { - "id": "r3", - "source": "e1", - "target": "e3", - "type": "KNOWS", - }, - ] - - # Complete knowledge graph - self.kg = { - "entities": self.entities, - "relationships": self.relationships, - "metadata": {"version": "1.0", "created": "2024-01-01"}, - } - - def tearDown(self): - """Clean up test fixtures.""" - shutil.rmtree(self.test_dir) - - def test_exporter_initialization(self): - """Test exporter initialization with default and custom parameters.""" - # Default initialization - exporter = ArangoAQLExporter() - self.assertEqual(exporter.vertex_collection, "vertices") - self.assertEqual(exporter.edge_collection, "edges") - self.assertEqual(exporter.batch_size, 1000) - self.assertTrue(exporter.include_collection_creation) - - # Custom initialization - exporter = ArangoAQLExporter( - vertex_collection="nodes", - edge_collection="links", - batch_size=500, - include_collection_creation=False, - ) - self.assertEqual(exporter.vertex_collection, "nodes") - self.assertEqual(exporter.edge_collection, "links") - self.assertEqual(exporter.batch_size, 500) - self.assertFalse(exporter.include_collection_creation) - - def test_export_knowledge_graph(self): - """Test exporting a complete knowledge graph.""" - exporter = ArangoAQLExporter() - output_path = Path(self.test_dir) / "graph.aql" - - exporter.export_knowledge_graph(self.kg, str(output_path)) - - # Verify file was created - self.assertTrue(output_path.exists()) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Check for collection creation comments - self.assertIn("Create vertex collection", content) - self.assertIn("Create edge collection", content) - - # Check for INSERT statements - self.assertIn("INSERT doc INTO vertices", content) - self.assertIn("INSERT doc INTO edges", content) - - # Check for entity data - self.assertIn("Alice", content) - self.assertIn("Acme Corp", content) - self.assertIn("Bob", content) - self.assertIn("Person", content) - self.assertIn("Organization", content) - - # Check for relationship data - self.assertIn("WORKS_FOR", content) - self.assertIn("KNOWS", content) - self.assertIn("vertices/e1", content) - self.assertIn("vertices/e2", content) - - def test_export_entities_only(self): - """Test exporting only entities (vertices).""" - exporter = ArangoAQLExporter() - output_path = Path(self.test_dir) / "entities.aql" - - exporter.export_entities(self.entities, str(output_path)) - - # Verify file was created - self.assertTrue(output_path.exists()) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Check for vertices INSERT - self.assertIn("INSERT doc INTO vertices", content) - self.assertIn("Alice", content) - self.assertIn("Acme Corp", content) - - # Check that edges are NOT present (empty relationships) - self.assertIn("Attempting to insert 0 edges", content) - - def test_export_relationships_only(self): - """Test exporting only relationships (edges).""" - exporter = ArangoAQLExporter() - output_path = Path(self.test_dir) / "relationships.aql" - - exporter.export_relationships(self.relationships, str(output_path)) - - # Verify file was created - self.assertTrue(output_path.exists()) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Check for edges INSERT - self.assertIn("INSERT doc INTO edges", content) - self.assertIn("WORKS_FOR", content) - self.assertIn("KNOWS", content) - - # Check that vertices section indicates 0 vertices - self.assertIn("Inserting 0 vertices", content) - - def test_custom_collection_names(self): - """Test exporting with custom collection names.""" - exporter = ArangoAQLExporter( - vertex_collection="custom_nodes", edge_collection="custom_edges" - ) - output_path = Path(self.test_dir) / "custom.aql" - - exporter.export(self.kg, str(output_path)) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Check for custom collection names - self.assertIn("INSERT doc INTO custom_nodes", content) - self.assertIn("INSERT doc INTO custom_edges", content) - self.assertIn("custom_nodes/e1", content) - - def test_special_characters_in_properties(self): - """Test handling of special characters in node and edge properties.""" - special_entities = [ - { - "id": "special_1", - "type": "Person", - "name": "O'Brien", - "properties": {"quote": 'She said "hello"', "path": "C:\\Users\\test"}, - } - ] - - special_relationships = [ - { - "id": "special_r1", - "source": "special_1", - "target": "e1", - "type": "KNOWS", - "properties": {"note": "Uses 'quotes' and \"escapes\""}, - } - ] - - kg = {"entities": special_entities, "relationships": special_relationships} - - exporter = ArangoAQLExporter() - output_path = Path(self.test_dir) / "special_chars.aql" - - exporter.export(kg, str(output_path)) - - # Verify file was created and is valid JSON structure - self.assertTrue(output_path.exists()) - - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # The content should contain the special characters properly escaped in JSON - self.assertIn("O'Brien", content) - self.assertIn("She said", content) - - def test_empty_collections(self): - """Test handling of empty entity and relationship collections.""" - empty_kg = {"entities": [], "relationships": []} - - exporter = ArangoAQLExporter() - output_path = Path(self.test_dir) / "empty.aql" - - exporter.export(empty_kg, str(output_path)) - - # Verify file was created - self.assertTrue(output_path.exists()) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Should have collection creation comments but no INSERT statements - self.assertIn("Inserting 0 vertices", content) - self.assertIn("Attempting to insert 0 edges", content) - - def test_missing_source_or_target(self): - """Test handling of edges with missing source or target.""" - invalid_relationships = [ - {"id": "r_invalid_1", "source": "e1", "type": "KNOWS"}, # Missing target - {"id": "r_invalid_2", "target": "e2", "type": "RELATED"}, # Missing source - { - "id": "r_valid", - "source": "e1", - "target": "e2", - "type": "VALID", - }, # Valid - ] - - kg = {"entities": self.entities, "relationships": invalid_relationships} - - exporter = ArangoAQLExporter() - output_path = Path(self.test_dir) / "invalid_edges.aql" - - # Should not raise an exception, but should skip invalid edges - exporter.export(kg, str(output_path)) - - # Verify file was created - self.assertTrue(output_path.exists()) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # The data section should only contain the valid relationship - self.assertIn("VALID", content) - # The header comment should count all input edges, including invalid ones - self.assertIn("Attempting to insert 3 edges", content) - - def test_key_sanitization(self): - """Test sanitization of keys with invalid characters.""" - entities_with_invalid_keys = [ - { - "id": "e1@domain.com", - "type": "Email", - "name": "Test Email", - }, - { - "id": "user/123/profile", - "type": "Profile", - "name": "User Profile", - }, - { - "id": "_system", - "type": "System", - "name": "System Node", - }, - ] - - kg = {"entities": entities_with_invalid_keys, "relationships": []} - - exporter = ArangoAQLExporter() - output_path = Path(self.test_dir) / "sanitized.aql" - - exporter.export(kg, str(output_path)) - - # Verify file was created - self.assertTrue(output_path.exists()) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Keys should be sanitized (@ and / replaced with _) - self.assertIn("e1_domain_com", content) - self.assertIn("user_123_profile", content) - # _system should become k_system (no leading underscore) - self.assertIn("k_system", content) - - def test_batch_processing(self): - """Test batch processing with small batch size.""" - # Create many entities to test batching - many_entities = [ - {"id": f"e{i}", "type": "Node", "name": f"Node {i}"} for i in range(250) - ] - - kg = {"entities": many_entities, "relationships": []} - - # Use small batch size - exporter = ArangoAQLExporter(batch_size=100) - output_path = Path(self.test_dir) / "batched.aql" - - exporter.export(kg, str(output_path)) - - # Verify file was created - self.assertTrue(output_path.exists()) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Count number of INSERT statements (should be 3 batches: 100, 100, 50) - insert_count = content.count("INSERT doc INTO vertices") - self.assertEqual(insert_count, 3, "Should have 3 batches for 250 entities") - - def test_nested_properties(self): - """Test handling of nested dictionaries and lists in properties.""" - entities_with_nested = [ - { - "id": "complex_1", - "type": "ComplexNode", - "name": "Complex", - "properties": { - "nested_dict": {"key1": "value1", "key2": "value2"}, - "nested_list": [1, 2, 3, 4, 5], - "mixed": {"list": [1, 2], "value": "test"}, - }, - } - ] - - kg = {"entities": entities_with_nested, "relationships": []} - - exporter = ArangoAQLExporter() - output_path = Path(self.test_dir) / "nested.aql" - - exporter.export(kg, str(output_path)) - - # Verify file was created - self.assertTrue(output_path.exists()) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Nested structures should be properly serialized as JSON - self.assertIn("nested_dict", content) - self.assertIn("nested_list", content) - # Check for the list values (may be formatted on separate lines) - self.assertIn("1", content) - self.assertIn("2", content) - self.assertIn("3", content) - self.assertIn("4", content) - self.assertIn("5", content) - - def test_aql_syntax_validity(self): - """Test that generated AQL has valid syntax structure.""" - exporter = ArangoAQLExporter() - output_path = Path(self.test_dir) / "syntax_test.aql" - - exporter.export(self.kg, str(output_path)) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Check for valid AQL structure - # Should have FOR doc IN [array] INSERT doc INTO collection pattern - # Using re.search with DOTALL flag - pattern = r"FOR doc IN \[.*?\]\s+INSERT doc INTO \w+" - self.assertTrue( - re.search(pattern, content, re.DOTALL), - f"Pattern '{pattern}' not found in generated AQL", - ) - - # Check that JSON arrays in the INSERT statements are valid - # Find all JSON arrays in the content - json_arrays = re.findall(r"FOR doc IN (\[.*?\])\s+INSERT", content, re.DOTALL) - for json_array_match in json_arrays: - # Extract just the array part - json_str = json_array_match.replace("\n INSERT", "").strip() - try: - # This should parse without errors - parsed = json.loads(json_str) - self.assertIsInstance(parsed, list) - except json.JSONDecodeError as e: - self.fail(f"Invalid JSON in AQL: {e}") - - def test_without_collection_creation(self): - """Test export without collection creation statements.""" - exporter = ArangoAQLExporter(include_collection_creation=False) - output_path = Path(self.test_dir) / "no_creation.aql" - - exporter.export(self.kg, str(output_path)) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Should still have INSERT statements - self.assertIn("INSERT doc INTO vertices", content) - self.assertIn("INSERT doc INTO edges", content) - - def test_export_with_nodes_edges_keys(self): - """Test export using 'nodes' and 'edges' keys. - - Instead of 'entities' and 'relationships'. - """ - kg_alt = { - "nodes": self.entities, - "edges": self.relationships, - } - - exporter = ArangoAQLExporter() - output_path = Path(self.test_dir) / "nodes_edges.aql" - - exporter.export(kg_alt, str(output_path)) - - # Verify file was created - self.assertTrue(output_path.exists()) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Should have all the data - self.assertIn("Alice", content) - self.assertIn("WORKS_FOR", content) - - def test_override_collection_names_in_export(self): - """Test overriding collection names via export options.""" - exporter = ArangoAQLExporter( - vertex_collection="default_v", edge_collection="default_e" - ) - output_path = Path(self.test_dir) / "override.aql" - - # Override via options - exporter.export( - self.kg, - str(output_path), - vertex_collection="override_v", - edge_collection="override_e", - ) - - # Read and verify content - with open(output_path, "r", encoding="utf-8") as f: - content = f.read() - - # Should use overridden names - self.assertIn("INSERT doc INTO override_v", content) - self.assertIn("INSERT doc INTO override_e", content) - self.assertIn("override_v/e1", content) - - def test_invalid_collection_name_on_init(self): - """Test that invalid collection names raise ValueError on initialization.""" - # Test collection name starting with number - with self.assertRaises(ValueError) as context: - ArangoAQLExporter(vertex_collection="123invalid") - self.assertIn("must start with a letter or underscore", str(context.exception)) - - # Test collection name with invalid characters - with self.assertRaises(ValueError) as context: - ArangoAQLExporter(edge_collection="invalid@name") - self.assertIn("contains invalid character", str(context.exception)) - - # Test empty collection name - with self.assertRaises(ValueError) as context: - ArangoAQLExporter(vertex_collection="") - self.assertIn("cannot be empty", str(context.exception)) - - # Test too long collection name - with self.assertRaises(ValueError) as context: - ArangoAQLExporter(vertex_collection="a" * 257) - self.assertIn("exceeds maximum length", str(context.exception)) - - def test_invalid_collection_name_on_export(self): - """Test invalid collection names raise ValueError when overriding.""" - exporter = ArangoAQLExporter() - output_path = Path(self.test_dir) / "test.aql" - - kg = {"entities": self.entities, "relationships": self.relationships} - - # Test invalid vertex collection override - with self.assertRaises(ValueError) as context: - exporter.export(kg, str(output_path), vertex_collection="123invalid") - self.assertIn("must start with a letter or underscore", str(context.exception)) - - # Test invalid edge collection override - with self.assertRaises(ValueError) as context: - exporter.export(kg, str(output_path), edge_collection="invalid@name") - self.assertIn("contains invalid character", str(context.exception)) - - -if __name__ == "__main__": - unittest.main() +""" +Tests for ArangoDB AQL Exporter Module + +This module contains comprehensive tests for the ArangoDB AQL exporter, +validating AQL syntax generation, node and edge handling, and edge cases. +""" + +import json +import re +import shutil +import tempfile +import unittest +from pathlib import Path + +from semantica.export import ArangoAQLExporter + + +class TestArangoAQLExporter(unittest.TestCase): + """Test cases for ArangoDB AQL Exporter.""" + + def setUp(self): + """Set up test fixtures.""" + self.test_dir = tempfile.mkdtemp() + + # Sample entities for testing + self.entities = [ + { + "id": "e1", + "type": "Person", + "name": "Alice", + "label": "Alice", + "properties": {"age": 30, "email": "alice@example.com"}, + }, + { + "id": "e2", + "type": "Organization", + "name": "Acme Corp", + "label": "Acme Corp", + "properties": {"location": "New York", "founded": 2010}, + }, + { + "id": "e3", + "type": "Person", + "name": "Bob", + "label": "Bob", + "properties": {"age": 25}, + }, + ] + + # Sample relationships for testing + self.relationships = [ + { + "id": "r1", + "source": "e1", + "target": "e2", + "type": "WORKS_FOR", + "properties": {"role": "Engineer", "since": 2020}, + }, + { + "id": "r2", + "source": "e3", + "target": "e2", + "type": "WORKS_FOR", + "properties": {"role": "Manager"}, + }, + { + "id": "r3", + "source": "e1", + "target": "e3", + "type": "KNOWS", + }, + ] + + # Complete knowledge graph + self.kg = { + "entities": self.entities, + "relationships": self.relationships, + "metadata": {"version": "1.0", "created": "2024-01-01"}, + } + + def tearDown(self): + """Clean up test fixtures.""" + shutil.rmtree(self.test_dir) + + def test_exporter_initialization(self): + """Test exporter initialization with default and custom parameters.""" + # Default initialization + exporter = ArangoAQLExporter() + self.assertEqual(exporter.vertex_collection, "vertices") + self.assertEqual(exporter.edge_collection, "edges") + self.assertEqual(exporter.batch_size, 1000) + self.assertTrue(exporter.include_collection_creation) + + # Custom initialization + exporter = ArangoAQLExporter( + vertex_collection="nodes", + edge_collection="links", + batch_size=500, + include_collection_creation=False, + ) + self.assertEqual(exporter.vertex_collection, "nodes") + self.assertEqual(exporter.edge_collection, "links") + self.assertEqual(exporter.batch_size, 500) + self.assertFalse(exporter.include_collection_creation) + + def test_export_knowledge_graph(self): + """Test exporting a complete knowledge graph.""" + exporter = ArangoAQLExporter() + output_path = Path(self.test_dir) / "graph.aql" + + exporter.export_knowledge_graph(self.kg, str(output_path)) + + # Verify file was created + self.assertTrue(output_path.exists()) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Check for collection creation comments + self.assertIn("Create vertex collection", content) + self.assertIn("Create edge collection", content) + + # Check for INSERT statements + self.assertIn("INSERT doc INTO vertices", content) + self.assertIn("INSERT doc INTO edges", content) + + # Check for entity data + self.assertIn("Alice", content) + self.assertIn("Acme Corp", content) + self.assertIn("Bob", content) + self.assertIn("Person", content) + self.assertIn("Organization", content) + + # Check for relationship data + self.assertIn("WORKS_FOR", content) + self.assertIn("KNOWS", content) + self.assertIn("vertices/e1", content) + self.assertIn("vertices/e2", content) + + def test_export_entities_only(self): + """Test exporting only entities (vertices).""" + exporter = ArangoAQLExporter() + output_path = Path(self.test_dir) / "entities.aql" + + exporter.export_entities(self.entities, str(output_path)) + + # Verify file was created + self.assertTrue(output_path.exists()) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Check for vertices INSERT + self.assertIn("INSERT doc INTO vertices", content) + self.assertIn("Alice", content) + self.assertIn("Acme Corp", content) + + # Check that edges are NOT present (empty relationships) + self.assertIn("Attempting to insert 0 edges", content) + + def test_export_relationships_only(self): + """Test exporting only relationships (edges).""" + exporter = ArangoAQLExporter() + output_path = Path(self.test_dir) / "relationships.aql" + + exporter.export_relationships(self.relationships, str(output_path)) + + # Verify file was created + self.assertTrue(output_path.exists()) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Check for edges INSERT + self.assertIn("INSERT doc INTO edges", content) + self.assertIn("WORKS_FOR", content) + self.assertIn("KNOWS", content) + + # Check that vertices section indicates 0 vertices + self.assertIn("Inserting 0 vertices", content) + + def test_custom_collection_names(self): + """Test exporting with custom collection names.""" + exporter = ArangoAQLExporter( + vertex_collection="custom_nodes", edge_collection="custom_edges" + ) + output_path = Path(self.test_dir) / "custom.aql" + + exporter.export(self.kg, str(output_path)) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Check for custom collection names + self.assertIn("INSERT doc INTO custom_nodes", content) + self.assertIn("INSERT doc INTO custom_edges", content) + self.assertIn("custom_nodes/e1", content) + + def test_special_characters_in_properties(self): + """Test handling of special characters in node and edge properties.""" + special_entities = [ + { + "id": "special_1", + "type": "Person", + "name": "O'Brien", + "properties": {"quote": 'She said "hello"', "path": "C:\\Users\\test"}, + } + ] + + special_relationships = [ + { + "id": "special_r1", + "source": "special_1", + "target": "e1", + "type": "KNOWS", + "properties": {"note": "Uses 'quotes' and \"escapes\""}, + } + ] + + kg = {"entities": special_entities, "relationships": special_relationships} + + exporter = ArangoAQLExporter() + output_path = Path(self.test_dir) / "special_chars.aql" + + exporter.export(kg, str(output_path)) + + # Verify file was created and is valid JSON structure + self.assertTrue(output_path.exists()) + + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # The content should contain the special characters properly escaped in JSON + self.assertIn("O'Brien", content) + self.assertIn("She said", content) + + def test_empty_collections(self): + """Test handling of empty entity and relationship collections.""" + empty_kg = {"entities": [], "relationships": []} + + exporter = ArangoAQLExporter() + output_path = Path(self.test_dir) / "empty.aql" + + exporter.export(empty_kg, str(output_path)) + + # Verify file was created + self.assertTrue(output_path.exists()) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Should have collection creation comments but no INSERT statements + self.assertIn("Inserting 0 vertices", content) + self.assertIn("Attempting to insert 0 edges", content) + + def test_missing_source_or_target(self): + """Test handling of edges with missing source or target.""" + invalid_relationships = [ + {"id": "r_invalid_1", "source": "e1", "type": "KNOWS"}, # Missing target + {"id": "r_invalid_2", "target": "e2", "type": "RELATED"}, # Missing source + { + "id": "r_valid", + "source": "e1", + "target": "e2", + "type": "VALID", + }, # Valid + ] + + kg = {"entities": self.entities, "relationships": invalid_relationships} + + exporter = ArangoAQLExporter() + output_path = Path(self.test_dir) / "invalid_edges.aql" + + # Should not raise an exception, but should skip invalid edges + exporter.export(kg, str(output_path)) + + # Verify file was created + self.assertTrue(output_path.exists()) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # The data section should only contain the valid relationship + self.assertIn("VALID", content) + # The header comment should count all input edges, including invalid ones + self.assertIn("Attempting to insert 3 edges", content) + + def test_key_sanitization(self): + """Test sanitization of keys with invalid characters.""" + entities_with_invalid_keys = [ + { + "id": "e1@domain.com", + "type": "Email", + "name": "Test Email", + }, + { + "id": "user/123/profile", + "type": "Profile", + "name": "User Profile", + }, + { + "id": "_system", + "type": "System", + "name": "System Node", + }, + ] + + kg = {"entities": entities_with_invalid_keys, "relationships": []} + + exporter = ArangoAQLExporter() + output_path = Path(self.test_dir) / "sanitized.aql" + + exporter.export(kg, str(output_path)) + + # Verify file was created + self.assertTrue(output_path.exists()) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Keys should be sanitized (@ and / replaced with _) + self.assertIn("e1_domain_com", content) + self.assertIn("user_123_profile", content) + # _system should become k_system (no leading underscore) + self.assertIn("k_system", content) + + def test_batch_processing(self): + """Test batch processing with small batch size.""" + # Create many entities to test batching + many_entities = [ + {"id": f"e{i}", "type": "Node", "name": f"Node {i}"} for i in range(250) + ] + + kg = {"entities": many_entities, "relationships": []} + + # Use small batch size + exporter = ArangoAQLExporter(batch_size=100) + output_path = Path(self.test_dir) / "batched.aql" + + exporter.export(kg, str(output_path)) + + # Verify file was created + self.assertTrue(output_path.exists()) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Count number of INSERT statements (should be 3 batches: 100, 100, 50) + insert_count = content.count("INSERT doc INTO vertices") + self.assertEqual(insert_count, 3, "Should have 3 batches for 250 entities") + + def test_nested_properties(self): + """Test handling of nested dictionaries and lists in properties.""" + entities_with_nested = [ + { + "id": "complex_1", + "type": "ComplexNode", + "name": "Complex", + "properties": { + "nested_dict": {"key1": "value1", "key2": "value2"}, + "nested_list": [1, 2, 3, 4, 5], + "mixed": {"list": [1, 2], "value": "test"}, + }, + } + ] + + kg = {"entities": entities_with_nested, "relationships": []} + + exporter = ArangoAQLExporter() + output_path = Path(self.test_dir) / "nested.aql" + + exporter.export(kg, str(output_path)) + + # Verify file was created + self.assertTrue(output_path.exists()) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Nested structures should be properly serialized as JSON + self.assertIn("nested_dict", content) + self.assertIn("nested_list", content) + # Check for the list values (may be formatted on separate lines) + self.assertIn("1", content) + self.assertIn("2", content) + self.assertIn("3", content) + self.assertIn("4", content) + self.assertIn("5", content) + + def test_aql_syntax_validity(self): + """Test that generated AQL has valid syntax structure.""" + exporter = ArangoAQLExporter() + output_path = Path(self.test_dir) / "syntax_test.aql" + + exporter.export(self.kg, str(output_path)) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Check for valid AQL structure + # Should have FOR doc IN [array] INSERT doc INTO collection pattern + # Using re.search with DOTALL flag + pattern = r"FOR doc IN \[.*?\]\s+INSERT doc INTO \w+" + self.assertTrue( + re.search(pattern, content, re.DOTALL), + f"Pattern '{pattern}' not found in generated AQL", + ) + + # Check that JSON arrays in the INSERT statements are valid + # Find all JSON arrays in the content + json_arrays = re.findall(r"FOR doc IN (\[.*?\])\s+INSERT", content, re.DOTALL) + for json_array_match in json_arrays: + # Extract just the array part + json_str = json_array_match.replace("\n INSERT", "").strip() + try: + # This should parse without errors + parsed = json.loads(json_str) + self.assertIsInstance(parsed, list) + except json.JSONDecodeError as e: + self.fail(f"Invalid JSON in AQL: {e}") + + def test_without_collection_creation(self): + """Test export without collection creation statements.""" + exporter = ArangoAQLExporter(include_collection_creation=False) + output_path = Path(self.test_dir) / "no_creation.aql" + + exporter.export(self.kg, str(output_path)) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Should still have INSERT statements + self.assertIn("INSERT doc INTO vertices", content) + self.assertIn("INSERT doc INTO edges", content) + + def test_export_with_nodes_edges_keys(self): + """Test export using 'nodes' and 'edges' keys. + + Instead of 'entities' and 'relationships'. + """ + kg_alt = { + "nodes": self.entities, + "edges": self.relationships, + } + + exporter = ArangoAQLExporter() + output_path = Path(self.test_dir) / "nodes_edges.aql" + + exporter.export(kg_alt, str(output_path)) + + # Verify file was created + self.assertTrue(output_path.exists()) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Should have all the data + self.assertIn("Alice", content) + self.assertIn("WORKS_FOR", content) + + def test_override_collection_names_in_export(self): + """Test overriding collection names via export options.""" + exporter = ArangoAQLExporter( + vertex_collection="default_v", edge_collection="default_e" + ) + output_path = Path(self.test_dir) / "override.aql" + + # Override via options + exporter.export( + self.kg, + str(output_path), + vertex_collection="override_v", + edge_collection="override_e", + ) + + # Read and verify content + with open(output_path, "r", encoding="utf-8") as f: + content = f.read() + + # Should use overridden names + self.assertIn("INSERT doc INTO override_v", content) + self.assertIn("INSERT doc INTO override_e", content) + self.assertIn("override_v/e1", content) + + def test_invalid_collection_name_on_init(self): + """Test that invalid collection names raise ValueError on initialization.""" + # Test collection name starting with number + with self.assertRaises(ValueError) as context: + ArangoAQLExporter(vertex_collection="123invalid") + self.assertIn("must start with a letter or underscore", str(context.exception)) + + # Test collection name with invalid characters + with self.assertRaises(ValueError) as context: + ArangoAQLExporter(edge_collection="invalid@name") + self.assertIn("contains invalid character", str(context.exception)) + + # Test empty collection name + with self.assertRaises(ValueError) as context: + ArangoAQLExporter(vertex_collection="") + self.assertIn("cannot be empty", str(context.exception)) + + # Test too long collection name + with self.assertRaises(ValueError) as context: + ArangoAQLExporter(vertex_collection="a" * 257) + self.assertIn("exceeds maximum length", str(context.exception)) + + def test_invalid_collection_name_on_export(self): + """Test invalid collection names raise ValueError when overriding.""" + exporter = ArangoAQLExporter() + output_path = Path(self.test_dir) / "test.aql" + + kg = {"entities": self.entities, "relationships": self.relationships} + + # Test invalid vertex collection override + with self.assertRaises(ValueError) as context: + exporter.export(kg, str(output_path), vertex_collection="123invalid") + self.assertIn("must start with a letter or underscore", str(context.exception)) + + # Test invalid edge collection override + with self.assertRaises(ValueError) as context: + exporter.export(kg, str(output_path), edge_collection="invalid@name") + self.assertIn("contains invalid character", str(context.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_parquet_exporter.py b/tests/test_parquet_exporter.py index 83ca3a35..f8782afa 100644 --- a/tests/test_parquet_exporter.py +++ b/tests/test_parquet_exporter.py @@ -1,546 +1,546 @@ -""" -Unit tests for Apache Parquet exporter module. - -Tests schema validation, data export, pandas conversion, empty inputs, -and minimal graph structures. -""" - -import shutil -import tempfile -import unittest -from pathlib import Path - -# Try to import pyarrow -try: - import pyarrow.parquet as pq - - PARQUET_AVAILABLE = True -except ImportError: - PARQUET_AVAILABLE = False - -from semantica.export import ParquetExporter -from semantica.utils.exceptions import ValidationError - - -@unittest.skipIf(not PARQUET_AVAILABLE, "pyarrow not installed") -class TestParquetExporter(unittest.TestCase): - """Test cases for ParquetExporter class.""" - - def setUp(self): - """Set up test fixtures.""" - self.test_dir = tempfile.mkdtemp() - - # Sample entities with various field names - self.entities = [ - { - "id": "e1", - "type": "Person", - "name": "Alice", - "label": "Alice", - "confidence": 0.95, - "start": 0, - "end": 5, - "metadata": {"age": 30, "city": "NYC"}, - }, - { - "id": "e2", - "type": "Organization", - "text": "Acme Corp", - "entity_type": "ORG", - "confidence": 0.88, - "start_offset": 10, - "end_offset": 19, - "metadata": {"location": "NY", "employees": 100}, - }, - ] - - # Sample relationships - self.relationships = [ - { - "id": "r1", - "source": "e1", - "target": "e2", - "type": "WORKS_FOR", - "confidence": 0.92, - "metadata": {"role": "Engineer", "since": 2020}, - }, - { - "source_id": "e2", - "target_id": "e1", - "relationship_type": "EMPLOYS", - "confidence": 0.90, - }, - ] - - # Knowledge graph - self.kg = { - "entities": self.entities, - "relationships": self.relationships, - "metadata": {"version": "1.0", "created": "2024-01-01"}, - } - - def tearDown(self): - """Clean up test directory.""" - shutil.rmtree(self.test_dir) - - def test_initialization(self): - """Test ParquetExporter initialization.""" - exporter = ParquetExporter() - self.assertIsNotNone(exporter) - self.assertEqual(exporter.compression, "snappy") - - # Test with different compression - exporter_gzip = ParquetExporter(compression="gzip") - self.assertEqual(exporter_gzip.compression, "gzip") - - exporter_none = ParquetExporter(compression="none") - self.assertIsNone(exporter_none.compression) - - def test_initialization_without_pyarrow(self): - """Test initialization fails gracefully without pyarrow.""" - # This test verifies the constant is correctly set - if not PARQUET_AVAILABLE: - self.assertFalse(PARQUET_AVAILABLE) - - def test_export_entities_basic(self): - """Test basic entity export to Parquet.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "entities.parquet" - - exporter.export_entities(self.entities, str(output_path)) - self.assertTrue(output_path.exists()) - - # Read and verify Parquet file - table = pq.read_table(str(output_path)) - - # Verify schema - from semantica.export.parquet_exporter import ENTITY_SCHEMA - - self.assertEqual(table.schema, ENTITY_SCHEMA) - - # Verify data - self.assertEqual(table.num_rows, 2) - self.assertEqual(table.column("id")[0].as_py(), "e1") - self.assertEqual(table.column("type")[0].as_py(), "Person") - self.assertEqual(table.column("confidence")[0].as_py(), 0.95) - - def test_export_entities_field_normalization(self): - """Test entity field name normalization.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "entities_normalized.parquet" - - # Entities with various field name variations - varied_entities = [ - {"id": "e1", "text": "Entity 1", "type": "TYPE1"}, - {"entity_id": "e2", "label": "Entity 2", "entity_type": "TYPE2"}, - {"id": "e3", "name": "Entity 3", "type": "TYPE3"}, - ] - - exporter.export_entities(varied_entities, str(output_path)) - self.assertTrue(output_path.exists()) - - # Read and verify normalization - table = pq.read_table(str(output_path)) - - self.assertEqual(table.num_rows, 3) - self.assertEqual(table.column("id")[0].as_py(), "e1") - self.assertEqual(table.column("text")[0].as_py(), "Entity 1") - self.assertEqual(table.column("text")[1].as_py(), "Entity 2") - self.assertEqual(table.column("text")[2].as_py(), "Entity 3") - - def test_export_entities_with_compression(self): - """Test entity export with different compression codecs.""" - import pyarrow.parquet as pq_module - - for compression in ["snappy", "gzip", "brotli", "zstd", "lz4", "none"]: - with self.subTest(compression=compression): - # Check if codec is available in this pyarrow build - try: - # Test codec availability by checking compression opts - if compression != "none": - codec_available = compression.upper() in dir( - pq_module.lib.Codec - ) - if not codec_available: - self.skipTest( - f"Codec {compression} not available in " "pyarrow build" - ) - except AttributeError: - # If we can't check, just try and skip on error - pass - - try: - exporter = ParquetExporter(compression=compression) - output_path = ( - Path(self.test_dir) / f"entities_{compression}.parquet" - ) - - exporter.export_entities(self.entities, str(output_path)) - self.assertTrue(output_path.exists()) - - # Verify file can be read - table = pq.read_table(str(output_path)) - self.assertEqual(table.num_rows, 2) - except (ImportError, RuntimeError, OSError) as e: - if "codec" in str(e).lower() or "compression" in str(e).lower(): - self.skipTest(f"Codec {compression} not available: {e}") - raise - - def test_export_entities_empty(self): - """Test exporting empty entities list raises ValidationError.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "empty_entities.parquet" - - with self.assertRaises(ValidationError): - exporter.export_entities([], str(output_path)) - - def test_export_entities_metadata_handling(self): - """Test entity metadata is correctly serialized.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "entities_metadata.parquet" - - exporter.export_entities(self.entities, str(output_path)) - - # Read and verify metadata - table = pq.read_table(str(output_path)) - metadata_col = table.column("metadata") - - # First entity should have metadata - first_metadata = metadata_col[0].as_py() - self.assertIsNotNone(first_metadata) - self.assertIn("keys", first_metadata) - self.assertIn("values", first_metadata) - self.assertEqual(set(first_metadata["keys"]), {"age", "city"}) - - def test_export_relationships_basic(self): - """Test basic relationship export to Parquet.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "relationships.parquet" - - exporter.export_relationships(self.relationships, str(output_path)) - self.assertTrue(output_path.exists()) - - # Read and verify Parquet file - table = pq.read_table(str(output_path)) - - # Verify schema - from semantica.export.parquet_exporter import RELATIONSHIP_SCHEMA - - self.assertEqual(table.schema, RELATIONSHIP_SCHEMA) - - # Verify data - self.assertEqual(table.num_rows, 2) - self.assertEqual(table.column("id")[0].as_py(), "r1") - self.assertEqual(table.column("source_id")[0].as_py(), "e1") - self.assertEqual(table.column("target_id")[0].as_py(), "e2") - self.assertEqual(table.column("type")[0].as_py(), "WORKS_FOR") - - def test_export_relationships_field_normalization(self): - """Test relationship field name normalization.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "relationships_normalized.parquet" - - # Relationships with various field name variations - varied_rels = [ - {"id": "r1", "source": "e1", "target": "e2", "type": "TYPE1"}, - {"source_id": "e2", "target_id": "e3", "relationship_type": "TYPE2"}, - {"from_id": "e3", "to_id": "e1", "relation_type": "TYPE3"}, - ] - - exporter.export_relationships(varied_rels, str(output_path)) - self.assertTrue(output_path.exists()) - - # Read and verify normalization - table = pq.read_table(str(output_path)) - - self.assertEqual(table.num_rows, 3) - self.assertEqual(table.column("source_id")[0].as_py(), "e1") - self.assertEqual(table.column("target_id")[0].as_py(), "e2") - - def test_export_relationships_empty(self): - """Test exporting empty relationships list raises ValidationError.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "empty_relationships.parquet" - - with self.assertRaises(ValidationError): - exporter.export_relationships([], str(output_path)) - - def test_export_relationships_auto_id_generation(self): - """Test relationship ID is auto-generated when missing.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "relationships_auto_id.parquet" - - # Relationships without IDs - rels_no_id = [ - {"source_id": "e1", "target_id": "e2", "type": "REL1"}, - {"source_id": "e2", "target_id": "e3", "type": "REL2"}, - ] - - exporter.export_relationships(rels_no_id, str(output_path)) - - # Read and verify IDs were generated - table = pq.read_table(str(output_path)) - self.assertEqual(table.num_rows, 2) - self.assertIsNotNone(table.column("id")[0].as_py()) - self.assertIsNotNone(table.column("id")[1].as_py()) - - def test_export_knowledge_graph_basic(self): - """Test basic knowledge graph export to multiple Parquet files.""" - exporter = ParquetExporter() - base_path = Path(self.test_dir) / "kg" - - exporter.export_knowledge_graph(self.kg, str(base_path)) - - # Verify files were created - entities_path = Path(self.test_dir) / "kg_entities.parquet" - rels_path = Path(self.test_dir) / "kg_relationships.parquet" - - self.assertTrue(entities_path.exists()) - self.assertTrue(rels_path.exists()) - - # Verify entities - entities_table = pq.read_table(str(entities_path)) - self.assertEqual(entities_table.num_rows, 2) - - # Verify relationships - rels_table = pq.read_table(str(rels_path)) - self.assertEqual(rels_table.num_rows, 2) - - def test_export_knowledge_graph_invalid_input(self): - """Test knowledge graph export with invalid input raises ValidationError.""" - exporter = ParquetExporter() - base_path = Path(self.test_dir) / "kg_invalid" - - # Not a dictionary - with self.assertRaises(ValidationError): - exporter.export_knowledge_graph("not a dict", str(base_path)) - - # Missing both entities and relationships - with self.assertRaises(ValidationError): - exporter.export_knowledge_graph({"metadata": {}}, str(base_path)) - - def test_export_knowledge_graph_partial(self): - """Test knowledge graph export with only entities or relationships.""" - exporter = ParquetExporter() - - # Only entities - kg_entities_only = {"entities": self.entities} - base_path_ent = Path(self.test_dir) / "kg_entities_only" - exporter.export_knowledge_graph(kg_entities_only, str(base_path_ent)) - - entities_path = Path(self.test_dir) / "kg_entities_only_entities.parquet" - self.assertTrue(entities_path.exists()) - - # Only relationships - kg_rels_only = {"relationships": self.relationships} - base_path_rel = Path(self.test_dir) / "kg_rels_only" - exporter.export_knowledge_graph(kg_rels_only, str(base_path_rel)) - - rels_path = Path(self.test_dir) / "kg_rels_only_relationships.parquet" - self.assertTrue(rels_path.exists()) - - def test_export_generic_list(self): - """Test generic export with list of dictionaries.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "generic_list.parquet" - - exporter.export(self.entities, str(output_path), schema=None) - self.assertTrue(output_path.exists()) - - # Verify file can be read (schema auto-selected based on data structure) - table = pq.read_table(str(output_path)) - self.assertEqual(table.num_rows, 2) - - def test_export_generic_dict(self): - """Test generic export with dictionary (multiple files).""" - exporter = ParquetExporter() - base_path = Path(self.test_dir) / "generic_dict" - - data_dict = {"entities": self.entities, "relationships": self.relationships} - - exporter.export(data_dict, str(base_path)) - - # Verify both files were created - entities_path = Path(self.test_dir) / "generic_dict_entities.parquet" - rels_path = Path(self.test_dir) / "generic_dict_relationships.parquet" - - self.assertTrue(entities_path.exists()) - self.assertTrue(rels_path.exists()) - - def test_export_invalid_data_type(self): - """Test export with invalid data type raises ValidationError.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "invalid.parquet" - - with self.assertRaises(ValidationError): - exporter.export("invalid string data", str(output_path)) - - def test_pandas_compatibility(self): - """Test exported Parquet files can be read by pandas.""" - try: - import pandas as pd - except ImportError: - self.skipTest("pandas not installed") - - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "pandas_test.parquet" - - exporter.export_entities(self.entities, str(output_path)) - - # Read with pandas - df = pd.read_parquet(str(output_path)) - - self.assertEqual(len(df), 2) - self.assertIn("id", df.columns) - self.assertIn("text", df.columns) - self.assertIn("type", df.columns) - - def test_file_size_comparison(self): - """Test Parquet file sizes with different compression.""" - exporter_snappy = ParquetExporter(compression="snappy") - exporter_gzip = ParquetExporter(compression="gzip") - exporter_none = ParquetExporter(compression="none") - - path_snappy = Path(self.test_dir) / "size_snappy.parquet" - path_gzip = Path(self.test_dir) / "size_gzip.parquet" - path_none = Path(self.test_dir) / "size_none.parquet" - - # Create larger dataset for meaningful comparison - large_entities = self.entities * 100 - - exporter_snappy.export_entities(large_entities, str(path_snappy)) - exporter_gzip.export_entities(large_entities, str(path_gzip)) - exporter_none.export_entities(large_entities, str(path_none)) - - size_snappy = path_snappy.stat().st_size - size_none = path_none.stat().st_size - - # Uncompressed should be largest - self.assertGreater(size_none, size_snappy) - - # All files should be readable - for path in [path_snappy, path_gzip, path_none]: - table = pq.read_table(str(path)) - self.assertEqual(table.num_rows, 200) - - def test_entity_missing_id_skipped(self): - """Test entities without IDs are skipped with warning.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "entities_missing_id.parquet" - - # Mix of entities with and without IDs - entities_mixed = [ - {"id": "e1", "text": "Valid Entity"}, - {"text": "Missing ID"}, # No ID - {"id": "e2", "text": "Another Valid"}, - ] - - exporter.export_entities(entities_mixed, str(output_path)) - - # Only 2 entities should be exported (one without ID is skipped) - table = pq.read_table(str(output_path)) - self.assertEqual(table.num_rows, 2) - - def test_relationship_missing_source_target_skipped(self): - """Test relationships without source/target are skipped with warning.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "rels_missing.parquet" - - # Mix of valid and invalid relationships - rels_mixed = [ - {"id": "r1", "source_id": "e1", "target_id": "e2"}, - {"id": "r2", "target_id": "e2"}, # Missing source - {"id": "r3", "source_id": "e1"}, # Missing target - {"id": "r4", "source_id": "e3", "target_id": "e4"}, - ] - - exporter.export_relationships(rels_mixed, str(output_path)) - - # Only 2 valid relationships should be exported - table = pq.read_table(str(output_path)) - self.assertEqual(table.num_rows, 2) - - def test_all_entities_skipped_raises_error(self): - """Test that exporting entities with all skipped raises ValidationError.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "all_skipped.parquet" - - # All entities missing IDs - bad_entities = [ - {"text": "No ID 1"}, - {"text": "No ID 2"}, - "not a dict", - ] - - with self.assertRaises(ValidationError) as cm: - exporter.export_entities(bad_entities, str(output_path)) - - self.assertIn("No valid entities", str(cm.exception)) - - def test_all_relationships_skipped_raises_error(self): - """Test that exporting relationships with all skipped raises ValidationError.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "all_rels_skipped.parquet" - - # All relationships missing source or target - bad_rels = [ - {"id": "r1", "source_id": "e1"}, # Missing target - {"id": "r2", "target_id": "e2"}, # Missing source - "not a dict", - ] - - with self.assertRaises(ValidationError) as cm: - exporter.export_relationships(bad_rels, str(output_path)) - - self.assertIn("No valid relationships", str(cm.exception)) - - def test_invalid_confidence_values_handled(self): - """Test that invalid confidence values are handled gracefully.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "invalid_confidence.parquet" - - entities_with_invalid_conf = [ - {"id": "e1", "text": "Valid", "confidence": 0.9}, - {"id": "e2", "text": "String conf", "confidence": "invalid"}, - {"id": "e3", "text": "None conf", "confidence": None}, - ] - - exporter.export_entities(entities_with_invalid_conf, str(output_path)) - - table = pq.read_table(str(output_path)) - self.assertEqual(table.num_rows, 3) - # First entity has valid confidence - self.assertEqual(table.column("confidence")[0].as_py(), 0.9) - # Second entity has invalid confidence (should be None) - self.assertIsNone(table.column("confidence")[1].as_py()) - # Third entity has None confidence - self.assertIsNone(table.column("confidence")[2].as_py()) - - def test_invalid_start_end_values_handled(self): - """Test that invalid start/end offset values are handled gracefully.""" - exporter = ParquetExporter() - output_path = Path(self.test_dir) / "invalid_offsets.parquet" - - entities_with_invalid_offsets = [ - {"id": "e1", "text": "Valid", "start": 0, "end": 10}, - {"id": "e2", "text": "String offsets", "start": "abc", "end": "def"}, - {"id": "e3", "text": "None offsets", "start": None, "end": None}, - ] - - exporter.export_entities(entities_with_invalid_offsets, str(output_path)) - - table = pq.read_table(str(output_path)) - self.assertEqual(table.num_rows, 3) - # First entity has valid offsets - self.assertEqual(table.column("start")[0].as_py(), 0) - self.assertEqual(table.column("end")[0].as_py(), 10) - # Second entity has invalid offsets (should be None) - self.assertIsNone(table.column("start")[1].as_py()) - self.assertIsNone(table.column("end")[1].as_py()) - # Third entity has None offsets - self.assertIsNone(table.column("start")[2].as_py()) - self.assertIsNone(table.column("end")[2].as_py()) - - -if __name__ == "__main__": - unittest.main() +""" +Unit tests for Apache Parquet exporter module. + +Tests schema validation, data export, pandas conversion, empty inputs, +and minimal graph structures. +""" + +import shutil +import tempfile +import unittest +from pathlib import Path + +# Try to import pyarrow +try: + import pyarrow.parquet as pq + + PARQUET_AVAILABLE = True +except ImportError: + PARQUET_AVAILABLE = False + +from semantica.export import ParquetExporter +from semantica.utils.exceptions import ValidationError + + +@unittest.skipIf(not PARQUET_AVAILABLE, "pyarrow not installed") +class TestParquetExporter(unittest.TestCase): + """Test cases for ParquetExporter class.""" + + def setUp(self): + """Set up test fixtures.""" + self.test_dir = tempfile.mkdtemp() + + # Sample entities with various field names + self.entities = [ + { + "id": "e1", + "type": "Person", + "name": "Alice", + "label": "Alice", + "confidence": 0.95, + "start": 0, + "end": 5, + "metadata": {"age": 30, "city": "NYC"}, + }, + { + "id": "e2", + "type": "Organization", + "text": "Acme Corp", + "entity_type": "ORG", + "confidence": 0.88, + "start_offset": 10, + "end_offset": 19, + "metadata": {"location": "NY", "employees": 100}, + }, + ] + + # Sample relationships + self.relationships = [ + { + "id": "r1", + "source": "e1", + "target": "e2", + "type": "WORKS_FOR", + "confidence": 0.92, + "metadata": {"role": "Engineer", "since": 2020}, + }, + { + "source_id": "e2", + "target_id": "e1", + "relationship_type": "EMPLOYS", + "confidence": 0.90, + }, + ] + + # Knowledge graph + self.kg = { + "entities": self.entities, + "relationships": self.relationships, + "metadata": {"version": "1.0", "created": "2024-01-01"}, + } + + def tearDown(self): + """Clean up test directory.""" + shutil.rmtree(self.test_dir) + + def test_initialization(self): + """Test ParquetExporter initialization.""" + exporter = ParquetExporter() + self.assertIsNotNone(exporter) + self.assertEqual(exporter.compression, "snappy") + + # Test with different compression + exporter_gzip = ParquetExporter(compression="gzip") + self.assertEqual(exporter_gzip.compression, "gzip") + + exporter_none = ParquetExporter(compression="none") + self.assertIsNone(exporter_none.compression) + + def test_initialization_without_pyarrow(self): + """Test initialization fails gracefully without pyarrow.""" + # This test verifies the constant is correctly set + if not PARQUET_AVAILABLE: + self.assertFalse(PARQUET_AVAILABLE) + + def test_export_entities_basic(self): + """Test basic entity export to Parquet.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "entities.parquet" + + exporter.export_entities(self.entities, str(output_path)) + self.assertTrue(output_path.exists()) + + # Read and verify Parquet file + table = pq.read_table(str(output_path)) + + # Verify schema + from semantica.export.parquet_exporter import ENTITY_SCHEMA + + self.assertEqual(table.schema, ENTITY_SCHEMA) + + # Verify data + self.assertEqual(table.num_rows, 2) + self.assertEqual(table.column("id")[0].as_py(), "e1") + self.assertEqual(table.column("type")[0].as_py(), "Person") + self.assertEqual(table.column("confidence")[0].as_py(), 0.95) + + def test_export_entities_field_normalization(self): + """Test entity field name normalization.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "entities_normalized.parquet" + + # Entities with various field name variations + varied_entities = [ + {"id": "e1", "text": "Entity 1", "type": "TYPE1"}, + {"entity_id": "e2", "label": "Entity 2", "entity_type": "TYPE2"}, + {"id": "e3", "name": "Entity 3", "type": "TYPE3"}, + ] + + exporter.export_entities(varied_entities, str(output_path)) + self.assertTrue(output_path.exists()) + + # Read and verify normalization + table = pq.read_table(str(output_path)) + + self.assertEqual(table.num_rows, 3) + self.assertEqual(table.column("id")[0].as_py(), "e1") + self.assertEqual(table.column("text")[0].as_py(), "Entity 1") + self.assertEqual(table.column("text")[1].as_py(), "Entity 2") + self.assertEqual(table.column("text")[2].as_py(), "Entity 3") + + def test_export_entities_with_compression(self): + """Test entity export with different compression codecs.""" + import pyarrow.parquet as pq_module + + for compression in ["snappy", "gzip", "brotli", "zstd", "lz4", "none"]: + with self.subTest(compression=compression): + # Check if codec is available in this pyarrow build + try: + # Test codec availability by checking compression opts + if compression != "none": + codec_available = compression.upper() in dir( + pq_module.lib.Codec + ) + if not codec_available: + self.skipTest( + f"Codec {compression} not available in " "pyarrow build" + ) + except AttributeError: + # If we can't check, just try and skip on error + pass + + try: + exporter = ParquetExporter(compression=compression) + output_path = ( + Path(self.test_dir) / f"entities_{compression}.parquet" + ) + + exporter.export_entities(self.entities, str(output_path)) + self.assertTrue(output_path.exists()) + + # Verify file can be read + table = pq.read_table(str(output_path)) + self.assertEqual(table.num_rows, 2) + except (ImportError, RuntimeError, OSError) as e: + if "codec" in str(e).lower() or "compression" in str(e).lower(): + self.skipTest(f"Codec {compression} not available: {e}") + raise + + def test_export_entities_empty(self): + """Test exporting empty entities list raises ValidationError.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "empty_entities.parquet" + + with self.assertRaises(ValidationError): + exporter.export_entities([], str(output_path)) + + def test_export_entities_metadata_handling(self): + """Test entity metadata is correctly serialized.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "entities_metadata.parquet" + + exporter.export_entities(self.entities, str(output_path)) + + # Read and verify metadata + table = pq.read_table(str(output_path)) + metadata_col = table.column("metadata") + + # First entity should have metadata + first_metadata = metadata_col[0].as_py() + self.assertIsNotNone(first_metadata) + self.assertIn("keys", first_metadata) + self.assertIn("values", first_metadata) + self.assertEqual(set(first_metadata["keys"]), {"age", "city"}) + + def test_export_relationships_basic(self): + """Test basic relationship export to Parquet.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "relationships.parquet" + + exporter.export_relationships(self.relationships, str(output_path)) + self.assertTrue(output_path.exists()) + + # Read and verify Parquet file + table = pq.read_table(str(output_path)) + + # Verify schema + from semantica.export.parquet_exporter import RELATIONSHIP_SCHEMA + + self.assertEqual(table.schema, RELATIONSHIP_SCHEMA) + + # Verify data + self.assertEqual(table.num_rows, 2) + self.assertEqual(table.column("id")[0].as_py(), "r1") + self.assertEqual(table.column("source_id")[0].as_py(), "e1") + self.assertEqual(table.column("target_id")[0].as_py(), "e2") + self.assertEqual(table.column("type")[0].as_py(), "WORKS_FOR") + + def test_export_relationships_field_normalization(self): + """Test relationship field name normalization.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "relationships_normalized.parquet" + + # Relationships with various field name variations + varied_rels = [ + {"id": "r1", "source": "e1", "target": "e2", "type": "TYPE1"}, + {"source_id": "e2", "target_id": "e3", "relationship_type": "TYPE2"}, + {"from_id": "e3", "to_id": "e1", "relation_type": "TYPE3"}, + ] + + exporter.export_relationships(varied_rels, str(output_path)) + self.assertTrue(output_path.exists()) + + # Read and verify normalization + table = pq.read_table(str(output_path)) + + self.assertEqual(table.num_rows, 3) + self.assertEqual(table.column("source_id")[0].as_py(), "e1") + self.assertEqual(table.column("target_id")[0].as_py(), "e2") + + def test_export_relationships_empty(self): + """Test exporting empty relationships list raises ValidationError.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "empty_relationships.parquet" + + with self.assertRaises(ValidationError): + exporter.export_relationships([], str(output_path)) + + def test_export_relationships_auto_id_generation(self): + """Test relationship ID is auto-generated when missing.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "relationships_auto_id.parquet" + + # Relationships without IDs + rels_no_id = [ + {"source_id": "e1", "target_id": "e2", "type": "REL1"}, + {"source_id": "e2", "target_id": "e3", "type": "REL2"}, + ] + + exporter.export_relationships(rels_no_id, str(output_path)) + + # Read and verify IDs were generated + table = pq.read_table(str(output_path)) + self.assertEqual(table.num_rows, 2) + self.assertIsNotNone(table.column("id")[0].as_py()) + self.assertIsNotNone(table.column("id")[1].as_py()) + + def test_export_knowledge_graph_basic(self): + """Test basic knowledge graph export to multiple Parquet files.""" + exporter = ParquetExporter() + base_path = Path(self.test_dir) / "kg" + + exporter.export_knowledge_graph(self.kg, str(base_path)) + + # Verify files were created + entities_path = Path(self.test_dir) / "kg_entities.parquet" + rels_path = Path(self.test_dir) / "kg_relationships.parquet" + + self.assertTrue(entities_path.exists()) + self.assertTrue(rels_path.exists()) + + # Verify entities + entities_table = pq.read_table(str(entities_path)) + self.assertEqual(entities_table.num_rows, 2) + + # Verify relationships + rels_table = pq.read_table(str(rels_path)) + self.assertEqual(rels_table.num_rows, 2) + + def test_export_knowledge_graph_invalid_input(self): + """Test knowledge graph export with invalid input raises ValidationError.""" + exporter = ParquetExporter() + base_path = Path(self.test_dir) / "kg_invalid" + + # Not a dictionary + with self.assertRaises(ValidationError): + exporter.export_knowledge_graph("not a dict", str(base_path)) + + # Missing both entities and relationships + with self.assertRaises(ValidationError): + exporter.export_knowledge_graph({"metadata": {}}, str(base_path)) + + def test_export_knowledge_graph_partial(self): + """Test knowledge graph export with only entities or relationships.""" + exporter = ParquetExporter() + + # Only entities + kg_entities_only = {"entities": self.entities} + base_path_ent = Path(self.test_dir) / "kg_entities_only" + exporter.export_knowledge_graph(kg_entities_only, str(base_path_ent)) + + entities_path = Path(self.test_dir) / "kg_entities_only_entities.parquet" + self.assertTrue(entities_path.exists()) + + # Only relationships + kg_rels_only = {"relationships": self.relationships} + base_path_rel = Path(self.test_dir) / "kg_rels_only" + exporter.export_knowledge_graph(kg_rels_only, str(base_path_rel)) + + rels_path = Path(self.test_dir) / "kg_rels_only_relationships.parquet" + self.assertTrue(rels_path.exists()) + + def test_export_generic_list(self): + """Test generic export with list of dictionaries.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "generic_list.parquet" + + exporter.export(self.entities, str(output_path), schema=None) + self.assertTrue(output_path.exists()) + + # Verify file can be read (schema auto-selected based on data structure) + table = pq.read_table(str(output_path)) + self.assertEqual(table.num_rows, 2) + + def test_export_generic_dict(self): + """Test generic export with dictionary (multiple files).""" + exporter = ParquetExporter() + base_path = Path(self.test_dir) / "generic_dict" + + data_dict = {"entities": self.entities, "relationships": self.relationships} + + exporter.export(data_dict, str(base_path)) + + # Verify both files were created + entities_path = Path(self.test_dir) / "generic_dict_entities.parquet" + rels_path = Path(self.test_dir) / "generic_dict_relationships.parquet" + + self.assertTrue(entities_path.exists()) + self.assertTrue(rels_path.exists()) + + def test_export_invalid_data_type(self): + """Test export with invalid data type raises ValidationError.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "invalid.parquet" + + with self.assertRaises(ValidationError): + exporter.export("invalid string data", str(output_path)) + + def test_pandas_compatibility(self): + """Test exported Parquet files can be read by pandas.""" + try: + import pandas as pd + except ImportError: + self.skipTest("pandas not installed") + + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "pandas_test.parquet" + + exporter.export_entities(self.entities, str(output_path)) + + # Read with pandas + df = pd.read_parquet(str(output_path)) + + self.assertEqual(len(df), 2) + self.assertIn("id", df.columns) + self.assertIn("text", df.columns) + self.assertIn("type", df.columns) + + def test_file_size_comparison(self): + """Test Parquet file sizes with different compression.""" + exporter_snappy = ParquetExporter(compression="snappy") + exporter_gzip = ParquetExporter(compression="gzip") + exporter_none = ParquetExporter(compression="none") + + path_snappy = Path(self.test_dir) / "size_snappy.parquet" + path_gzip = Path(self.test_dir) / "size_gzip.parquet" + path_none = Path(self.test_dir) / "size_none.parquet" + + # Create larger dataset for meaningful comparison + large_entities = self.entities * 100 + + exporter_snappy.export_entities(large_entities, str(path_snappy)) + exporter_gzip.export_entities(large_entities, str(path_gzip)) + exporter_none.export_entities(large_entities, str(path_none)) + + size_snappy = path_snappy.stat().st_size + size_none = path_none.stat().st_size + + # Uncompressed should be largest + self.assertGreater(size_none, size_snappy) + + # All files should be readable + for path in [path_snappy, path_gzip, path_none]: + table = pq.read_table(str(path)) + self.assertEqual(table.num_rows, 200) + + def test_entity_missing_id_skipped(self): + """Test entities without IDs are skipped with warning.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "entities_missing_id.parquet" + + # Mix of entities with and without IDs + entities_mixed = [ + {"id": "e1", "text": "Valid Entity"}, + {"text": "Missing ID"}, # No ID + {"id": "e2", "text": "Another Valid"}, + ] + + exporter.export_entities(entities_mixed, str(output_path)) + + # Only 2 entities should be exported (one without ID is skipped) + table = pq.read_table(str(output_path)) + self.assertEqual(table.num_rows, 2) + + def test_relationship_missing_source_target_skipped(self): + """Test relationships without source/target are skipped with warning.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "rels_missing.parquet" + + # Mix of valid and invalid relationships + rels_mixed = [ + {"id": "r1", "source_id": "e1", "target_id": "e2"}, + {"id": "r2", "target_id": "e2"}, # Missing source + {"id": "r3", "source_id": "e1"}, # Missing target + {"id": "r4", "source_id": "e3", "target_id": "e4"}, + ] + + exporter.export_relationships(rels_mixed, str(output_path)) + + # Only 2 valid relationships should be exported + table = pq.read_table(str(output_path)) + self.assertEqual(table.num_rows, 2) + + def test_all_entities_skipped_raises_error(self): + """Test that exporting entities with all skipped raises ValidationError.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "all_skipped.parquet" + + # All entities missing IDs + bad_entities = [ + {"text": "No ID 1"}, + {"text": "No ID 2"}, + "not a dict", + ] + + with self.assertRaises(ValidationError) as cm: + exporter.export_entities(bad_entities, str(output_path)) + + self.assertIn("No valid entities", str(cm.exception)) + + def test_all_relationships_skipped_raises_error(self): + """Test that exporting relationships with all skipped raises ValidationError.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "all_rels_skipped.parquet" + + # All relationships missing source or target + bad_rels = [ + {"id": "r1", "source_id": "e1"}, # Missing target + {"id": "r2", "target_id": "e2"}, # Missing source + "not a dict", + ] + + with self.assertRaises(ValidationError) as cm: + exporter.export_relationships(bad_rels, str(output_path)) + + self.assertIn("No valid relationships", str(cm.exception)) + + def test_invalid_confidence_values_handled(self): + """Test that invalid confidence values are handled gracefully.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "invalid_confidence.parquet" + + entities_with_invalid_conf = [ + {"id": "e1", "text": "Valid", "confidence": 0.9}, + {"id": "e2", "text": "String conf", "confidence": "invalid"}, + {"id": "e3", "text": "None conf", "confidence": None}, + ] + + exporter.export_entities(entities_with_invalid_conf, str(output_path)) + + table = pq.read_table(str(output_path)) + self.assertEqual(table.num_rows, 3) + # First entity has valid confidence + self.assertEqual(table.column("confidence")[0].as_py(), 0.9) + # Second entity has invalid confidence (should be None) + self.assertIsNone(table.column("confidence")[1].as_py()) + # Third entity has None confidence + self.assertIsNone(table.column("confidence")[2].as_py()) + + def test_invalid_start_end_values_handled(self): + """Test that invalid start/end offset values are handled gracefully.""" + exporter = ParquetExporter() + output_path = Path(self.test_dir) / "invalid_offsets.parquet" + + entities_with_invalid_offsets = [ + {"id": "e1", "text": "Valid", "start": 0, "end": 10}, + {"id": "e2", "text": "String offsets", "start": "abc", "end": "def"}, + {"id": "e3", "text": "None offsets", "start": None, "end": None}, + ] + + exporter.export_entities(entities_with_invalid_offsets, str(output_path)) + + table = pq.read_table(str(output_path)) + self.assertEqual(table.num_rows, 3) + # First entity has valid offsets + self.assertEqual(table.column("start")[0].as_py(), 0) + self.assertEqual(table.column("end")[0].as_py(), 10) + # Second entity has invalid offsets (should be None) + self.assertIsNone(table.column("start")[1].as_py()) + self.assertIsNone(table.column("end")[1].as_py()) + # Third entity has None offsets + self.assertIsNone(table.column("start")[2].as_py()) + self.assertIsNone(table.column("end")[2].as_py()) + + +if __name__ == "__main__": + unittest.main()