# PyArrow Multi-Key Vectorization API Reference

This document provides detailed API reference for the PyArrow multi-key vectorization helper functions and updates to existing functions to support multi-column keys.

## Core Helper Functions

### `_create_composite_key_array()`

Create a StructArray representing composite keys for efficient comparison. This function enables vectorized multi-column key operations by staying entirely in Arrow space.

```python
def _create_composite_key_array(
    table: pa.Table,
    key_columns: list[str]
) -> pa.Array
```

**Parameters:**

| Parameter | Type | Description |
|-----------|------|-------------|
| `table` | `pa.Table` | PyArrow table containing the key columns. |
| `key_columns` | `list[str]` | List of column names to include in the composite key. |

**Returns:**

- `pa.StructArray`: A StructArray where each element represents a composite key.

**Raises:**

- `ValueError`: If `key_columns` is empty.
- `KeyError`: If any key column is not found in the table.

**Examples:**

```python
import pyarrow as pa
from fsspeckit.datasets.pyarrow.dataset import _create_composite_key_array

# Create sample table
table = pa.Table.from_pydict({
    "tenant_id": [1, 1, 2, 2],
    "user_id": [100, 101, 200, 201],
    "record_id": [1, 2, 1, 2],
    "value": [10, 20, 30, 40]
})

# Create composite key
composite_keys = _create_composite_key_array(
    table=table,
    key_columns=["tenant_id", "user_id", "record_id"]
)

print(f"Composite key type: {type(composite_keys)}")
print(f"Composite keys: {composite_keys.to_pylist()}")
```

**Performance Notes:**

- Uses `pa.StructArray.from_arrays()` for zero-copy operations
- Handles ChunkedArrays by combining them efficiently
- Maintains Arrow-native types for optimal comparison performance

---

### `_filter_by_key_membership()`

Filter table rows based on key membership using PyArrow joins. This is the core function for vectorized multi-key filtering.

```python
def _filter_by_key_membership(
    table: pa.Table,
    key_columns: list[str],
    reference_keys: pa.Table,
    keep_matches: bool = True
) -> pa.Table
```

**Parameters:**

| Parameter | Type | Description |
|-----------|------|-------------|
| `table` | `pa.Table` | Table to filter. |
| `key_columns` | `list[str]` | List of column names to use as keys. |
| `reference_keys` | `pa.Table` | Table containing the keys to match against. |
| `keep_matches` | `bool` | If True, keep rows present in reference_keys (semi-join). If False, keep rows NOT present in reference_keys (anti-join). Default: `True`. |

**Returns:**

- `pa.Table`: Filtered PyArrow Table.

**Behavior:**

1. **Primary Path**: Uses `pa.Table.join()` with `join_type="semi"` or `"anti"` for optimal performance
2. **Fallback Path**: Falls back to string-based keys if native join fails due to type compatibility issues

**Examples:**

```python
import pyarrow as pa
from fsspeckit.datasets.pyarrow.dataset import _filter_by_key_membership

# Source data
source_table = pa.Table.from_pydict({
    "tenant_id": [1, 1, 2, 2, 3],
    "user_id": [100, 101, 200, 201, 300],
    "value": [10, 20, 30, 40, 50]
})

# Reference keys to match against
reference_table = pa.Table.from_pydict({
    "tenant_id": [1, 2],
    "user_id": [100, 200]
})

# Keep only matching rows (semi-join)
matched = _filter_by_key_membership(
    table=source_table,
    key_columns=["tenant_id", "user_id"],
    reference_keys=reference_table,
    keep_matches=True
)

print("Matching rows:")
print(matched.to_pandas())

# Keep only non-matching rows (anti-join)
non_matched = _filter_by_key_membership(
    table=source_table,
    key_columns=["tenant_id", "user_id"],
    reference_keys=reference_table,
    keep_matches=False
)

print("\nNon-matching rows:")
print(non_matched.to_pandas())
```

**Performance Characteristics:**

- Native Arrow joins handle multi-column keys efficiently
- Fallback to string-based keys provides compatibility for complex type combinations
- Zero-copy operations when possible
- Automatic logging of fallback usage

---

### `_create_string_key_array()`

Create a string representation of composite keys as a fallback mechanism. Used when StructArray or Join operations fail due to heterogeneous type combinations.

```python
def _create_string_key_array(
    table: pa.Table,
    key_columns: list[str]
) -> pa.Array
```

**Parameters:**

| Parameter | Type | Description |
|-----------|------|-------------|
| `table` | `pa.Table` | PyArrow table containing the key columns. |
| `key_columns` | `list[str]` | List of column names to include in the composite key. |

**Returns:**

- `pa.StringArray`: A StringArray where each element represents a joined string key.

**Implementation Details:**

- Casts each key column to string
- Handles null values with `"__NULL__"` placeholder
- Joins components with `\x1f` (ASCII unit separator) delimiter
- Optimized for single-key case (no joining needed)

**Examples:**

```python
import pyarrow as pa
from fsspeckit.datasets.pyarrow.dataset import _create_string_key_array

# Table with mixed types
table = pa.Table.from_pydict({
    "tenant_id": [1, 2, None, 4],
    "record_id": ["A001", "B002", "C003", "D004"],
    "timestamp": [1704067200, 1704067260, None, 1704067380]
})

# Create string composite keys
string_keys = _create_string_key_array(
    table=table,
    key_columns=["tenant_id", "record_id", "timestamp"]
)

print("String keys:")
for i, key in enumerate(string_keys):
    print(f"  Row {i}: '{key.as_py()}'")

# Single column case (no joining)
single_key = _create_string_key_array(table, ["record_id"])
print(f"\nSingle column keys: {single_key.to_pylist()}")
```

**Use Cases:**

- Fallback mechanism for incompatible type combinations
- Manual key creation when needed
- Debugging and key inspection

---

## Updated Existing Functions

### `deduplicate_pyarrow()`

The deduplication function now supports multi-column keys with vectorized processing.

**New Parameters:**

| Parameter | Type | Description |
|-----------|------|-------------|
| `key_columns` | `list[str]` | Key columns for deduplication. Supports single or multiple columns. |
| `dedup_order_by` | `str | list[str] | None` | Columns to order by for deduplication when duplicates are found. Defaults to key_columns. |

**Enhanced Behavior:**

1. **Single Column**: Uses existing optimized single-key logic
2. **Multiple Columns**: Uses vectorized multi-key processing with `_create_composite_key_array()`
3. **Chunked Processing**: Handles large datasets with memory-efficient chunking

**Examples:**

```python
# Single column deduplication (unchanged behavior)
single_unique = deduplicate_pyarrow(
    table=table,
    key_columns=["id"]
)

# Multi-column deduplication (new vectorized behavior)
multi_unique = deduplicate_pyarrow(
    table=table,
    key_columns=["tenant_id", "user_id", "record_id"],
    dedup_order_by=["timestamp"],
    keep="last"
)

# With chunking for large datasets
large_unique = deduplicate_pyarrow(
    table=large_table,
    key_columns=["tenant_id", "user_id", "record_id"],
    chunk_size="100MB",
    memory_monitor=memory_monitor
)
```

---

### `merge_parquet_dataset_pyarrow()`

Enhanced to support composite keys in all merge strategies (insert, update, upsert, deduplicate).

**Updated Parameters:**

| Parameter | Type | Description |
|-----------|------|-------------|
| `key_columns` | `str | list[str]` | Key columns for merge. Now supports multiple columns for composite business keys. |

**Enhanced Features:**

1. **Vectorized Key Matching**: Uses `_filter_by_key_membership()` for efficient multi-key operations
2. **Fallback Support**: Graceful degradation for complex type combinations
3. **Memory Efficiency**: Processes large datasets with optimized chunking

**Examples:**

```python
# Single key merge (existing behavior)
stats_single = merge_parquet_dataset_pyarrow(
    data=data,
    path="dataset/",
    strategy="upsert",
    key_columns="id"
)

# Composite key merge (new vectorized behavior)
stats_composite = merge_parquet_dataset_pyarrow(
    data=data,
    path="dataset/",
    strategy="upsert",
    key_columns=["tenant_id", "customer_id", "order_id"],
    dedup_order_by=["updated_at"]
)

# With memory monitoring
stats_monitored = merge_parquet_dataset_pyarrow(
    data=data,
    path="dataset/",
    strategy="upsert",
    key_columns=["tenant_id", "customer_id", "order_id"],
    chunk_size="50MB",
    memory_monitor=memory_monitor
)
```

---

## Supported Key Types and Combinations

### Homogeneous Type Keys (Maximum Performance)

| Type Combination | Example | Performance Level |
|------------------|---------|-------------------|
| All Integers | `["id", "sub_id", "sequence"]` | ✅ Maximum |
| All Strings | `["tenant", "category", "subtype"]` | ✅ Maximum |
| All Timestamps | `["created", "updated", "deleted"]` | ✅ Maximum |
| Mixed Primitives | `["int_col", "string_col", "timestamp"]` | ✅ Excellent |

### Heterogeneous Type Keys (With Fallback)

| Type Combination | Example | Performance Level |
|------------------|---------|-------------------|
| + Binary Data | `["id", "binary_col", "status"]` | ⚠️ Good (85-90%) |
| + Complex Objects | `["id", "struct_col", "list_col"]` | ⚠️ Fair (70-80%) |
| + Nested Structs | `["id", "nested.struct.field"]` | ⚠️ Fair (70-80%) |

### Null Handling

- **StructArray Path**: Preserves null semantics in native Arrow operations
- **String Fallback**: Converts nulls to `"__NULL__"` placeholder
- **Consistent Behavior**: Both approaches maintain identical null semantics

---

## Error Handling and Fallback

### Automatic Fallback Behavior

```python
# The system automatically detects type compatibility issues
try:
    # Try native Arrow operations first
    result = deduplicate_pyarrow(table, key_columns=["mixed_types"])
except Exception as e:
    # Fallback automatically triggered if needed
    logger.warning("Using string-based fallback for heterogeneous types")
    # Operation continues with string-based keys
```

### Fallback Logging

When fallback is triggered, the system logs detailed information:

```
WARNING: Primary join approach failed, falling back to string-based keys.
This can happen with heterogeneous type combinations. Error: Schema error: [...]
```

### Manual Fallback Control

```python
# Force string-based keys if needed
from fsspeckit.datasets.pyarrow.dataset import _create_string_key_array

string_keys = _create_string_key_array(table, key_columns)
# Use string_keys for manual operations
```

---

## Performance Optimization Parameters

### Chunk Size Selection

| Dataset Size | Recommended Chunk Size | Memory Usage | Performance |
|--------------|----------------------|--------------|-------------|
| < 1M rows | "50MB" | Low | Good |
| 1-10M rows | "100MB" | Medium | Excellent |
| 10-100M rows | "200MB" | High | Excellent |
| > 100M rows | "500MB" | Very High | Optimal |

### Memory Monitoring Integration

```python
from fsspeckit.datasets.pyarrow.memory import MemoryMonitor

monitor = MemoryMonitor(
    max_pyarrow_mb=2048,           # Limit PyArrow memory
    max_process_memory_mb=8192,    # Limit total process memory
    min_system_available_mb=1024   # Keep system memory free
)

# Use with any multi-key operation
result = deduplicate_pyarrow(
    table=large_table,
    key_columns=["tenant_id", "user_id", "record_id"],
    chunk_size="100MB",
    memory_monitor=monitor
)
```

### Performance Monitoring

```python
from fsspeckit.datasets.pyarrow.dataset import PerformanceMonitor

monitor = PerformanceMonitor()

# Monitor operations
monitor.start_op("vectorized_deduplication")
result = deduplicate_pyarrow(table, key_columns=composite_keys)
monitor.end_op()

# Get detailed metrics
metrics = monitor.get_metrics(
    total_rows_before=table.num_rows,
    total_rows_after=result.num_rows,
    total_bytes=table.nbytes
)

print(f"Performance: {metrics['rows_per_sec']:.0f} rows/sec")
print(f"Memory peak: {metrics['memory_peak_mb']:.1f} MB")
```

---

## Migration Guide

### From Single-Column to Multi-Column Keys

**Before:**
```python
# Legacy single-key approach
result = deduplicate_pyarrow(table, key_columns=["id"])
```

**After:**
```python
# Enhanced multi-key support
result = deduplicate_pyarrow(
    table=table,
    key_columns=["tenant_id", "user_id", "record_id"],  # Composite key
    dedup_order_by=["timestamp"],
    keep="last"
)
```

### Performance Impact

| Operation | Single Key | Multi Key (Legacy) | Multi Key (Vectorized) |
|-----------|------------|--------------------|------------------------|
| Deduplication | 1.0x | 0.2-0.3x | 0.9-1.0x |
| Merge - Key Matching | 1.0x | 0.1-0.2x | 0.8-0.9x |
| Memory Usage | 1.0x | 3-5x | 1.0-1.2x |

---

## Best Practices

### 1. Key Column Ordering

```python
# Optimal: High cardinality first
optimal_keys = ["user_id", "tenant_id", "status"]

# Less optimal: Low cardinality first
suboptimal_keys = ["status", "tenant_id", "user_id"]
```

### 2. Type Selection

```python
# Prefer compatible types when possible
good_types = ["user_id:int64", "tenant_id:int64", "sequence:int64"]
avoid_mixed = ["user_id:int64", "tenant_id:string", "timestamp:timestamp"]
```

### 3. Chunk Size Optimization

```python
# Start with conservative chunk size
safe_chunk_size = "100MB"

# Increase based on memory availability
if available_memory_gb > 16:
    optimal_chunk_size = "500MB"
```

### 4. Error Handling

```python
def robust_composite_operation(table, key_columns):
    try:
        # Use vectorized multi-key operations
        return deduplicate_pyarrow(table, key_columns=key_columns)
    except Exception as e:
        logger.warning(f"Vectorized operation failed: {e}")
        # Fallback to smaller chunks or single-key operations
        return fallback_strategy(table, key_columns)
```

---

## Integration with Existing APIs

All existing fsspeckit APIs continue to work unchanged. The multi-key vectorization is additive and backward-compatible:

```python
# Existing code continues to work
result = deduplicate_pyarrow(table, key_columns=["id"])

# New multi-key capabilities available
result = deduplicate_pyarrow(table, key_columns=["tenant_id", "user_id", "record_id"])
```

The helper functions are also available for advanced use cases:

```python
# Direct access to vectorization functions
from fsspeckit.datasets.pyarrow.dataset import (
    _create_composite_key_array,
    _filter_by_key_membership,
    _create_string_key_array
)
```

For practical usage examples, see [Multi-Key Usage Examples](./multi-key-examples.md).

For performance characteristics and optimization details, see [Multi-Key Performance Guide](./multi-key-performance.md).