Spaces:
Running
Running
File size: 4,470 Bytes
09801ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import pandas as pd
from typing import List, Dict, Any, Tuple
import logging
from utils.paths import get_user_paths
logger = logging.getLogger(__name__)
class MultiDatasetEngine:
"""Intelligent engine for analyzing and joining multiple datasets"""
@staticmethod
def infer_joins(df1: pd.DataFrame, df2: pd.DataFrame, df1_name: str, df2_name: str) -> List[Dict[str, Any]]:
"""
Intelligently infer possible join keys between two dataframes based on column names and content.
Returns a list of suggested joins ranked by confidence.
"""
suggestions = []
cols1 = set(df1.columns)
cols2 = set(df2.columns)
# 1. Exact match on ID-like columns
id_keywords = ['id', 'uuid', 'key']
exact_matches = cols1.intersection(cols2)
for col in exact_matches:
if any(kw in col.lower() for kw in id_keywords):
# Verify type compatibility
if df1[col].dtype == df2[col].dtype:
# Calculate overlap
overlap = len(set(df1[col]).intersection(set(df2[col])))
if overlap > 0:
confidence = 0.95 if overlap > len(df1) * 0.1 or overlap > len(df2) * 0.1 else 0.7
suggestions.append({
"left_key": col,
"right_key": col,
"confidence": confidence,
"reason": f"Exact match on ID column '{col}' with data overlap."
})
# 2. Heuristic match (e.g. 'customer_id' in df1, 'id' in df2 where df2 is 'customers')
for col1 in cols1:
for col2 in cols2:
# If col1 is like 'table2_id' and col2 is 'id'
if col1.lower() == f"{df2_name.lower().replace('.csv', '')}_id" and col2.lower() == 'id':
overlap = len(set(df1[col1]).intersection(set(df2[col2])))
if overlap > 0:
suggestions.append({
"left_key": col1,
"right_key": col2,
"confidence": 0.85,
"reason": f"Inferred foreign key relationship '{col1}' -> '{col2}'."
})
# If col2 is like 'table1_id' and col1 is 'id'
elif col2.lower() == f"{df1_name.lower().replace('.csv', '')}_id" and col1.lower() == 'id':
overlap = len(set(df1[col1]).intersection(set(df2[col2])))
if overlap > 0:
suggestions.append({
"left_key": col1,
"right_key": col2,
"confidence": 0.85,
"reason": f"Inferred foreign key relationship '{col1}' <- '{col2}'."
})
# Sort by confidence descending
suggestions.sort(key=lambda x: x['confidence'], reverse=True)
return suggestions
@staticmethod
def join_datasets(
df1: pd.DataFrame,
df2: pd.DataFrame,
left_on: str,
right_on: str,
how: str = 'left'
) -> pd.DataFrame:
"""
Execute the join between two datasets.
Handles duplicates and renaming intelligently.
"""
# Ensure compatible types for joining
if df1[left_on].dtype != df2[right_on].dtype:
try:
# Attempt to coerce to string for joining if types mismatch
df1[left_on] = df1[left_on].astype(str)
df2[right_on] = df2[right_on].astype(str)
except Exception as e:
logger.warning(f"Could not coerce types for join: {e}")
# Perform merge
merged = pd.merge(df1, df2, left_on=left_on, right_on=right_on, how=how, suffixes=('_left', '_right'))
# If right_on is different from left_on, we might have redundant columns, but pandas merge handles it.
# If they are different names, both are kept. We can drop the right_on if it's redundant.
if left_on != right_on and right_on in merged.columns:
# Only drop if the values are identical (as in inner/left join)
if how in ['left', 'inner']:
merged = merged.drop(columns=[right_on])
return merged
|