mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-08-29 04:26:20 +00:00
feat: add Apache Arrow exporter
This commit is contained in:
@@ -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
|
||||
@@ -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()
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
@@ -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],
|
||||
|
||||
@@ -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()
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user