Coverage for src/dataknobs_data/pandas/metadata.py: 0%
122 statements
« prev ^ index » next coverage.py v7.10.3, created at 2025-08-17 18:57 -0500
1"""Metadata preservation for DataKnobs-Pandas conversions."""
3from dataclasses import dataclass
4from enum import Enum
5from typing import Any, Dict, List, Optional
7import pandas as pd
9from dataknobs_data.records import Record
10from dataknobs_data.fields import Field, FieldType
13class MetadataStrategy(Enum):
14 """Strategy for handling metadata during conversion."""
15 NONE = "none" # Don't preserve metadata
16 ATTRS = "attrs" # Store in DataFrame.attrs
17 COLUMNS = "columns" # Store as additional columns
18 MULTI_INDEX = "multi_index" # Use multi-level column index
21@dataclass
22class MetadataConfig:
23 """Configuration for metadata handling."""
24 strategy: MetadataStrategy = MetadataStrategy.ATTRS
25 include_record_metadata: bool = True
26 include_field_metadata: bool = True
27 metadata_prefix: str = "_meta_"
28 preserve_record_ids: bool = True
31class MetadataHandler:
32 """Handles metadata preservation during conversions."""
34 def __init__(self, config: Optional[MetadataConfig] = None):
35 """Initialize metadata handler.
37 Args:
38 config: Metadata configuration
39 """
40 self.config = config or MetadataConfig()
42 def extract_metadata_from_records(self, records: List[Record]) -> Dict[str, Any]:
43 """Extract metadata from records.
45 Args:
46 records: List of records
48 Returns:
49 Dictionary of metadata
50 """
51 metadata = {
52 "record_count": len(records),
53 "has_record_ids": all(r.id for r in records),
54 "field_names": self._get_all_field_names(records),
55 "field_types": self._get_field_types(records),
56 }
58 if self.config.include_record_metadata:
59 metadata["record_metadata"] = self._extract_record_metadata(records)
61 if self.config.include_field_metadata:
62 metadata["field_metadata"] = self._extract_field_metadata(records)
64 return metadata
66 def apply_metadata_to_dataframe(
67 self,
68 df: pd.DataFrame,
69 metadata: Dict[str, Any],
70 records: Optional[List[Record]] = None
71 ) -> pd.DataFrame:
72 """Apply metadata to DataFrame based on strategy.
74 Args:
75 df: Target DataFrame
76 metadata: Metadata to apply
77 records: Original records (for additional metadata)
79 Returns:
80 DataFrame with metadata
81 """
82 if self.config.strategy == MetadataStrategy.NONE:
83 return df
85 elif self.config.strategy == MetadataStrategy.ATTRS:
86 df.attrs.update(metadata)
87 if records and self.config.preserve_record_ids:
88 record_ids = [r.id for r in records]
89 df.attrs["record_ids"] = record_ids
91 elif self.config.strategy == MetadataStrategy.COLUMNS:
92 # Add metadata as columns
93 if self.config.include_record_metadata and records:
94 for key, values in self._get_record_metadata_columns(records).items():
95 col_name = f"{self.config.metadata_prefix}{key}"
96 df[col_name] = values
98 elif self.config.strategy == MetadataStrategy.MULTI_INDEX:
99 # Create multi-level column index with metadata
100 if "field_types" in metadata:
101 arrays = [
102 df.columns.tolist(),
103 [metadata["field_types"].get(col, "unknown") for col in df.columns]
104 ]
105 df.columns = pd.MultiIndex.from_arrays(
106 arrays,
107 names=["field_name", "field_type"]
108 )
110 return df
112 def extract_metadata_from_dataframe(self, df: pd.DataFrame) -> Dict[str, Any]:
113 """Extract metadata from DataFrame.
115 Args:
116 df: Source DataFrame
118 Returns:
119 Dictionary of metadata
120 """
121 metadata = {}
123 if self.config.strategy == MetadataStrategy.ATTRS:
124 metadata.update(df.attrs)
126 elif self.config.strategy == MetadataStrategy.COLUMNS:
127 # Extract from metadata columns
128 meta_cols = [col for col in df.columns if col.startswith(self.config.metadata_prefix)]
129 for col in meta_cols:
130 key = col.replace(self.config.metadata_prefix, "")
131 metadata[key] = df[col].tolist()
133 elif self.config.strategy == MetadataStrategy.MULTI_INDEX:
134 # Extract from multi-level index
135 if isinstance(df.columns, pd.MultiIndex):
136 metadata["field_names"] = df.columns.get_level_values(0).tolist()
137 if df.columns.nlevels > 1:
138 metadata["field_types"] = df.columns.get_level_values(1).tolist()
140 return metadata
142 def create_records_with_metadata(
143 self,
144 df: pd.DataFrame,
145 base_records: List[Record],
146 metadata: Optional[Dict[str, Any]] = None
147 ) -> List[Record]:
148 """Create records with preserved metadata.
150 Args:
151 df: Source DataFrame
152 base_records: Base records from conversion
153 metadata: Additional metadata
155 Returns:
156 Records with metadata
157 """
158 if not metadata:
159 metadata = self.extract_metadata_from_dataframe(df)
161 # Apply record IDs if preserved
162 if "record_ids" in metadata and len(metadata["record_ids"]) == len(base_records):
163 for record, record_id in zip(base_records, metadata["record_ids"]):
164 if record_id:
165 record.id = record_id
167 # Apply record metadata if present
168 if "record_metadata" in metadata:
169 record_meta = metadata["record_metadata"]
170 for i, record in enumerate(base_records):
171 if i < len(record_meta) and record_meta[i]:
172 record.metadata = record_meta[i]
174 # Apply field metadata if present
175 if "field_metadata" in metadata:
176 field_meta = metadata["field_metadata"]
177 for record in base_records:
178 for field_name, field in record.fields.items():
179 if field_name in field_meta:
180 field.metadata = field_meta[field_name]
182 return base_records
184 def _get_all_field_names(self, records: List[Record]) -> List[str]:
185 """Get all unique field names from records."""
186 field_names = set()
187 for record in records:
188 field_names.update(record.fields.keys())
189 return sorted(field_names)
191 def _get_field_types(self, records: List[Record]) -> Dict[str, str]:
192 """Get field types from records."""
193 field_types = {}
194 for record in records:
195 for field_name, field in record.fields.items():
196 if field_name not in field_types and field.type:
197 field_types[field_name] = field.type.value
198 return field_types
200 def _extract_record_metadata(self, records: List[Record]) -> List[Dict[str, Any]]:
201 """Extract metadata from each record."""
202 return [r.metadata if r.metadata else {} for r in records]
204 def _extract_field_metadata(self, records: List[Record]) -> Dict[str, Dict[str, Any]]:
205 """Extract metadata from fields."""
206 field_metadata = {}
207 for record in records:
208 for field_name, field in record.fields.items():
209 if field.metadata and field_name not in field_metadata:
210 field_metadata[field_name] = field.metadata
211 return field_metadata
213 def _get_record_metadata_columns(self, records: List[Record]) -> Dict[str, List]:
214 """Get record metadata as column data."""
215 columns = {}
217 # Collect all metadata keys
218 all_keys = set()
219 for record in records:
220 if record.metadata:
221 all_keys.update(record.metadata.keys())
223 # Create column for each metadata key
224 for key in all_keys:
225 values = []
226 for record in records:
227 value = record.metadata.get(key) if record.metadata else None
228 values.append(value)
229 columns[key] = values
231 return columns
233 def clean_dataframe_columns(self, df: pd.DataFrame) -> pd.DataFrame:
234 """Remove metadata columns from DataFrame.
236 Args:
237 df: DataFrame to clean
239 Returns:
240 DataFrame without metadata columns
241 """
242 if self.config.strategy == MetadataStrategy.COLUMNS:
243 # Remove metadata columns
244 meta_cols = [col for col in df.columns if col.startswith(self.config.metadata_prefix)]
245 return df.drop(columns=meta_cols)
247 elif self.config.strategy == MetadataStrategy.MULTI_INDEX:
248 # Flatten multi-index to single level
249 if isinstance(df.columns, pd.MultiIndex):
250 df.columns = df.columns.get_level_values(0)
252 return df