mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-08-29 04:26:20 +00:00
* feat: implement Apache Arrow and Feather file ingestion support (#235) * fix(arrow): eliminate double full-scan and clean up reader wrapper - Replace _read_batches with _read_batches_with_info which collects batch metadata (total_rows, record_batches) during the same pass as the data read, so ingest_file no longer calls _file_metadata before _read_batches. For a limit=1 read on a large file this previously scanned every batch twice; now it stops after the first batch. - _file_metadata is now only invoked for include_data=False (where a full scan is unavoidable to report accurate row counts). - Remove the dead num_record_batches property from _ArrowReaderWrapper; it was never called by production code and its is_table branch materialised all batches just to count them. - Fix _open_file exception chain: raise ... from file_err instead of from feather_err so the most diagnostic IPC error appears in the Python traceback chain, not the least informative fallback error. * docs(changelog): add [Unreleased] entries for Arrow ingestion (#705) --------- Co-authored-by: KaifAhmad1 <kaifahmad087@gmail.com>
This commit is contained in:
@@ -11,6 +11,30 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
### Added
|
||||
|
||||
- **Apache Arrow & Feather File Ingestion** (#705) by @Luffy2208
|
||||
- Added `ArrowIngestor` (`semantica/ingest/arrow_ingestor.py`) for reading `.arrow`, `.feather`, and `.ipc` files via PyArrow
|
||||
- Supports Arrow IPC File format (random-access), Arrow IPC Stream format, Feather v1 and v2
|
||||
- Selective column reads, optional row limits, and batch-aware iteration that stops early without scanning the full file
|
||||
- `extract_schema()` and `extract_metadata()` convenience methods for schema/metadata inspection without reading row data
|
||||
- `_ArrowReaderWrapper` provides a unified interface across all three reader types, preventing stream exhaustion during schema inspection
|
||||
- `ingest_arrow()` convenience function and `ingest(..., source_type="arrow")` unified dispatch
|
||||
- Automatic Arrow format detection in `ingest()` by file extension (`.arrow`, `.feather`, `.ipc`) and by Arrow IPC magic bytes (`ARROW1\x00\x00`) in `FileTypeDetector`
|
||||
- Registry integration under the `arrow` task namespace with `file`, `schema`, and `metadata` methods
|
||||
- Lazy-import exports of `ArrowIngestor`, `ArrowData`, and `ingest_arrow` from `semantica.ingest`
|
||||
- Optional dependency group: `pip install semantica[ingest-arrow]`; included in `pip install semantica[all]`
|
||||
- 34 tests covering schema extraction, metadata inspection, row limits, column selection, multi-batch reading, IPC stream format, Feather ingestion, empty datasets, null values, magic-byte detection, and failure modes
|
||||
|
||||
### Fixed
|
||||
|
||||
- **Arrow ingestion double full-scan on every data read** (#705) by @KaifAhmad1
|
||||
- `ingest_file` previously called `_file_metadata` (a full batch scan) before `_read_batches`, meaning every read scanned the entire file twice; for a `limit=1` read on a large file the metadata pass visited every batch while the data pass read only one; replaced with a single-pass `_read_batches_with_info` that collects batch metadata as a side effect of the data read; `_file_metadata` is now only invoked for `include_data=False`
|
||||
|
||||
- **Dead `num_record_batches` property on `_ArrowReaderWrapper` materialised all table batches** (#705) by @KaifAhmad1
|
||||
- The property was never called by production code but its `is_table` branch called `to_batches()` purely to take `len()`, materialising the entire table in memory just for a count; property removed
|
||||
|
||||
- **Arrow `_open_file` chained the wrong exception** (#705) by @KaifAhmad1
|
||||
- The fallback cascade (IPC file → IPC stream → Feather) raised `from feather_err`, surfacing the least diagnostic error in the Python traceback chain; changed to `from file_err` so the IPC file open error — the most informative signal for unrecognised formats — appears as `__cause__`
|
||||
|
||||
- **Neo4j Bulk CSV Export** (#665) by @Luffy2208
|
||||
- Added `Neo4jCSVExporter` for generating Neo4j bulk-import CSV files compatible with `neo4j-admin database import`
|
||||
- Produces deterministic `nodes.csv` and `relationships.csv` with stable node IDs — reuses existing graph IDs or derives reproducible SHA-256 content-based IDs when none are present
|
||||
|
||||
+3
-2
@@ -120,6 +120,7 @@ parse-docling = ["docling>=1.0.0"]
|
||||
db-snowflake = ["snowflake-connector-python>=4.5.0", "cryptography>=49.0.0"]
|
||||
db-arrow = ["pyarrow>=21.0.0"]
|
||||
ingest-parquet = ["pyarrow>=21.0.0"]
|
||||
ingest-arrow = ["pyarrow>=21.0.0"]
|
||||
|
||||
db-all = [
|
||||
"semantica[db-snowflake,db-arrow]"
|
||||
@@ -232,8 +233,8 @@ explorer-lite = [
|
||||
|
||||
# Everything (cross-platform — gpu excluded; install semantica[gpu] separately on Linux)
|
||||
all = [
|
||||
"semantica[dev,viz,infra,cloud,monitoring,watch,llm-all,models-huggingface,split-all,graph-all,vectorstore-all,parse-docling,ingest-parquet,explorer]",
|
||||
"semantica[dev,viz,infra,cloud,monitoring,watch,llm-all,models-huggingface,split-all,graph-all,vectorstore-all,parse-docling,ingest-parquet,agno]"
|
||||
"semantica[dev,viz,infra,cloud,monitoring,watch,llm-all,models-huggingface,split-all,graph-all,vectorstore-all,parse-docling,ingest-parquet,ingest-arrow,explorer]",
|
||||
"semantica[dev,viz,infra,cloud,monitoring,watch,llm-all,models-huggingface,split-all,graph-all,vectorstore-all,parse-docling,ingest-parquet,ingest-arrow,agno]"
|
||||
]
|
||||
|
||||
# ---------------- ENTRYPOINTS ----------------
|
||||
|
||||
@@ -96,6 +96,7 @@ Main Classes:
|
||||
- DBIngestor: Database export handling
|
||||
- OntologyIngestor: Ontology file processing
|
||||
- ParquetIngestor: Apache Parquet file and partitioned dataset processing
|
||||
- ArrowIngestor: Apache Arrow IPC and Feather file processing
|
||||
- XMLIngestor: XML file parsing, validation, and metadata extraction
|
||||
- MethodRegistry: Registry for custom ingestion methods
|
||||
- IngestConfig: Configuration manager for ingest module
|
||||
@@ -112,6 +113,7 @@ Convenience Functions:
|
||||
- ingest_database: Database ingestion wrapper
|
||||
- ingest_ontology: Ontology ingestion wrapper
|
||||
- ingest_parquet: Parquet ingestion wrapper
|
||||
- ingest_arrow: Arrow IPC/Feather ingestion wrapper
|
||||
- ingest_xml: XML ingestion wrapper
|
||||
|
||||
|
||||
@@ -140,6 +142,7 @@ from .file_ingestor import (
|
||||
from .methods import (
|
||||
get_ingest_method,
|
||||
ingest,
|
||||
ingest_arrow,
|
||||
ingest_database,
|
||||
ingest_email,
|
||||
ingest_feed,
|
||||
@@ -218,6 +221,9 @@ _LAZY_EXPORTS: Dict[str, Tuple[str, str]] = {
|
||||
# Parquet ingestion
|
||||
"ParquetIngestor": (".parquet_ingestor", "ParquetIngestor"),
|
||||
"ParquetData": (".parquet_ingestor", "ParquetData"),
|
||||
# Arrow IPC / Feather ingestion
|
||||
"ArrowIngestor": (".arrow_ingestor", "ArrowIngestor"),
|
||||
"ArrowData": (".arrow_ingestor", "ArrowData"),
|
||||
# XML ingestion
|
||||
"XMLIngestor": (".xml_ingestor", "XMLIngestor"),
|
||||
"XMLIngestionData": (".xml_ingestor", "XMLIngestionData"),
|
||||
@@ -244,6 +250,10 @@ _OPTIONAL_DEPENDENCY_MESSAGES = {
|
||||
"Parquet ingestion requires optional dependency 'pyarrow'. "
|
||||
"Install it before importing ParquetIngestor or using ingest_parquet()."
|
||||
),
|
||||
".arrow_ingestor": (
|
||||
"Arrow ingestion requires optional dependency 'pyarrow'. "
|
||||
"Install it before importing ArrowIngestor or using ingest_arrow()."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@@ -334,6 +344,9 @@ __all__ = [
|
||||
# Parquet ingestion
|
||||
"ParquetIngestor",
|
||||
"ParquetData",
|
||||
# Arrow IPC / Feather ingestion
|
||||
"ArrowIngestor",
|
||||
"ArrowData",
|
||||
# XML ingestion
|
||||
"XMLIngestor",
|
||||
"XMLIngestionData",
|
||||
@@ -349,6 +362,7 @@ __all__ = [
|
||||
"ingest_email",
|
||||
"ingest_database",
|
||||
"ingest_ontology",
|
||||
"ingest_arrow",
|
||||
"ingest_parquet",
|
||||
"ingest_public_api",
|
||||
"ingest_xml",
|
||||
|
||||
@@ -0,0 +1,506 @@
|
||||
"""
|
||||
Apache Arrow IPC Ingestion Module
|
||||
|
||||
This module provides dedicated Arrow IPC ingestion for local ``.arrow``,
|
||||
``.feather``, and ``.ipc`` files. It uses PyArrow when available so callers
|
||||
can read selected columns, inspect schemas and file metadata, and process
|
||||
record batches without unnecessary full-table copies.
|
||||
|
||||
Supported formats:
|
||||
- Arrow IPC File (``*.arrow``, ``*.ipc``)
|
||||
- Feather v1/v2 (``*.feather``) — Feather v2 is the Arrow IPC format
|
||||
|
||||
Example Usage:
|
||||
>>> from semantica.ingest import ArrowIngestor
|
||||
>>> ingestor = ArrowIngestor()
|
||||
>>> data = ingestor.ingest_file("events.arrow", columns=["id", "event_type"])
|
||||
>>> schema = ingestor.extract_schema("events.arrow")
|
||||
>>> metadata = ingestor.extract_metadata("events.feather")
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Sequence, Union
|
||||
|
||||
try:
|
||||
import pyarrow as pa
|
||||
import pyarrow.feather as pf
|
||||
import pyarrow.ipc as ipc
|
||||
|
||||
ARROW_AVAILABLE = True
|
||||
except (ImportError, OSError):
|
||||
pa = None # type: ignore[assignment]
|
||||
pf = None # type: ignore[assignment]
|
||||
ipc = None # type: ignore[assignment]
|
||||
ARROW_AVAILABLE = False
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
from ..utils.progress_tracker import get_progress_tracker
|
||||
|
||||
# Extensions recognised as Arrow IPC / Feather files.
|
||||
_ARROW_EXTENSIONS = {".arrow", ".feather", ".ipc"}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ArrowData:
|
||||
"""Arrow IPC ingestion result."""
|
||||
|
||||
data: List[Dict[str, Any]]
|
||||
row_count: int
|
||||
columns: List[str]
|
||||
schema: Dict[str, Any]
|
||||
source: str
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
ingested_at: datetime = field(default_factory=datetime.now)
|
||||
|
||||
|
||||
class _ArrowReaderWrapper:
|
||||
"""Unified wrapper for Arrow File, Stream, and Feather readers."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
reader_or_table: Any,
|
||||
is_stream: bool = False,
|
||||
is_table: bool = False,
|
||||
):
|
||||
self.obj = reader_or_table
|
||||
self.is_stream = is_stream
|
||||
self.is_table = is_table
|
||||
|
||||
@property
|
||||
def schema(self) -> Any:
|
||||
return self.obj.schema
|
||||
|
||||
def iter_batches(self) -> Any:
|
||||
if self.is_table:
|
||||
for batch in self.obj.to_batches():
|
||||
yield batch
|
||||
elif self.is_stream:
|
||||
for batch in self.obj:
|
||||
yield batch
|
||||
else:
|
||||
for i in range(self.obj.num_record_batches):
|
||||
yield self.obj.get_batch(i)
|
||||
|
||||
|
||||
class ArrowIngestor:
|
||||
"""
|
||||
Dedicated Arrow IPC / Feather ingestion handler.
|
||||
|
||||
Features:
|
||||
- Single Arrow IPC or Feather file ingestion
|
||||
- Selective column reads
|
||||
- Schema and file metadata extraction
|
||||
- Batch-aware reading with optional row limits
|
||||
- Memory-efficient iteration over record batches
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs):
|
||||
"""
|
||||
Initialize Arrow ingestor.
|
||||
|
||||
Args:
|
||||
config: Optional configuration dictionary
|
||||
**kwargs: Additional configuration options
|
||||
|
||||
Raises:
|
||||
ImportError: If pyarrow is not installed
|
||||
"""
|
||||
if not ARROW_AVAILABLE:
|
||||
raise ImportError(
|
||||
"pyarrow is required for ArrowIngestor. "
|
||||
"Install it with: pip install pyarrow"
|
||||
)
|
||||
|
||||
self.logger = get_logger("arrow_ingestor")
|
||||
self.config = config or {}
|
||||
self.config.update(kwargs)
|
||||
self.progress_tracker = get_progress_tracker()
|
||||
if not self.progress_tracker.enabled:
|
||||
self.progress_tracker.enabled = True
|
||||
|
||||
self.logger.debug("Arrow ingestor initialized")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def ingest(
|
||||
self,
|
||||
source: Union[str, Path],
|
||||
columns: Optional[Union[str, Sequence[str]]] = None,
|
||||
limit: Optional[int] = None,
|
||||
include_data: bool = True,
|
||||
**options,
|
||||
) -> ArrowData:
|
||||
"""
|
||||
Ingest an Arrow IPC or Feather file.
|
||||
|
||||
Args:
|
||||
source: Arrow / Feather file path
|
||||
columns: Optional column name or names to read
|
||||
limit: Optional maximum number of rows to return
|
||||
include_data: If False, return schema and metadata without rows
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
ArrowData: Ingested data and metadata
|
||||
"""
|
||||
return self.ingest_file(
|
||||
source,
|
||||
columns=columns,
|
||||
limit=limit,
|
||||
include_data=include_data,
|
||||
**options,
|
||||
)
|
||||
|
||||
def ingest_file(
|
||||
self,
|
||||
file_path: Union[str, Path],
|
||||
columns: Optional[Union[str, Sequence[str]]] = None,
|
||||
limit: Optional[int] = None,
|
||||
include_data: bool = True,
|
||||
**options,
|
||||
) -> ArrowData:
|
||||
"""
|
||||
Ingest a single Arrow IPC or Feather file.
|
||||
|
||||
Args:
|
||||
file_path: Path to Arrow / Feather file
|
||||
columns: Optional column name or names to read
|
||||
limit: Optional maximum number of rows to return
|
||||
include_data: If False, skip reading row data
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
ArrowData: Ingested data, schema, and metadata
|
||||
"""
|
||||
file_path = Path(file_path)
|
||||
self._validate_file(file_path)
|
||||
|
||||
if limit is not None and limit < 0:
|
||||
raise ValidationError("limit must be greater than or equal to 0")
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
file=str(file_path),
|
||||
module="ingest",
|
||||
submodule="ArrowIngestor",
|
||||
message=f"Ingesting Arrow: {file_path.name}",
|
||||
)
|
||||
|
||||
try:
|
||||
reader = self._open_file(file_path)
|
||||
file_schema = reader.schema
|
||||
|
||||
selected_columns = self._normalize_columns(
|
||||
columns,
|
||||
[field.name for field in file_schema],
|
||||
)
|
||||
|
||||
if include_data:
|
||||
table, batch_info = self._read_batches_with_info(
|
||||
reader, selected_columns, limit
|
||||
)
|
||||
data = table.to_pylist()
|
||||
schema = self._schema_to_dict(table.schema)
|
||||
result_columns = list(table.column_names)
|
||||
metadata = {
|
||||
"format": "arrow",
|
||||
"source_type": "file",
|
||||
"file": str(file_path),
|
||||
"file_size": file_path.stat().st_size,
|
||||
"total_rows": sum(b["num_rows"] for b in batch_info),
|
||||
"num_record_batches": len(batch_info),
|
||||
"num_columns": len(file_schema),
|
||||
"record_batches": batch_info,
|
||||
"schema_metadata": self._decode_metadata_map(file_schema.metadata),
|
||||
"returned_rows": len(data),
|
||||
"selected_columns": result_columns,
|
||||
"limit": limit,
|
||||
"include_data": include_data,
|
||||
}
|
||||
else:
|
||||
metadata = self._file_metadata(file_path, reader)
|
||||
selected_schema = self._select_schema(file_schema, selected_columns)
|
||||
data = []
|
||||
schema = self._schema_to_dict(selected_schema)
|
||||
result_columns = [f.name for f in selected_schema]
|
||||
metadata.update(
|
||||
{
|
||||
"returned_rows": 0,
|
||||
"selected_columns": result_columns,
|
||||
"limit": limit,
|
||||
"include_data": include_data,
|
||||
}
|
||||
)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Ingested Arrow: {len(data)} rows",
|
||||
)
|
||||
|
||||
self.logger.info(
|
||||
f"Arrow ingestion completed: {len(data)} row(s) from {file_path}"
|
||||
)
|
||||
|
||||
return ArrowData(
|
||||
data=data,
|
||||
row_count=len(data),
|
||||
columns=result_columns,
|
||||
schema=schema,
|
||||
source=str(file_path),
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
except (ValidationError, ProcessingError):
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message="Arrow ingestion failed"
|
||||
)
|
||||
raise
|
||||
except Exception as e:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
self.logger.error(f"Failed to ingest Arrow {file_path}: {e}")
|
||||
raise ProcessingError(f"Failed to ingest Arrow file: {e}") from e
|
||||
|
||||
def extract_schema(self, source: Union[str, Path], **options) -> Dict[str, Any]:
|
||||
"""
|
||||
Extract schema from an Arrow IPC or Feather file.
|
||||
|
||||
Args:
|
||||
source: Arrow / Feather file path
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
dict: Schema with column names, types, nullability, and metadata
|
||||
"""
|
||||
source_path = Path(source)
|
||||
self._validate_file(source_path)
|
||||
reader = self._open_file(source_path)
|
||||
return self._schema_to_dict(reader.schema)
|
||||
|
||||
def extract_metadata(self, source: Union[str, Path], **options) -> Dict[str, Any]:
|
||||
"""
|
||||
Extract Arrow file metadata without reading row data.
|
||||
|
||||
Args:
|
||||
source: Arrow / Feather file path
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
dict: Row counts, record batches, column info, and file metadata
|
||||
"""
|
||||
source_path = Path(source)
|
||||
self._validate_file(source_path)
|
||||
reader = self._open_file(source_path)
|
||||
return self._file_metadata(source_path, reader)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _open_file(self, file_path: Path) -> Any:
|
||||
"""Open an Arrow IPC file, stream, or Feather file.
|
||||
|
||||
Tries opening as a File first, then Stream, then Feather.
|
||||
"""
|
||||
try:
|
||||
# Try opening as Arrow File (Random Access format)
|
||||
reader = ipc.open_file(str(file_path))
|
||||
return _ArrowReaderWrapper(reader, is_stream=False)
|
||||
except Exception as file_err:
|
||||
try:
|
||||
# Try opening as Arrow Stream
|
||||
reader = ipc.open_stream(str(file_path))
|
||||
return _ArrowReaderWrapper(reader, is_stream=True)
|
||||
except Exception as stream_err:
|
||||
try:
|
||||
# Try opening as Feather file
|
||||
table = pf.read_table(str(file_path))
|
||||
return _ArrowReaderWrapper(table, is_table=True)
|
||||
except Exception as feather_err:
|
||||
raise ProcessingError(
|
||||
f"Failed to open Arrow file {file_path}: "
|
||||
f"IPC file error: {file_err}. "
|
||||
f"IPC stream error: {stream_err}. "
|
||||
f"Feather error: {feather_err}."
|
||||
) from file_err
|
||||
|
||||
def _read_batches_with_info(
|
||||
self,
|
||||
reader: Any,
|
||||
columns: Optional[List[str]],
|
||||
limit: Optional[int],
|
||||
) -> tuple:
|
||||
"""Single-pass read: returns (Table, batch_info_list).
|
||||
|
||||
Collecting batch metadata here avoids a separate full-file scan before
|
||||
the data read, which would double I/O for large files with small limits.
|
||||
"""
|
||||
batch_info: List[Dict[str, Any]] = []
|
||||
|
||||
if limit == 0:
|
||||
return (
|
||||
pa.Table.from_batches(
|
||||
[],
|
||||
schema=self._select_schema(reader.schema, columns),
|
||||
),
|
||||
batch_info,
|
||||
)
|
||||
|
||||
batches: list = []
|
||||
remaining = limit # None means "all rows"
|
||||
|
||||
for i, batch in enumerate(reader.iter_batches()):
|
||||
if columns is not None:
|
||||
batch = pa.RecordBatch.from_arrays(
|
||||
[batch.column(c) for c in columns],
|
||||
names=columns,
|
||||
)
|
||||
|
||||
if remaining is not None:
|
||||
if batch.num_rows > remaining:
|
||||
batch = batch.slice(0, remaining)
|
||||
remaining -= batch.num_rows
|
||||
|
||||
batch_info.append(
|
||||
{
|
||||
"batch_index": i,
|
||||
"num_rows": batch.num_rows,
|
||||
"num_columns": batch.num_columns,
|
||||
}
|
||||
)
|
||||
batches.append(batch)
|
||||
|
||||
if remaining is not None and remaining <= 0:
|
||||
break
|
||||
|
||||
return (
|
||||
pa.Table.from_batches(
|
||||
batches,
|
||||
schema=self._select_schema(reader.schema, columns),
|
||||
),
|
||||
batch_info,
|
||||
)
|
||||
|
||||
def _validate_file(self, file_path: Path) -> None:
|
||||
"""Validate a local Arrow / Feather file path."""
|
||||
if not file_path.exists():
|
||||
raise ValidationError(f"Arrow file not found: {file_path}")
|
||||
if not file_path.is_file():
|
||||
raise ValidationError(f"Path is not a file: {file_path}")
|
||||
if file_path.suffix.lower() not in _ARROW_EXTENSIONS:
|
||||
raise ValidationError(
|
||||
f"File is not an Arrow/Feather file: {file_path}. "
|
||||
f"Supported extensions: {', '.join(sorted(_ARROW_EXTENSIONS))}"
|
||||
)
|
||||
|
||||
def _normalize_columns(
|
||||
self,
|
||||
columns: Optional[Union[str, Sequence[str]]],
|
||||
available_columns: Sequence[str],
|
||||
) -> Optional[List[str]]:
|
||||
"""Normalize and validate optional selected columns."""
|
||||
if columns is None:
|
||||
configured_columns = self.config.get("columns")
|
||||
if configured_columns is None:
|
||||
return None
|
||||
columns = configured_columns
|
||||
|
||||
if isinstance(columns, str):
|
||||
normalized = [columns]
|
||||
else:
|
||||
normalized = list(columns)
|
||||
|
||||
missing = [col for col in normalized if col not in available_columns]
|
||||
if missing:
|
||||
raise ValidationError(
|
||||
"Column(s) not found in Arrow schema: "
|
||||
f"{', '.join(missing)}. Available columns: "
|
||||
f"{', '.join(available_columns)}"
|
||||
)
|
||||
return normalized
|
||||
|
||||
def _select_schema(self, schema: Any, columns: Optional[List[str]]) -> Any:
|
||||
"""Return schema limited to selected columns when provided."""
|
||||
if columns is None:
|
||||
return schema
|
||||
fields = [schema.field(column) for column in columns]
|
||||
return pa.schema(fields, metadata=schema.metadata)
|
||||
|
||||
def _schema_to_dict(self, schema: Any) -> Dict[str, Any]:
|
||||
"""Convert PyArrow schema to serializable metadata."""
|
||||
fields = []
|
||||
for schema_field in schema:
|
||||
fields.append(
|
||||
{
|
||||
"name": schema_field.name,
|
||||
"type": str(schema_field.type),
|
||||
"nullable": schema_field.nullable,
|
||||
"metadata": self._decode_metadata_map(schema_field.metadata),
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"columns": [field_info["name"] for field_info in fields],
|
||||
"fields": fields,
|
||||
"metadata": self._decode_metadata_map(schema.metadata),
|
||||
}
|
||||
|
||||
def _file_metadata(self, file_path: Path, reader: Any) -> Dict[str, Any]:
|
||||
"""Extract metadata for a single Arrow IPC / Feather file."""
|
||||
total_rows = 0
|
||||
batch_info: List[Dict[str, Any]] = []
|
||||
|
||||
# If it's a stream, open a new one to avoid exhausting the caller's reader
|
||||
meta_reader = reader
|
||||
if reader.is_stream:
|
||||
meta_reader = self._open_file(file_path)
|
||||
|
||||
for i, batch in enumerate(meta_reader.iter_batches()):
|
||||
total_rows += batch.num_rows
|
||||
batch_info.append(
|
||||
{
|
||||
"batch_index": i,
|
||||
"num_rows": batch.num_rows,
|
||||
"num_columns": batch.num_columns,
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"format": "arrow",
|
||||
"source_type": "file",
|
||||
"file": str(file_path),
|
||||
"file_size": file_path.stat().st_size,
|
||||
"total_rows": total_rows,
|
||||
"num_record_batches": len(batch_info),
|
||||
"num_columns": len(reader.schema),
|
||||
"record_batches": batch_info,
|
||||
"schema_metadata": self._decode_metadata_map(reader.schema.metadata),
|
||||
}
|
||||
|
||||
def _decode_metadata_map(
|
||||
self, metadata: Optional[Dict[Any, Any]]
|
||||
) -> Dict[str, str]:
|
||||
"""Decode PyArrow metadata bytes to strings."""
|
||||
if not metadata:
|
||||
return {}
|
||||
|
||||
decoded = {}
|
||||
for key, value in metadata.items():
|
||||
decoded[self._decode_metadata_value(key)] = self._decode_metadata_value(
|
||||
value
|
||||
)
|
||||
return decoded
|
||||
|
||||
def _decode_metadata_value(self, value: Any) -> str:
|
||||
"""Decode a metadata key or value."""
|
||||
if isinstance(value, bytes):
|
||||
return value.decode("utf-8", errors="replace")
|
||||
return str(value)
|
||||
@@ -196,6 +196,7 @@ class FileTypeDetector:
|
||||
b"\x47\x49\x46\x38": "gif", # GIF image
|
||||
b"PK\x03\x04": "zip", # ZIP (alternative)
|
||||
b"PAR1": "parquet", # Apache Parquet
|
||||
b"ARROW1\x00\x00": "arrow", # Apache Arrow IPC
|
||||
}
|
||||
|
||||
# Check if content starts with any known magic number
|
||||
|
||||
@@ -19,6 +19,11 @@ Parquet Ingestion:
|
||||
- "schema": Parquet schema extraction
|
||||
- "metadata": Parquet file or directory metadata extraction
|
||||
|
||||
Arrow IPC / Feather Ingestion:
|
||||
- "file": Single Arrow IPC or Feather file ingestion
|
||||
- "schema": Arrow schema extraction
|
||||
- "metadata": Arrow file metadata extraction
|
||||
|
||||
XML Ingestion:
|
||||
- "file": Single XML file ingestion with structured parsing
|
||||
- "directory": Directory ingestion for XML files
|
||||
@@ -151,6 +156,7 @@ Main Functions:
|
||||
- ingest_email: Email ingestion wrapper
|
||||
- ingest_database: Database ingestion wrapper
|
||||
- ingest_parquet: Parquet ingestion wrapper
|
||||
- ingest_arrow: Arrow IPC / Feather ingestion wrapper
|
||||
- ingest_xml: XML ingestion wrapper
|
||||
- ingest: Unified ingestion function with source type dispatch
|
||||
- get_ingest_method: Get ingestion method by name
|
||||
@@ -179,6 +185,7 @@ from .registry import method_registry
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .api_ingestor import APIData
|
||||
from .arrow_ingestor import ArrowData
|
||||
from .db_ingestor import TableData
|
||||
from .email_ingestor import EmailData
|
||||
from .feed_ingestor import FeedData
|
||||
@@ -350,6 +357,83 @@ def ingest_parquet(
|
||||
raise
|
||||
|
||||
|
||||
def ingest_arrow(
|
||||
source: Union[str, Path, List[Union[str, Path]]],
|
||||
method: str = "file",
|
||||
**kwargs,
|
||||
) -> Union[ArrowData, List[ArrowData], Dict[str, Any]]:
|
||||
"""
|
||||
Ingest Apache Arrow IPC or Feather files.
|
||||
|
||||
Args:
|
||||
source: Arrow / Feather file path or list of paths
|
||||
method: Ingestion method:
|
||||
- "file": Single Arrow IPC / Feather file ingestion
|
||||
- "schema": Extract schema without reading data
|
||||
- "metadata": Extract file metadata without reading data
|
||||
**kwargs: Additional options passed to ArrowIngestor
|
||||
|
||||
Returns:
|
||||
ArrowData, list of ArrowData, or metadata/schema dictionary
|
||||
|
||||
Examples:
|
||||
>>> from semantica.ingest.methods import ingest_arrow
|
||||
>>> data = ingest_arrow("events.arrow", columns=["id"], limit=100)
|
||||
>>> schema = ingest_arrow("events.arrow", method="schema")
|
||||
>>> feather = ingest_arrow("data.feather")
|
||||
"""
|
||||
custom_method = method_registry.get("arrow", method)
|
||||
if custom_method and custom_method != ingest_arrow:
|
||||
try:
|
||||
return custom_method(source, **kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Custom method {method} failed: {e}, falling back to default"
|
||||
)
|
||||
|
||||
try:
|
||||
try:
|
||||
from .arrow_ingestor import ArrowIngestor
|
||||
except ModuleNotFoundError as exc:
|
||||
if _is_missing_dependency(exc, "pyarrow"):
|
||||
raise _missing_optional_dependency(
|
||||
"Arrow ingestion",
|
||||
"pyarrow",
|
||||
) from exc
|
||||
raise
|
||||
|
||||
config = ingest_config.get_method_config("arrow")
|
||||
config.update(kwargs)
|
||||
try:
|
||||
ingestor = ArrowIngestor(**config)
|
||||
except ImportError as exc:
|
||||
raise _missing_optional_dependency(
|
||||
"Arrow ingestion",
|
||||
"pyarrow",
|
||||
) from exc
|
||||
|
||||
def _run_single(path: Union[str, Path]) -> Union[ArrowData, Dict[str, Any]]:
|
||||
source_path = Path(path)
|
||||
|
||||
if method == "schema":
|
||||
return ingestor.extract_schema(source_path, **kwargs)
|
||||
if method == "metadata":
|
||||
return ingestor.extract_metadata(source_path, **kwargs)
|
||||
|
||||
return ingestor.ingest_file(source_path, **kwargs)
|
||||
|
||||
if isinstance(source, list):
|
||||
return [_run_single(path) for path in source]
|
||||
|
||||
return _run_single(source)
|
||||
|
||||
except ConfigurationError:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to ingest Arrow: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def ingest_xml(
|
||||
source: Union[str, Path, List[Union[str, Path]]],
|
||||
method: str = "file",
|
||||
@@ -1256,6 +1340,8 @@ def ingest(
|
||||
source_type = "ontology"
|
||||
elif source_str_lower.endswith((".parquet", ".pq")):
|
||||
source_type = "parquet"
|
||||
elif source_str_lower.endswith((".arrow", ".feather", ".ipc")):
|
||||
source_type = "arrow"
|
||||
elif source_str_lower.endswith(".xml"):
|
||||
source_type = "xml"
|
||||
else:
|
||||
@@ -1268,6 +1354,15 @@ def ingest(
|
||||
)
|
||||
):
|
||||
source_type = "parquet"
|
||||
elif (
|
||||
isinstance(sources, list)
|
||||
and sources
|
||||
and all(
|
||||
str(source).lower().endswith((".arrow", ".feather", ".ipc"))
|
||||
for source in sources
|
||||
)
|
||||
):
|
||||
source_type = "arrow"
|
||||
elif (
|
||||
isinstance(sources, list)
|
||||
and sources
|
||||
@@ -1306,6 +1401,8 @@ def ingest(
|
||||
return {"data": ingest_database(sources, method=method, **kwargs)}
|
||||
elif source_type == "parquet":
|
||||
return {"data": ingest_parquet(sources, method=method or "file", **kwargs)}
|
||||
elif source_type == "arrow":
|
||||
return {"data": ingest_arrow(sources, method=method or "file", **kwargs)}
|
||||
elif source_type == "xml":
|
||||
return {"xml": ingest_xml(sources, method=method or "file", **kwargs)}
|
||||
elif source_type == "ontology":
|
||||
@@ -1403,6 +1500,11 @@ method_registry.register("parquet", "directory", ingest_parquet)
|
||||
method_registry.register("parquet", "schema", ingest_parquet)
|
||||
method_registry.register("parquet", "metadata", ingest_parquet)
|
||||
method_registry.register("file", "parquet", ingest_parquet)
|
||||
method_registry.register("arrow", "default", ingest_arrow)
|
||||
method_registry.register("arrow", "file", ingest_arrow)
|
||||
method_registry.register("arrow", "schema", ingest_arrow)
|
||||
method_registry.register("arrow", "metadata", ingest_arrow)
|
||||
method_registry.register("file", "arrow", ingest_arrow)
|
||||
method_registry.register("xml", "default", ingest_xml)
|
||||
method_registry.register("xml", "file", ingest_xml)
|
||||
method_registry.register("xml", "directory", ingest_xml)
|
||||
|
||||
@@ -64,6 +64,7 @@ class MethodRegistry:
|
||||
"public_api": {},
|
||||
"mcp": {},
|
||||
"parquet": {},
|
||||
"arrow": {},
|
||||
"xml": {},
|
||||
"ingest": {},
|
||||
}
|
||||
|
||||
@@ -59,6 +59,9 @@ SUPPORTED_DOCUMENT_FORMATS = [
|
||||
"pptx",
|
||||
"parquet",
|
||||
"pq",
|
||||
"arrow",
|
||||
"feather",
|
||||
"ipc",
|
||||
]
|
||||
|
||||
SUPPORTED_IMAGE_FORMATS = ["jpg", "jpeg", "png", "gif", "bmp", "tiff", "webp", "svg"]
|
||||
|
||||
@@ -0,0 +1,546 @@
|
||||
"""Comprehensive tests for the Arrow IPC / Feather ingestor."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pa = pytest.importorskip("pyarrow")
|
||||
ipc = pytest.importorskip("pyarrow.ipc")
|
||||
|
||||
from semantica.ingest import ( # noqa: E402
|
||||
ArrowData,
|
||||
ArrowIngestor,
|
||||
ingest,
|
||||
ingest_arrow,
|
||||
ingest_file,
|
||||
list_available_methods,
|
||||
)
|
||||
from semantica.ingest.file_ingestor import FileTypeDetector # noqa: E402
|
||||
from semantica.utils.exceptions import ProcessingError, ValidationError # noqa: E402
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _write_arrow_ipc(path: Path, table: pa.Table) -> Path:
|
||||
"""Write a table as an Arrow IPC file."""
|
||||
with pa.OSFile(str(path), "wb") as sink:
|
||||
writer = ipc.new_file(sink, table.schema)
|
||||
writer.write_table(table)
|
||||
writer.close()
|
||||
return path
|
||||
|
||||
|
||||
def _write_arrow_ipc_batches(path: Path, batches, schema: pa.Schema) -> Path:
|
||||
"""Write multiple record batches as an Arrow IPC file."""
|
||||
with pa.OSFile(str(path), "wb") as sink:
|
||||
writer = ipc.new_file(sink, schema)
|
||||
for batch in batches:
|
||||
writer.write_batch(batch)
|
||||
writer.close()
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_arrow(tmp_path: Path) -> Path:
|
||||
"""Simple Arrow IPC file with 4 columns and 3 rows."""
|
||||
table = pa.table(
|
||||
{
|
||||
"id": [1, 2, 3],
|
||||
"name": ["alpha", "beta", "gamma"],
|
||||
"score": [0.7, 0.8, 0.9],
|
||||
"city": ["Pune", "Delhi", "Mumbai"],
|
||||
}
|
||||
)
|
||||
return _write_arrow_ipc(tmp_path / "events.arrow", table)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_feather(tmp_path: Path) -> Path:
|
||||
"""Feather v2 file (Arrow IPC format)."""
|
||||
import pyarrow.feather as pf
|
||||
|
||||
table = pa.table({"x": [10, 20], "y": ["a", "b"]})
|
||||
path = tmp_path / "data.feather"
|
||||
pf.write_feather(table, str(path))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def multi_type_arrow(tmp_path: Path) -> Path:
|
||||
"""Arrow file with multiple data types."""
|
||||
table = pa.table(
|
||||
{
|
||||
"int_col": pa.array([1, 2, 3], type=pa.int64()),
|
||||
"float_col": pa.array([1.1, 2.2, 3.3], type=pa.float64()),
|
||||
"str_col": pa.array(["a", "b", "c"], type=pa.string()),
|
||||
"bool_col": pa.array([True, False, True], type=pa.bool_()),
|
||||
}
|
||||
)
|
||||
return _write_arrow_ipc(tmp_path / "multi.arrow", table)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def multi_batch_arrow(tmp_path: Path) -> Path:
|
||||
"""Arrow file with multiple record batches."""
|
||||
schema = pa.schema([("id", pa.int64()), ("value", pa.string())])
|
||||
batch1 = pa.record_batch({"id": [1, 2], "value": ["a", "b"]}, schema=schema)
|
||||
batch2 = pa.record_batch({"id": [3, 4], "value": ["c", "d"]}, schema=schema)
|
||||
batch3 = pa.record_batch({"id": [5], "value": ["e"]}, schema=schema)
|
||||
return _write_arrow_ipc_batches(
|
||||
tmp_path / "batched.arrow", [batch1, batch2, batch3], schema
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def empty_table_arrow(tmp_path: Path) -> Path:
|
||||
"""Arrow file with zero rows."""
|
||||
table = pa.table(
|
||||
{"id": pa.array([], type=pa.int64()), "name": pa.array([], type=pa.string())}
|
||||
)
|
||||
return _write_arrow_ipc(tmp_path / "empty.arrow", table)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def nullable_arrow(tmp_path: Path) -> Path:
|
||||
"""Arrow file with null values."""
|
||||
table = pa.table(
|
||||
{
|
||||
"id": [1, 2, 3],
|
||||
"name": ["alpha", None, "gamma"],
|
||||
"score": [0.7, None, 0.9],
|
||||
}
|
||||
)
|
||||
return _write_arrow_ipc(tmp_path / "nullable.arrow", table)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def metadata_arrow(tmp_path: Path) -> Path:
|
||||
"""Arrow file with schema-level and field-level metadata."""
|
||||
field_id = pa.field("id", pa.int64(), metadata={b"description": b"primary key"})
|
||||
field_name = pa.field("name", pa.string())
|
||||
schema = pa.schema(
|
||||
[field_id, field_name],
|
||||
metadata={b"author": b"test", b"version": b"1.0"},
|
||||
)
|
||||
table = pa.table({"id": [1, 2], "name": ["a", "b"]}, schema=schema)
|
||||
return _write_arrow_ipc(tmp_path / "meta.arrow", table)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_metadata_arrow(tmp_path: Path) -> Path:
|
||||
"""Arrow file with no schema metadata."""
|
||||
table = pa.table({"x": [1]})
|
||||
return _write_arrow_ipc(tmp_path / "nometa.arrow", table)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Happy-path tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_arrow_file_ingestion_reads_data_schema_and_metadata(
|
||||
sample_arrow: Path,
|
||||
) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(sample_arrow)
|
||||
|
||||
assert isinstance(result, ArrowData)
|
||||
assert result.row_count == 3
|
||||
assert result.columns == ["id", "name", "score", "city"]
|
||||
assert result.data[0]["name"] == "alpha"
|
||||
assert result.schema["columns"] == ["id", "name", "score", "city"]
|
||||
assert result.schema["fields"][0]["type"] == "int64"
|
||||
assert result.metadata["total_rows"] == 3
|
||||
assert result.metadata["format"] == "arrow"
|
||||
|
||||
|
||||
def test_feather_file_ingestion(sample_feather: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(sample_feather)
|
||||
|
||||
assert isinstance(result, ArrowData)
|
||||
assert result.row_count == 2
|
||||
assert result.columns == ["x", "y"]
|
||||
assert result.data == [{"x": 10, "y": "a"}, {"x": 20, "y": "b"}]
|
||||
|
||||
|
||||
def test_multi_type_columns(multi_type_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(multi_type_arrow)
|
||||
|
||||
assert result.row_count == 3
|
||||
type_map = {f["name"]: f["type"] for f in result.schema["fields"]}
|
||||
assert type_map["int_col"] == "int64"
|
||||
assert type_map["float_col"] == "double"
|
||||
assert type_map["str_col"] == "string"
|
||||
assert type_map["bool_col"] == "bool"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Schema tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_extract_schema(sample_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
schema = ingestor.extract_schema(sample_arrow)
|
||||
|
||||
assert schema["columns"] == ["id", "name", "score", "city"]
|
||||
assert len(schema["fields"]) == 4
|
||||
assert all("nullable" in f for f in schema["fields"])
|
||||
|
||||
|
||||
def test_schema_metadata_extraction(metadata_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
schema = ingestor.extract_schema(metadata_arrow)
|
||||
|
||||
assert schema["metadata"]["author"] == "test"
|
||||
assert schema["metadata"]["version"] == "1.0"
|
||||
# Field-level metadata
|
||||
id_field = next(f for f in schema["fields"] if f["name"] == "id")
|
||||
assert id_field["metadata"]["description"] == "primary key"
|
||||
|
||||
|
||||
def test_metadata_extraction(sample_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
metadata = ingestor.extract_metadata(sample_arrow)
|
||||
|
||||
assert metadata["total_rows"] == 3
|
||||
assert metadata["num_columns"] == 4
|
||||
assert metadata["format"] == "arrow"
|
||||
assert metadata["source_type"] == "file"
|
||||
assert metadata["file_size"] > 0
|
||||
assert len(metadata["record_batches"]) >= 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Selective columns and limits
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_selective_column_reading(sample_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(sample_arrow, columns=["id", "name"])
|
||||
|
||||
assert result.columns == ["id", "name"]
|
||||
assert result.row_count == 3
|
||||
assert all(set(row.keys()) == {"id", "name"} for row in result.data)
|
||||
|
||||
|
||||
def test_row_limit(sample_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(sample_arrow, limit=2)
|
||||
|
||||
assert result.row_count == 2
|
||||
assert result.metadata["limit"] == 2
|
||||
|
||||
|
||||
def test_limit_zero(sample_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(sample_arrow, limit=0)
|
||||
|
||||
assert result.row_count == 0
|
||||
assert result.data == []
|
||||
|
||||
|
||||
def test_selective_columns_with_limit(sample_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(sample_arrow, columns=["id", "name"], limit=2)
|
||||
|
||||
assert result.row_count == 2
|
||||
assert result.columns == ["id", "name"]
|
||||
assert result.data == [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batch tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_multi_batch_reading(multi_batch_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(multi_batch_arrow)
|
||||
|
||||
assert result.row_count == 5
|
||||
assert result.columns == ["id", "value"]
|
||||
assert result.metadata["num_record_batches"] == 3
|
||||
assert [row["id"] for row in result.data] == [1, 2, 3, 4, 5]
|
||||
|
||||
|
||||
def test_multi_batch_with_limit(multi_batch_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(multi_batch_arrow, limit=3)
|
||||
|
||||
assert result.row_count == 3
|
||||
assert [row["id"] for row in result.data] == [1, 2, 3]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# include_data=False
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_include_data_false(sample_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(sample_arrow, include_data=False)
|
||||
|
||||
assert result.row_count == 0
|
||||
assert result.data == []
|
||||
assert result.metadata["include_data"] is False
|
||||
assert result.schema["columns"] == ["id", "name", "score", "city"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Edge cases
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_empty_table(empty_table_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(empty_table_arrow)
|
||||
|
||||
assert result.row_count == 0
|
||||
assert result.data == []
|
||||
assert result.columns == ["id", "name"]
|
||||
assert result.schema["columns"] == ["id", "name"]
|
||||
|
||||
|
||||
def test_null_values(nullable_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
result = ingestor.ingest_file(nullable_arrow)
|
||||
|
||||
assert result.row_count == 3
|
||||
assert result.data[1]["name"] is None
|
||||
assert result.data[1]["score"] is None
|
||||
assert result.data[0]["name"] == "alpha"
|
||||
assert result.data[2]["score"] == 0.9
|
||||
|
||||
|
||||
def test_missing_metadata_no_exception(no_metadata_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
metadata = ingestor.extract_metadata(no_metadata_arrow)
|
||||
schema = ingestor.extract_schema(no_metadata_arrow)
|
||||
|
||||
# Should return empty dict, not raise
|
||||
assert isinstance(metadata["schema_metadata"], dict)
|
||||
assert isinstance(schema["metadata"], dict)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Failure cases
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_missing_file_raises_validation_error() -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
with pytest.raises(ValidationError, match="Arrow file not found"):
|
||||
ingestor.ingest_file("/nonexistent/path/data.arrow")
|
||||
|
||||
|
||||
def test_wrong_extension_raises_validation_error(tmp_path: Path) -> None:
|
||||
path = tmp_path / "data.txt"
|
||||
path.write_text("hello")
|
||||
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
with pytest.raises(ValidationError, match="not an Arrow/Feather file"):
|
||||
ingestor.ingest_file(path)
|
||||
|
||||
|
||||
def test_corrupted_file_raises_processing_error(tmp_path: Path) -> None:
|
||||
path = tmp_path / "bad.arrow"
|
||||
path.write_bytes(b"\x00\x01\x02\x03garbage data here")
|
||||
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
with pytest.raises(ProcessingError, match="Failed to open Arrow file"):
|
||||
ingestor.ingest_file(path)
|
||||
|
||||
|
||||
def test_invalid_text_file_raises_processing_error(tmp_path: Path) -> None:
|
||||
path = tmp_path / "invalid.arrow"
|
||||
path.write_text("this is not an arrow file")
|
||||
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
with pytest.raises(ProcessingError):
|
||||
ingestor.ingest_file(path)
|
||||
|
||||
|
||||
def test_negative_limit_raises_validation_error(sample_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
with pytest.raises(
|
||||
ValidationError, match="limit must be greater than or equal to 0"
|
||||
):
|
||||
ingestor.ingest_file(sample_arrow, limit=-1)
|
||||
|
||||
|
||||
def test_invalid_column_raises_validation_error(sample_arrow: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
|
||||
with pytest.raises(ValidationError, match="Column.*not found"):
|
||||
ingestor.ingest_file(sample_arrow, columns=["nonexistent_col"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Convenience functions and dispatch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_ingest_arrow_convenience_function(sample_arrow: Path) -> None:
|
||||
result = ingest_arrow(sample_arrow, columns=["name"])
|
||||
|
||||
assert isinstance(result, ArrowData)
|
||||
assert result.columns == ["name"]
|
||||
|
||||
|
||||
def test_ingest_arrow_schema_method(sample_arrow: Path) -> None:
|
||||
schema = ingest_arrow(sample_arrow, method="schema")
|
||||
|
||||
assert isinstance(schema, dict)
|
||||
assert "columns" in schema
|
||||
assert "fields" in schema
|
||||
|
||||
|
||||
def test_ingest_arrow_metadata_method(sample_arrow: Path) -> None:
|
||||
metadata = ingest_arrow(sample_arrow, method="metadata")
|
||||
|
||||
assert isinstance(metadata, dict)
|
||||
assert metadata["format"] == "arrow"
|
||||
|
||||
|
||||
def test_ingest_arrow_list_of_files(sample_arrow: Path, sample_feather: Path) -> None:
|
||||
results = ingest_arrow([sample_arrow, sample_feather])
|
||||
|
||||
assert isinstance(results, list)
|
||||
assert len(results) == 2
|
||||
assert all(isinstance(r, ArrowData) for r in results)
|
||||
|
||||
|
||||
def test_unified_ingest_auto_detects_arrow(sample_arrow: Path) -> None:
|
||||
result = ingest(sample_arrow)
|
||||
|
||||
assert isinstance(result, dict)
|
||||
assert "data" in result
|
||||
assert isinstance(result["data"], ArrowData)
|
||||
|
||||
|
||||
def test_unified_ingest_auto_detects_feather(sample_feather: Path) -> None:
|
||||
result = ingest(sample_feather)
|
||||
|
||||
assert isinstance(result, dict)
|
||||
assert "data" in result
|
||||
assert isinstance(result["data"], ArrowData)
|
||||
|
||||
|
||||
def test_ingest_file_method_arrow(sample_arrow: Path) -> None:
|
||||
result = ingest_file(sample_arrow, method="arrow")
|
||||
|
||||
assert isinstance(result, ArrowData)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Magic number detection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_file_type_detector_recognizes_arrow_magic_number() -> None:
|
||||
detector = FileTypeDetector()
|
||||
|
||||
# Arrow IPC files start with "ARROW1\x00\x00"
|
||||
arrow_content = b"ARROW1\x00\x00" + b"\x00" * 100
|
||||
assert detector.detect_type("noext", content=arrow_content) == "arrow"
|
||||
assert detector.is_supported("arrow")
|
||||
assert detector.is_supported("feather")
|
||||
assert detector.is_supported("ipc")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_arrow_methods_registered() -> None:
|
||||
methods = list_available_methods("arrow")
|
||||
|
||||
assert "arrow" in methods
|
||||
assert "file" in methods["arrow"]
|
||||
assert "schema" in methods["arrow"]
|
||||
assert "metadata" in methods["arrow"]
|
||||
|
||||
|
||||
def test_ingest_method_alias() -> None:
|
||||
"""ingest(method='arrow') in file context should dispatch to arrow."""
|
||||
ingestor = ArrowIngestor()
|
||||
# Just verify the ingestor can be instantiated — the actual dispatch
|
||||
# is tested via ingest_file(method="arrow") above.
|
||||
assert ingestor is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Streaming Format Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _write_arrow_stream(path: Path, table: pa.Table) -> Path:
|
||||
"""Write a table as an Arrow IPC Stream file."""
|
||||
with pa.OSFile(str(path), "wb") as sink:
|
||||
writer = ipc.new_stream(sink, table.schema)
|
||||
writer.write_table(table)
|
||||
writer.close()
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_arrow_stream(tmp_path: Path) -> Path:
|
||||
"""Arrow IPC Stream format file."""
|
||||
table = pa.table(
|
||||
{
|
||||
"id": [100, 200],
|
||||
"name": ["stream1", "stream2"],
|
||||
}
|
||||
)
|
||||
return _write_arrow_stream(tmp_path / "stream.ipc", table)
|
||||
|
||||
|
||||
def test_arrow_stream_file_ingestion(sample_arrow_stream: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
result = ingestor.ingest_file(sample_arrow_stream)
|
||||
|
||||
assert isinstance(result, ArrowData)
|
||||
assert result.row_count == 2
|
||||
assert result.columns == ["id", "name"]
|
||||
assert result.data == [
|
||||
{"id": 100, "name": "stream1"},
|
||||
{"id": 200, "name": "stream2"},
|
||||
]
|
||||
|
||||
|
||||
def test_arrow_stream_metadata_and_schema(sample_arrow_stream: Path) -> None:
|
||||
ingestor = ArrowIngestor()
|
||||
schema = ingestor.extract_schema(sample_arrow_stream)
|
||||
metadata = ingestor.extract_metadata(sample_arrow_stream)
|
||||
|
||||
assert schema["columns"] == ["id", "name"]
|
||||
assert metadata["format"] == "arrow"
|
||||
assert metadata["total_rows"] == 2
|
||||
Reference in New Issue
Block a user