diff --git a/docs/arrow_exporter.md b/docs/arrow_exporter.md new file mode 100644 index 00000000..64437a88 --- /dev/null +++ b/docs/arrow_exporter.md @@ -0,0 +1,269 @@ +# Apache Arrow Exporter + +## Overview + +The Apache Arrow exporter provides high-performance columnar data export for Semantica's knowledge graphs, entities, and relationships. It uses explicit schemas (no inference) and writes Arrow IPC files (.arrow) that are compatible with Pandas and DuckDB. + +## Features + +- **Explicit Schemas**: Pre-defined schemas for entities and relationships (no inference) +- **Columnar Format**: Efficient storage and fast analytics +- **Metadata Support**: Converts metadata dictionaries to Arrow struct fields +- **Field Normalization**: Handles various entity and relationship field name variations +- **Progress Tracking**: Integrated progress monitoring +- **Error Handling**: Structured error handling with detailed logging +- **Pandas/DuckDB Compatible**: Direct conversion to DataFrames and SQL queries + +## Installation + +The Arrow exporter requires PyArrow: + +```bash +pip install pyarrow +``` + +## Usage + +### Basic Usage + +```python +from semantica.export import ArrowExporter + +# Initialize exporter +exporter = ArrowExporter() + +# Export entities +entities = [ + {"id": "e1", "text": "Alice", "type": "Person", "confidence": 0.95}, + {"id": "e2", "text": "Acme Corp", "type": "Organization", "confidence": 0.88} +] +exporter.export_entities(entities, "entities.arrow") + +# Export relationships +relationships = [ + {"id": "r1", "source_id": "e1", "target_id": "e2", "type": "WORKS_FOR"} +] +exporter.export_relationships(relationships, "relationships.arrow") + +# Export knowledge graph +knowledge_graph = { + "entities": entities, + "relationships": relationships +} +exporter.export_knowledge_graph(knowledge_graph, "kg_base") +# Creates: kg_base_entities.arrow, kg_base_relationships.arrow +``` + +### Using Convenience Function + +```python +from semantica.export import export_arrow + +# Simple export +export_arrow(entities, "entities.arrow") + +# Export multiple types +data = { + "entities": entities, + "relationships": relationships +} +export_arrow(data, "output_base") +``` + +### With Compression + +```python +# Use LZ4 compression +exporter = ArrowExporter(compression="lz4") +exporter.export_entities(entities, "entities_compressed.arrow") +``` + +## Schemas + +### Entity Schema + +```python +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 + +```python +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), +]) +``` + +## Field Normalization + +The exporter automatically normalizes field names: + +**Entities:** +- `text`, `label`, `name` → `text` +- `type`, `entity_type` → `type` +- `id`, `entity_id` → `id` +- `start`, `start_offset` → `start` +- `end`, `end_offset` → `end` + +**Relationships:** +- `source`, `source_id` → `source_id` +- `target`, `target_id` → `target_id` +- `type`, `relationship_type` → `type` + +## Reading Arrow Files + +### With PyArrow + +```python +import pyarrow as pa +import pyarrow.ipc as ipc + +with pa.OSFile("entities.arrow", 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + print(table.schema) + print(table.to_pandas()) +``` + +### With Pandas + +```python +import pandas as pd +import pyarrow.ipc as ipc + +with ipc.open_file("entities.arrow") as reader: + df = reader.read_all().to_pandas() + print(df) +``` + +### With DuckDB + +```python +import duckdb + +# Query Arrow file directly +result = duckdb.query("SELECT * FROM 'entities.arrow' WHERE type = 'Person'") +print(result.df()) +``` + +## Methods + +### `export(data, file_path, schema=None, **options)` + +Generic export method that handles both single and multiple files. + +**Parameters:** +- `data`: List of dicts or dict with list values +- `file_path`: Output file path (base path for dict exports) +- `schema`: Optional Arrow schema (auto-detected if not provided) +- `**options`: Additional options + +### `export_entities(entities, file_path, **options)` + +Export entities to Arrow IPC file with normalization. + +**Parameters:** +- `entities`: List of entity dictionaries +- `file_path`: Output Arrow file path +- `**options`: Additional options + +### `export_relationships(relationships, file_path, **options)` + +Export relationships to Arrow IPC file with normalization. + +**Parameters:** +- `relationships`: List of relationship dictionaries +- `file_path`: Output Arrow file path +- `**options`: Additional options + +### `export_knowledge_graph(knowledge_graph, base_path, **options)` + +Export knowledge graph to multiple Arrow files. + +**Parameters:** +- `knowledge_graph`: Knowledge graph dictionary with 'entities' and 'relationships' +- `base_path`: Base path for output files (without extension) +- `**options`: Additional options + +## Examples + +See `examples/arrow_export_example.py` for comprehensive usage examples. + +## Testing + +Run the test suite: + +```bash +# All Arrow exporter tests +pytest tests/test_arrow_exporter.py -v + +# Integration tests +pytest tests/test_export_module.py::TestExportModule::test_arrow_exporter -v +``` + +## Performance Benefits + +- **Columnar Storage**: Faster analytics on specific columns +- **Compression**: Smaller file sizes (especially with LZ4/ZSTD) +- **Zero-Copy**: Memory-efficient data transfer +- **Cross-Language**: Works with Python, R, Julia, JavaScript, and more +- **SQL Queries**: Direct querying with DuckDB without loading into memory + +## Comparison with Other Formats + +| Feature | Arrow | CSV | JSON | +|---------|-------|-----|------| +| Type Safety | ✓ | ✗ | ✗ | +| Compression | ✓ | ✗ | ✗ | +| Schema Validation | ✓ | ✗ | ✗ | +| Pandas Compatible | ✓ | ✓ | ✓ | +| DuckDB Native | ✓ | ✓ | ✗ | +| Binary Format | ✓ | ✗ | ✗ | +| Human Readable | ✗ | ✓ | ✓ | + +## Architecture + +The Arrow exporter follows Semantica's export architecture: + +1. **Normalization**: Field names are normalized to consistent format +2. **Schema Application**: Explicit schemas ensure type safety +3. **Metadata Conversion**: Dicts converted to Arrow struct fields +4. **Progress Tracking**: Integrated with Semantica's progress tracker +5. **Error Handling**: Structured exceptions with detailed messages + +## Contributing + +When contributing to the Arrow exporter: + +1. Maintain explicit schemas (no inference) +2. Follow existing code style and patterns +3. Add comprehensive tests for new features +4. Update this documentation +5. Ensure Pandas/DuckDB compatibility + +## License + +MIT License - See LICENSE file for details. + +## Author + +Semantica Contributors diff --git a/examples/arrow_export_example.py b/examples/arrow_export_example.py new file mode 100644 index 00000000..b399951c --- /dev/null +++ b/examples/arrow_export_example.py @@ -0,0 +1,150 @@ +""" +Apache Arrow Exporter - Example Usage + +This script demonstrates how to use the ArrowExporter to export +knowledge graphs, entities, and relationships to Apache Arrow format. +""" + +from semantica.export import ArrowExporter, export_arrow +from pathlib import Path +import tempfile + +def main(): + print("=" * 70) + print("Apache Arrow 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 ArrowExporter class + print("Example 1: Export entities to Arrow") + print("-" * 70) + exporter = ArrowExporter() + entities_path = temp_dir / "entities.arrow" + 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 Arrow") + print("-" * 70) + rels_path = temp_dir / "relationships.arrow" + 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 Arrow 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.arrow" + kg_rels = temp_dir / "knowledge_graph_relationships.arrow" + print(f"✓ 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: Use convenience function + print("Example 4: Using export_arrow convenience function") + print("-" * 70) + export_path = temp_dir / "entities_via_function.arrow" + export_arrow(entities, export_path) + print(f"✓ Exported using convenience function: {export_path}") + print(f" File size: {export_path.stat().st_size} bytes\n") + + # Example 5: Read back with PyArrow (if available) + try: + import pyarrow as pa + import pyarrow.ipc as ipc + + print("Example 5: Reading Arrow file with PyArrow") + print("-" * 70) + with pa.OSFile(str(entities_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + print(f"✓ Table schema:") + print(f" {table.schema}") + print(f"\n✓ Table data ({table.num_rows} rows):") + print(f" {table.to_pandas()}\n") + + # Example 6: Convert to Pandas DataFrame + print("Example 6: Convert to Pandas DataFrame") + print("-" * 70) + df = table.to_pandas() + print(f"✓ DataFrame shape: {df.shape}") + print(f"✓ DataFrame columns: {list(df.columns)}") + print(f"\n{df}\n") + + except ImportError: + print("⚠ PyArrow not available for reading examples\n") + + print("=" * 70) + print("✅ All examples completed successfully!") + print("=" * 70) + print(f"\n💡 Tip: Arrow files are columnar and highly compressed,") + print(f" perfect for analytics and compatible with Pandas/DuckDB!\n") + + # Cleanup + import shutil + print(f"🗑 Cleaning up: {temp_dir}") + shutil.rmtree(temp_dir) + print("Done!\n") + +if __name__ == "__main__": + main() diff --git a/semantica/export/__init__.py b/semantica/export/__init__.py index c17c1edf..76c91913 100644 --- a/semantica/export/__init__.py +++ b/semantica/export/__init__.py @@ -120,12 +120,14 @@ Author: Semantica Contributors License: MIT """ +from .arrow_exporter import ArrowExporter from .config import ExportConfig, export_config from .csv_exporter import CSVExporter from .graph_exporter import GraphExporter from .json_exporter import JSONExporter from .lpg_exporter import LPGExporter from .methods import ( + export_arrow, export_csv, export_graph, export_json, @@ -153,6 +155,7 @@ __all__ = [ "NamespaceManager", "JSONExporter", "CSVExporter", + "ArrowExporter", "GraphExporter", "SemanticNetworkYAMLExporter", "YAMLSchemaExporter", @@ -166,6 +169,7 @@ __all__ = [ "export_rdf", "export_json", "export_csv", + "export_arrow", "export_graph", "export_yaml", "export_owl", diff --git a/semantica/export/arrow_exporter.py b/semantica/export/arrow_exporter.py new file mode 100644 index 00000000..92fbdd0c --- /dev/null +++ b/semantica/export/arrow_exporter.py @@ -0,0 +1,610 @@ +""" +Apache Arrow Exporter Module + +This module provides comprehensive Apache Arrow export capabilities for the +Semantica framework, enabling high-performance columnar data export for entities, +relationships, and knowledge graphs. + +Key Features: + - Arrow IPC file export (.arrow) + - Explicit schema definition (no inference) + - Entity and relationship export with metadata + - Knowledge graph export to multiple Arrow files + - Pandas and DuckDB compatible + - Batch export processing + - Structured metadata handling + +Example Usage: + >>> from semantica.export import ArrowExporter + >>> exporter = ArrowExporter() + >>> exporter.export_entities(entities, "entities.arrow") + >>> exporter.export_knowledge_graph(kg, "kg_base") + +Author: Semantica Contributors +License: MIT +""" + +import json +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +try: + import pyarrow as pa + import pyarrow.ipc as ipc + ARROW_AVAILABLE = True +except ImportError: + ARROW_AVAILABLE = False + +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 + + +# Explicit Arrow 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), +]) + + +class ArrowExporter: + """ + Apache Arrow exporter for knowledge graphs and structured data. + + This class provides comprehensive Arrow IPC export functionality for entities, + relationships, and knowledge graphs. Uses explicit schemas for type safety + and compatibility with Pandas and DuckDB. + + Features: + - Entity and relationship export + - Knowledge graph export to multiple Arrow files + - Explicit schema definition (no inference) + - Metadata serialization as Arrow struct fields + - Pandas and DuckDB compatible + - Progress tracking and error handling + + Example Usage: + >>> exporter = ArrowExporter() + >>> exporter.export_entities(entities, "entities.arrow") + >>> exporter.export_knowledge_graph(kg, "output_base") + """ + + def __init__( + self, + compression: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, + **kwargs, + ): + """ + Initialize Arrow exporter. + + Sets up the exporter with specified Arrow formatting options. + + Args: + compression: Compression codec (default: None) + - None: No compression + - "lz4": LZ4 compression + - "zstd": Zstandard compression + config: Optional configuration dictionary (merged with kwargs) + **kwargs: Additional configuration options + + Raises: + ImportError: If pyarrow is not installed + """ + if not ARROW_AVAILABLE: + raise ImportError( + "pyarrow is not installed. Please install it with: " + "pip install pyarrow" + ) + + self.logger = get_logger("arrow_exporter") + self.config = config or {} + self.config.update(kwargs) + + # Arrow configuration + self.compression = compression + + # 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"Arrow 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 Arrow IPC file(s). + + This method handles both single Arrow file export (from list) and multiple + Arrow file export (from dictionary with multiple keys). + + Args: + data: Data to export: + - List of dicts: Exports to single Arrow file + - Dict with list values: Exports each key as separate Arrow file + file_path: Output file path (base path for dict exports) + schema: Arrow schema to use (default: auto-select based on data) + **options: Additional options + + Raises: + ValidationError: If data type is unsupported + + Example: + >>> # Single Arrow file + >>> exporter.export([{"id": "1", "name": "A"}], "data.arrow") + >>> # Multiple Arrow files + >>> exporter.export( + ... {"entities": [...], "relationships": [...]}, + ... "output_base" + ... ) + """ + # Track Arrow export + tracking_id = self.progress_tracker.start_tracking( + file=str(file_path), + module="export", + submodule="ArrowExporter", + message=f"Exporting data to Arrow: {file_path}", + ) + + try: + file_path = Path(file_path) + ensure_directory(file_path.parent) + + self.logger.debug(f"Exporting data to Arrow: {file_path}") + + # Handle different data structures + if isinstance(data, dict): + # Export each key as separate Arrow 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}.arrow" + + # Select schema based on key + key_schema = schema + if key == "entities" and schema is None: + key_schema = ENTITY_SCHEMA + elif key == "relationships" and schema is None: + key_schema = RELATIONSHIP_SCHEMA + + self._write_arrow(value, output_path, schema=key_schema, **options) + exported_files.append(output_path) + else: + self.logger.warning( + f"Skipping key '{key}': value is not a list (type: {type(value)})" + ) + + self.logger.info( + f"Exported {len(exported_files)} Arrow 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)} Arrow files", + ) + elif isinstance(data, list): + # Single Arrow file + self.progress_tracker.update_tracking( + tracking_id, message=f"Exporting {len(data)} records..." + ) + self._write_arrow(data, file_path, schema=schema, **options) + self.logger.info(f"Exported Arrow to: {file_path}") + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Exported Arrow 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 Arrow IPC file. + + This method normalizes entity data to a consistent format and exports + to Arrow using the explicit ENTITY_SCHEMA. Handles various entity field + name variations and serializes metadata as Arrow struct fields. + + Normalized Fields: + - id: Entity identifier (string) + - text: Entity text/label/name (string) + - type: Entity type (string) + - confidence: Confidence score (float64) + - start: Start offset/position (int64) + - end: End offset/position (int64) + - metadata: Metadata as Arrow struct (keys and values lists) + + Args: + entities: List of entity dictionaries with various field names + file_path: Output Arrow file path + **options: Additional options + + 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.arrow") + """ + if not entities: + raise ValidationError("No entities to export. Entities list is empty.") + + self.logger.debug(f"Exporting {len(entities)} entity(ies) to Arrow") + + # 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", "") + text = entity.get("text") or entity.get("label") or entity.get("name") or "" + entity_type = entity.get("type") or entity.get("entity_type", "") + confidence = entity.get("confidence") + start = entity.get("start") or entity.get("start_offset") + end = entity.get("end") or entity.get("end_offset") + + # Convert confidence to float + if confidence is not None: + try: + confidence = float(confidence) + except (TypeError, ValueError): + self.logger.warning( + f"Invalid confidence value for entity {i}: {confidence}. Setting to None." + ) + confidence = None + + # Convert start/end to int + if start is not None: + try: + start = int(start) + except (TypeError, ValueError): + start = None + if end is not None: + try: + end = int(end) + except (TypeError, ValueError): + end = None + + # Convert metadata dict to Arrow struct + metadata = None + if "metadata" in entity and entity["metadata"]: + metadata = self._dict_to_struct(entity["metadata"]) + + normalized = { + "id": str(entity_id), + "text": str(text) if text else None, + "type": str(entity_type) if entity_type else None, + "confidence": confidence, + "start": start, + "end": end, + "metadata": metadata, + } + + normalized_entities.append(normalized) + + self.logger.debug( + f"Normalized {len(normalized_entities)} entity(ies) for Arrow export" + ) + + self._write_arrow(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 Arrow IPC file. + + This method normalizes relationship data to a consistent format and exports + to Arrow using the explicit RELATIONSHIP_SCHEMA. Handles various relationship + field name variations and serializes metadata as Arrow struct fields. + + Normalized Fields: + - id: Relationship identifier (string) + - source_id: Source entity identifier (string) + - target_id: Target entity identifier (string) + - type: Relationship type (string) + - confidence: Confidence score (float64) + - metadata: Metadata as Arrow struct (keys and values lists) + + Args: + relationships: List of relationship dictionaries with various field names + file_path: Output Arrow file path + **options: Additional options + + Raises: + ValidationError: If relationships list is empty + + Example: + >>> relationships = [ + ... {"id": "r1", "source_id": "e1", "target_id": "e2", "type": "RELATED_TO"}, + ... {"source": "e2", "target": "e3", "relationship_type": "CONTAINS"} + ... ] + >>> exporter.export_relationships(relationships, "relationships.arrow") + """ + if not relationships: + raise ValidationError( + "No relationships to export. Relationships list is empty." + ) + + self.logger.debug(f"Exporting {len(relationships)} relationship(s) to Arrow") + + # 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 and normalize fields + rel_id = rel.get("id", f"rel_{i}") + source_id = rel.get("source_id") or rel.get("source", "") + target_id = rel.get("target_id") or rel.get("target", "") + rel_type = rel.get("type") or rel.get("relationship_type", "") + confidence = rel.get("confidence") + + # Convert confidence to float + if confidence is not None: + try: + confidence = float(confidence) + except (TypeError, ValueError): + self.logger.warning( + f"Invalid confidence value for relationship {i}: {confidence}. Setting to None." + ) + confidence = None + + # Convert metadata dict to Arrow struct + metadata = None + if "metadata" in rel and rel["metadata"]: + metadata = self._dict_to_struct(rel["metadata"]) + + normalized = { + "id": str(rel_id), + "source_id": str(source_id), + "target_id": str(target_id), + "type": str(rel_type) if rel_type else None, + "confidence": confidence, + "metadata": metadata, + } + + normalized_rels.append(normalized) + + self.logger.debug( + f"Normalized {len(normalized_rels)} relationship(s) for Arrow export" + ) + + self._write_arrow(normalized_rels, file_path, schema=RELATIONSHIP_SCHEMA, **options) + + def export_knowledge_graph( + self, knowledge_graph: Dict[str, Any], base_path: Union[str, Path], **options + ) -> None: + """ + Export knowledge graph to multiple Arrow IPC files. + + This method exports a knowledge graph to separate Arrow files for entities + and relationships. Each component is exported to its own file with a naming + pattern: `{base_path}_entities.arrow`, etc. + + Exported Files: + - {base_path}_entities.arrow: Entity data + - {base_path}_relationships.arrow: Relationship data + + Args: + knowledge_graph: Knowledge graph dictionary containing: + - entities: List of entity dictionaries + - relationships: List of relationship dictionaries + base_path: Base path for output files (without extension) + **options: Additional options passed to export methods + + Example: + >>> kg = { + ... "entities": [...], + ... "relationships": [...] + ... } + >>> exporter.export_knowledge_graph(kg, "output_base") + >>> # Creates: output_base_entities.arrow, output_base_relationships.arrow + """ + base_path = Path(base_path) + + self.logger.debug( + f"Exporting knowledge graph to Arrow files: base_path={base_path}" + ) + + exported_files = [] + + # Export entities + entities = knowledge_graph.get("entities", []) + if entities: + entities_path = base_path.parent / f"{base_path.stem}_entities.arrow" + self.export_entities(entities, entities_path, **options) + exported_files.append(entities_path) + self.logger.debug( + f"Exported {len(entities)} entity(ies) to {entities_path}" + ) + else: + self.logger.debug("No entities found in knowledge graph") + + # Export relationships + relationships = knowledge_graph.get("relationships", []) + if relationships: + rels_path = base_path.parent / f"{base_path.stem}_relationships.arrow" + self.export_relationships(relationships, rels_path, **options) + exported_files.append(rels_path) + self.logger.debug( + f"Exported {len(relationships)} relationship(s) to {rels_path}" + ) + else: + self.logger.debug("No relationships found in knowledge graph") + + if exported_files: + self.logger.info( + f"Exported knowledge graph to {len(exported_files)} Arrow file(s): " + f"{', '.join(str(f) for f in exported_files)}" + ) + else: + self.logger.warning("No data found in knowledge graph to export") + + def _dict_to_struct(self, metadata: Dict[str, Any]) -> Dict[str, List]: + """ + Convert metadata dictionary to Arrow struct format. + + Converts a dictionary into a struct with keys and values lists, + suitable for Arrow struct fields. + + Args: + metadata: Metadata dictionary to convert + + Returns: + Dict with 'keys' and 'values' lists (both strings) + """ + if not isinstance(metadata, dict): + return None + + keys = [] + values = [] + + for k, v in metadata.items(): + keys.append(str(k)) + # Convert value to JSON string if it's not a simple type + if isinstance(v, (dict, list)): + try: + values.append(json.dumps(v)) + except (TypeError, ValueError): + values.append(str(v)) + else: + values.append(str(v) if v is not None else "") + + return {"keys": keys, "values": values} + + def _write_arrow( + self, + data: List[Dict[str, Any]], + file_path: Path, + schema: Optional[pa.Schema] = None, + **options, + ) -> None: + """ + Write data to Arrow IPC file. + + This internal method handles the actual Arrow file writing, including + schema validation and batch writing. + + Args: + data: List of dictionaries to write as Arrow records + file_path: Output Arrow file path + schema: Arrow schema to use (if None, will try to infer from data structure) + **options: Unused (for compatibility) + + Raises: + ValidationError: If data list is empty or schema cannot be determined + ProcessingError: If file writing fails + """ + if not data: + raise ValidationError("No data to write. Data list is empty.") + + # If schema is not provided, try to infer from data structure + if schema is None: + # Try to detect if data looks like entities or relationships + sample = data[0] if data else {} + + # Check for relationship-specific fields + has_source = any(k in sample for k in ['source_id', 'source']) + has_target = any(k in sample for k in ['target_id', 'target']) + + # Check for entity-specific fields + has_text = any(k in sample for k in ['text', 'label', 'name']) + + if has_source and has_target: + schema = RELATIONSHIP_SCHEMA + self.logger.debug("Auto-detected relationship schema") + elif has_text or 'type' in sample: + schema = ENTITY_SCHEMA + self.logger.debug("Auto-detected entity schema") + else: + raise ValidationError( + "Schema is required for Arrow export. " + "Cannot auto-detect schema from data structure. " + "Provide explicit schema or use export_entities/export_relationships methods." + ) + + self.logger.debug( + f"Writing Arrow IPC file: {len(data)} row(s), " + f"schema={schema}, file={file_path}" + ) + + try: + # Create Arrow table from data using explicit schema + table = pa.Table.from_pylist(data, schema=schema) + + # Write to Arrow IPC file + with pa.OSFile(str(file_path), 'wb') as sink: + with ipc.new_file(sink, schema) as writer: + writer.write_table(table) + + self.logger.debug(f"Successfully wrote Arrow IPC file: {file_path}") + + except pa.ArrowInvalid as e: + error_msg = f"Arrow schema validation failed for {file_path}: {e}" + self.logger.error(error_msg) + raise ValidationError(error_msg) from e + except IOError as e: + error_msg = f"Failed to write Arrow file {file_path}: {e}" + self.logger.error(error_msg) + raise ProcessingError(error_msg) from e + except Exception as e: + error_msg = f"Unexpected error writing Arrow file: {e}" + self.logger.error(error_msg) + raise ProcessingError(error_msg) from e diff --git a/semantica/export/methods.py b/semantica/export/methods.py index 821de83f..a0eb0fc5 100644 --- a/semantica/export/methods.py +++ b/semantica/export/methods.py @@ -155,6 +155,7 @@ from typing import Any, Callable, Dict, List, Optional, Union from ..utils.exceptions import ConfigurationError, ProcessingError from ..utils.logging import get_logger +from .arrow_exporter import ArrowExporter from .config import export_config from .csv_exporter import CSVExporter from .graph_exporter import GraphExporter @@ -317,6 +318,51 @@ def export_csv( raise +def export_arrow( + data: Union[List[Dict[str, Any]], Dict[str, List[Dict[str, Any]]]], + file_path: Union[str, Path], + method: str = "default", + **kwargs, +) -> None: + """ + Export data to Apache Arrow format (convenience function). + + This is a user-friendly wrapper that exports data to Arrow IPC format. + + Args: + data: Data to export (list of dicts or dict with list values) + file_path: Output Arrow file path (or base path for multiple files) + method: Export method (default: "default") + **kwargs: Additional options passed to ArrowExporter + + Examples: + >>> from semantica.export.methods import export_arrow + >>> export_arrow(entities, "entities.arrow") + >>> export_arrow({"entities": [...], "relationships": [...]}, "output_base") + """ + # Check for custom method in registry + custom_method = method_registry.get("arrow", method) + if custom_method and custom_method is not export_arrow: + try: + return custom_method(data, file_path, **kwargs) + except Exception as e: + logger.warning( + f"Custom method {method} failed: {e}, falling back to default" + ) + + try: + # Get config + config = export_config.get_method_config("arrow") + config.update(kwargs) + + exporter = ArrowExporter(**config) + exporter.export(data, file_path, **kwargs) + + except Exception as e: + logger.error(f"Failed to export Arrow: {e}") + raise + + def export_graph( graph_data: Dict[str, Any], file_path: Union[str, Path], diff --git a/tests/test_arrow_exporter.py b/tests/test_arrow_exporter.py new file mode 100644 index 00000000..4f626241 --- /dev/null +++ b/tests/test_arrow_exporter.py @@ -0,0 +1,479 @@ +""" +Unit tests for Apache Arrow exporter module. + +Tests schema validation, data export, Pandas conversion, empty inputs, +and minimal graph structures. +""" + +import os +import unittest +import shutil +import tempfile +from pathlib import Path +from unittest.mock import MagicMock, patch + +# Try to import pyarrow +try: + import pyarrow as pa + import pyarrow.ipc as ipc + ARROW_AVAILABLE = True +except ImportError: + ARROW_AVAILABLE = False + +from semantica.export import ArrowExporter +from semantica.utils.exceptions import ValidationError, ProcessingError + + +@unittest.skipIf(not ARROW_AVAILABLE, "pyarrow not installed") +class TestArrowExporter(unittest.TestCase): + """Test cases for ArrowExporter 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 ArrowExporter initialization.""" + exporter = ArrowExporter() + self.assertIsNotNone(exporter) + self.assertEqual(exporter.compression, None) + + # Test with compression + exporter_compressed = ArrowExporter(compression="lz4") + self.assertEqual(exporter_compressed.compression, "lz4") + + def test_initialization_without_pyarrow(self): + """Test initialization fails gracefully without pyarrow.""" + with patch.dict('sys.modules', {'pyarrow': None}): + # This test would need to reload the module + # For now, we just verify the constant + if not ARROW_AVAILABLE: + with self.assertRaises(ImportError): + from semantica.export.arrow_exporter import ArrowExporter as TestExporter + + def test_export_entities_basic(self): + """Test basic entity export to Arrow.""" + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "entities.arrow" + + exporter.export_entities(self.entities, str(output_path)) + self.assertTrue(output_path.exists()) + + # Read and verify Arrow file + with pa.OSFile(str(output_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + + # Verify schema + from semantica.export.arrow_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 = ArrowExporter() + output_path = Path(self.test_dir) / "entities_normalized.arrow" + + # 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 + with pa.OSFile(str(output_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + + self.assertEqual(table.num_rows, 3) + # All entities should have normalized fields + ids = table.column('id').to_pylist() + self.assertEqual(ids, ['e1', 'e2', 'e3']) + + texts = table.column('text').to_pylist() + self.assertEqual(texts, ['Entity 1', 'Entity 2', 'Entity 3']) + + def test_export_relationships_basic(self): + """Test basic relationship export to Arrow.""" + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "relationships.arrow" + + exporter.export_relationships(self.relationships, str(output_path)) + self.assertTrue(output_path.exists()) + + # Read and verify Arrow file + with pa.OSFile(str(output_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + + # 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_knowledge_graph(self): + """Test knowledge graph export to multiple Arrow files.""" + exporter = ArrowExporter() + base_path = Path(self.test_dir) / "kg_output" + + exporter.export_knowledge_graph(self.kg, str(base_path)) + + # Verify files created + entities_path = Path(self.test_dir) / "kg_output_entities.arrow" + relationships_path = Path(self.test_dir) / "kg_output_relationships.arrow" + + self.assertTrue(entities_path.exists()) + self.assertTrue(relationships_path.exists()) + + # Verify entities file + with pa.OSFile(str(entities_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + self.assertEqual(table.num_rows, 2) + + # Verify relationships file + with pa.OSFile(str(relationships_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + self.assertEqual(table.num_rows, 2) + + def test_export_empty_entities(self): + """Test export with empty entity list.""" + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "empty_entities.arrow" + + with self.assertRaises(ValidationError) as context: + exporter.export_entities([], str(output_path)) + + self.assertIn("No entities to export", str(context.exception)) + + def test_export_empty_relationships(self): + """Test export with empty relationship list.""" + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "empty_rels.arrow" + + with self.assertRaises(ValidationError) as context: + exporter.export_relationships([], str(output_path)) + + self.assertIn("No relationships to export", str(context.exception)) + + def test_export_minimal_graph(self): + """Test export with minimal knowledge graph.""" + exporter = ArrowExporter() + base_path = Path(self.test_dir) / "minimal_kg" + + # Minimal graph with only entities + minimal_kg = { + "entities": [{"id": "e1", "text": "Entity 1", "type": "TYPE1"}], + "relationships": [] + } + + exporter.export_knowledge_graph(minimal_kg, str(base_path)) + + # Only entities file should be created + entities_path = Path(self.test_dir) / "minimal_kg_entities.arrow" + self.assertTrue(entities_path.exists()) + + def test_metadata_serialization(self): + """Test metadata dictionary to Arrow struct conversion.""" + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "entities_metadata.arrow" + + entities_with_metadata = [ + { + "id": "e1", + "text": "Entity 1", + "type": "TYPE1", + "metadata": { + "key1": "value1", + "key2": 123, + "key3": {"nested": "dict"} + } + } + ] + + exporter.export_entities(entities_with_metadata, str(output_path)) + self.assertTrue(output_path.exists()) + + # Read and verify metadata structure + with pa.OSFile(str(output_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + + # Verify metadata exists and is a struct + metadata_col = table.column('metadata')[0] + if metadata_col.is_valid: + metadata = metadata_col.as_py() + self.assertIn('keys', metadata) + self.assertIn('values', metadata) + self.assertIsInstance(metadata['keys'], list) + self.assertIsInstance(metadata['values'], list) + + def test_schema_validation(self): + """Test explicit schema validation.""" + from semantica.export.arrow_exporter import ENTITY_SCHEMA, RELATIONSHIP_SCHEMA + + # Verify entity schema structure + self.assertIsNotNone(ENTITY_SCHEMA) + self.assertIsInstance(ENTITY_SCHEMA, pa.Schema) + + entity_fields = {field.name for field in ENTITY_SCHEMA} + expected_entity_fields = {'id', 'text', 'type', 'confidence', 'start', 'end', 'metadata'} + self.assertEqual(entity_fields, expected_entity_fields) + + # Verify relationship schema structure + self.assertIsNotNone(RELATIONSHIP_SCHEMA) + self.assertIsInstance(RELATIONSHIP_SCHEMA, pa.Schema) + + rel_fields = {field.name for field in RELATIONSHIP_SCHEMA} + expected_rel_fields = {'id', 'source_id', 'target_id', 'type', 'confidence', 'metadata'} + self.assertEqual(rel_fields, expected_rel_fields) + + def test_pandas_conversion(self): + """Test conversion to Pandas DataFrame.""" + try: + import pandas as pd + + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "entities_pandas.arrow" + + exporter.export_entities(self.entities, str(output_path)) + + # Read Arrow file and convert to Pandas + with pa.OSFile(str(output_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + df = table.to_pandas() + + # Verify DataFrame + self.assertIsInstance(df, pd.DataFrame) + self.assertEqual(len(df), 2) + self.assertEqual(df['id'].tolist(), ['e1', 'e2']) + + except ImportError: + self.skipTest("Pandas not installed") + + def test_invalid_data_type(self): + """Test export with invalid data type.""" + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "invalid.arrow" + + # Test with string instead of list/dict + with self.assertRaises(ValidationError): + exporter.export("invalid_data", str(output_path)) + + def test_confidence_conversion(self): + """Test confidence value conversion to float.""" + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "entities_confidence.arrow" + + entities_varied_confidence = [ + {"id": "e1", "text": "Entity 1", "confidence": "0.95"}, # String + {"id": "e2", "text": "Entity 2", "confidence": 0.88}, # Float + {"id": "e3", "text": "Entity 3", "confidence": None}, # None + ] + + exporter.export_entities(entities_varied_confidence, str(output_path)) + + # Read and verify confidence values + with pa.OSFile(str(output_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + + confidences = table.column('confidence').to_pylist() + # First should be converted to float, third should be None + self.assertIsInstance(confidences[0], float) + self.assertEqual(confidences[1], 0.88) + self.assertIsNone(confidences[2]) + + def test_offset_conversion(self): + """Test start/end offset conversion to int.""" + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "entities_offsets.arrow" + + entities_varied_offsets = [ + {"id": "e1", "text": "Entity 1", "start": "10", "end": "20"}, # Strings + {"id": "e2", "text": "Entity 2", "start": 30, "end": 40}, # Ints + {"id": "e3", "text": "Entity 3", "start": None, "end": None}, # None + ] + + exporter.export_entities(entities_varied_offsets, str(output_path)) + + # Read and verify offset values + with pa.OSFile(str(output_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + + starts = table.column('start').to_pylist() + ends = table.column('end').to_pylist() + + # First should be converted to int + self.assertEqual(starts[0], 10) + self.assertEqual(ends[0], 20) + self.assertEqual(starts[1], 30) + self.assertEqual(ends[1], 40) + self.assertIsNone(starts[2]) + self.assertIsNone(ends[2]) + + def test_dict_export(self): + """Test export with dictionary containing multiple lists.""" + exporter = ArrowExporter() + base_path = Path(self.test_dir) / "dict_export" + + data_dict = { + "entities": self.entities, + "relationships": self.relationships + } + + exporter.export(data_dict, str(base_path)) + + # Verify both files created + entities_path = Path(self.test_dir) / "dict_export_entities.arrow" + relationships_path = Path(self.test_dir) / "dict_export_relationships.arrow" + + self.assertTrue(entities_path.exists()) + self.assertTrue(relationships_path.exists()) + + def test_progress_tracking(self): + """Test progress tracker integration.""" + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "entities_progress.arrow" + + # Progress tracker should be enabled + self.assertTrue(exporter.progress_tracker.enabled) + + # Export should complete successfully + exporter.export_entities(self.entities, str(output_path)) + self.assertTrue(output_path.exists()) + + def test_non_dict_entity_skipping(self): + """Test skipping non-dictionary entities.""" + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "entities_mixed.arrow" + + mixed_entities = [ + {"id": "e1", "text": "Entity 1", "type": "TYPE1"}, + "invalid_entity", # Should be skipped + {"id": "e2", "text": "Entity 2", "type": "TYPE2"}, + ] + + exporter.export_entities(mixed_entities, str(output_path)) + + # Read and verify only valid entities exported + with pa.OSFile(str(output_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + self.assertEqual(table.num_rows, 2) + + def test_non_dict_relationship_skipping(self): + """Test skipping non-dictionary relationships.""" + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "rels_mixed.arrow" + + mixed_rels = [ + {"id": "r1", "source_id": "e1", "target_id": "e2", "type": "TYPE1"}, + None, # Should be skipped + {"id": "r2", "source_id": "e2", "target_id": "e3", "type": "TYPE2"}, + ] + + exporter.export_relationships(mixed_rels, str(output_path)) + + # Read and verify only valid relationships exported + with pa.OSFile(str(output_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + self.assertEqual(table.num_rows, 2) + + +class TestArrowExporterIntegration(unittest.TestCase): + """Integration tests for Arrow exporter with export methods.""" + + def setUp(self): + """Set up test fixtures.""" + self.test_dir = tempfile.mkdtemp() + self.entities = [ + {"id": "e1", "text": "Entity 1", "type": "PERSON"}, + {"id": "e2", "text": "Entity 2", "type": "ORG"} + ] + + def tearDown(self): + """Clean up test directory.""" + shutil.rmtree(self.test_dir) + + @unittest.skipIf(not ARROW_AVAILABLE, "pyarrow not installed") + def test_export_arrow_method(self): + """Test export_arrow convenience function.""" + from semantica.export.methods import export_arrow + + output_path = Path(self.test_dir) / "entities_method.arrow" + export_arrow(self.entities, str(output_path)) + + self.assertTrue(output_path.exists()) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_export_module.py b/tests/test_export_module.py index 6b53821e..b63743eb 100644 --- a/tests/test_export_module.py +++ b/tests/test_export_module.py @@ -8,6 +8,7 @@ from pathlib import Path from unittest.mock import MagicMock, patch from semantica.export import ( + ArrowExporter, JSONExporter, CSVExporter, RDFExporter, @@ -255,6 +256,41 @@ class TestExportModule(unittest.TestCase): generator.generate_report(self.kg, str(output_path), format="html") self.assertTrue(output_path.exists()) + def test_arrow_exporter(self): + """Test Arrow exporter basic functionality.""" + try: + import pyarrow as pa + import pyarrow.ipc as ipc + + exporter = ArrowExporter() + output_path = Path(self.test_dir) / "entities.arrow" + + # Test export_entities + exporter.export_entities(self.entities, str(output_path)) + self.assertTrue(output_path.exists()) + + # Verify Arrow file + with pa.OSFile(str(output_path), 'rb') as source: + with ipc.open_file(source) as reader: + table = reader.read_all() + self.assertEqual(table.num_rows, 2) + self.assertIn('id', table.column_names) + self.assertIn('text', table.column_names) + self.assertIn('type', table.column_names) + + # Test export_knowledge_graph + base_path = Path(self.test_dir) / "kg_arrow" + exporter.export_knowledge_graph(self.kg, str(base_path)) + + entities_path = Path(self.test_dir) / "kg_arrow_entities.arrow" + rels_path = Path(self.test_dir) / "kg_arrow_relationships.arrow" + + self.assertTrue(entities_path.exists()) + self.assertTrue(rels_path.exists()) + + except ImportError: + print("Skipping Arrow test due to missing pyarrow") + def test_registry(self): def dummy_method(data, path, **kwargs): with open(path, 'w') as f: