PM-pinou
351 lines
13 KiB
Python
351 lines
13 KiB
Python
"""Data type classifier for column-level type inference.
|
|
|
|
Provides :class:`DataType` enumeration and functions to classify Polars
|
|
columns based on name heuristics and value sampling.
|
|
|
|
Typical usage::
|
|
|
|
from analysis.type_classifier import classify_column, classify_schema, DataType
|
|
|
|
# Single column
|
|
dtype = classify_column('src_ip', series, config_type='ipv4')
|
|
|
|
# Full schema
|
|
type_map = classify_schema(lazyframe, config_overrides={'src_ip': 'ipv4'})
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from enum import Enum
|
|
|
|
import polars as pl
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DataType enumeration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class DataType(Enum):
|
|
"""Enumerated column types recognised by the classifier engine."""
|
|
|
|
INT = 1
|
|
FLOAT = 2
|
|
ENUM = 3
|
|
HEX = 4
|
|
URL = 5
|
|
IPv4 = 6
|
|
LAT_LON = 7
|
|
TIMESTAMP = 8
|
|
BOOL_ENUM = 9
|
|
STRING = 10
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Detection patterns (raw strings for regex)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
HEX_PATTERN = r'^[0-9a-f]{2}(?:\s[0-9a-f]{2})*$'
|
|
"""Matches ``"0e a6 3f"``, ``"ab cd ef"`` etc. (space-separated hex pairs)."""
|
|
|
|
URL_PATTERN = r'^https?://'
|
|
"""Matches URLs starting with ``http://`` or ``https://``."""
|
|
|
|
IPv4_PATTERN = r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$'
|
|
"""Matches ``"192.168.1.1"`` style IPv4 addresses."""
|
|
|
|
MAC_PATTERN = r'^([0-9A-Fa-f]{2}[:.-]){5}[0-9A-Fa-f]{2}$'
|
|
"""Matches MAC addresses with ``:``, ``-``, or ``.`` separators."""
|
|
|
|
BOOL_ENUM_VALUES = (
|
|
'', ' ', '+', '-', '0', '1',
|
|
'true', 'false', 'yes', 'no', 't', 'f',
|
|
)
|
|
"""String values considered boolean / binary indicators."""
|
|
|
|
LAT_LON_NAMES = ('lat', 'latitude', 'lon', 'lng', 'longitude')
|
|
|
|
# Name-to-type heuristics (checked after config override, before values)
|
|
_NAME_TYPE_MAP: list[tuple[re.Pattern, DataType]] = [
|
|
# User column names - EXACT match only (no generic fuzzy patterns)
|
|
# IPv4 columns
|
|
(re.compile(r'^:ips$|^:ipd$|^server.ip$|^client.ip$', re.I), DataType.IPv4),
|
|
# ENUM columns
|
|
(re.compile(r'^:prs$|^:prd$', re.I), DataType.ENUM),
|
|
(re.compile(r'^(?:scnt|dcnt|cnrs|isrs|1ipp|4dbn)$', re.I), DataType.ENUM),
|
|
(re.compile(r'^0ver$', re.I), DataType.ENUM),
|
|
# FLOAT columns
|
|
(re.compile(r'^4dur$|^8ses$|^2tmo$|^0rnt$', re.I), DataType.FLOAT),
|
|
# INT columns
|
|
(re.compile(r'^4ksz$|^8ack$|^8ppk$|^row$', re.I), DataType.INT),
|
|
# HEX columns
|
|
(re.compile(r'^0cph$|^0crv$', re.I), DataType.HEX),
|
|
# TIMESTAMP columns
|
|
(re.compile(r'^8dbd$|^time$|^timestamp$', re.I), DataType.TIMESTAMP),
|
|
# Dot-notation suffixes for lat/lon (e.g. :ips.latd, :ipd.lond)
|
|
(re.compile(r'\.latd$', re.I), DataType.LAT_LON),
|
|
(re.compile(r'\.lond$', re.I), DataType.LAT_LON),
|
|
# Dot-notation STRING suffixes (e.g. :ips.ispn, :ipd.city)
|
|
(re.compile(r'\.(ispn|orgn|city)$', re.I), DataType.STRING),
|
|
# Standalone string columns - exact match
|
|
(re.compile(r'^(?:cnam|snam|tabl|name|0rnd|source.node|ecdhe.named.curve)$', re.I), DataType.STRING),
|
|
]
|
|
|
|
# Mapping from config.yaml type strings to DataType
|
|
_CONFIG_TYPE_MAP: dict[str, DataType] = {
|
|
'int': DataType.INT,
|
|
'float': DataType.FLOAT,
|
|
'enum': DataType.ENUM,
|
|
'hex': DataType.HEX,
|
|
'url': DataType.URL,
|
|
'ipv4': DataType.IPv4,
|
|
'lat_lon': DataType.LAT_LON,
|
|
'timestamp': DataType.TIMESTAMP,
|
|
'bool_enum': DataType.BOOL_ENUM,
|
|
'string': DataType.STRING,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Column classification (single column)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _name_to_type(name: str) -> DataType | None:
|
|
"""Match column name against known heuristics.
|
|
|
|
Returns a :class:`DataType` if the name matches strongly, otherwise
|
|
``None`` (fall through to value-based detection).
|
|
"""
|
|
for pattern, dtype in _NAME_TYPE_MAP:
|
|
if pattern.search(name):
|
|
return dtype
|
|
return None
|
|
|
|
|
|
def _is_port_value(v: str) -> bool:
|
|
"""Check if *v* is an integer string in the registered port range (1-65535)."""
|
|
try:
|
|
n = int(v)
|
|
return 1 <= n <= 65535
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def _can_float(v: str) -> bool:
|
|
"""Return ``True`` when *v* can be parsed as a Python float."""
|
|
try:
|
|
float(v)
|
|
return True
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def _values_to_type(series: pl.Series) -> DataType:
|
|
"""Infer column type by sampling non-null values.
|
|
|
|
Detection order: IPv4 → MAC guard → HEX → URL → PORT → BOOL_ENUM
|
|
→ LAT_LON → ENUM → FLOAT → STRING
|
|
"""
|
|
dtype = series.dtype
|
|
|
|
# Boolean dtype ──────────────────────────────────────────────────────
|
|
if dtype == pl.Boolean:
|
|
return DataType.BOOL_ENUM
|
|
|
|
# Numeric dtypes ─────────────────────────────────────────────────────
|
|
if dtype in (pl.Float32, pl.Float64):
|
|
# Check LAT_LON possibility before returning FLOAT
|
|
non_null = series.drop_nulls()
|
|
if len(non_null) > 0:
|
|
sample = non_null.head(100)
|
|
sample_strs = [str(v).strip().lower() for v in sample]
|
|
non_blank = [v for v in sample_strs if v not in ('', ' ', '+', '-')]
|
|
if non_blank:
|
|
try:
|
|
floats = [float(v) for v in non_blank]
|
|
max_abs = max(abs(f) for f in floats)
|
|
if (
|
|
all(-180 <= f <= 180 for f in floats)
|
|
and any('.' in v for v in non_blank)
|
|
and max_abs > 10
|
|
):
|
|
return DataType.LAT_LON
|
|
except ValueError:
|
|
pass
|
|
return DataType.FLOAT
|
|
|
|
if dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
|
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64):
|
|
return DataType.INT
|
|
|
|
# String / categorical types ─────────────────────────────────────────
|
|
non_null = series.drop_nulls()
|
|
if len(non_null) == 0:
|
|
return DataType.STRING
|
|
|
|
sample = non_null.head(100)
|
|
# Convert all to string representations for pattern matching
|
|
sample_strs = [str(v).strip().lower() for v in sample]
|
|
|
|
# ── 1. IPv4 detection ──────────────────────────────────────────────
|
|
if all(re.match(IPv4_PATTERN, v) for v in sample_strs):
|
|
if all(_valid_ip_octets(v) for v in sample_strs):
|
|
return DataType.IPv4
|
|
|
|
# ── 2. MAC guard (skip HEX if values look like MAC addresses) ──────
|
|
is_mac = all(re.match(MAC_PATTERN, v) for v in sample_strs)
|
|
|
|
# ── 3. HEX detection ───────────────────────────────────────────────
|
|
if not is_mac and all(re.match(HEX_PATTERN, v) for v in sample_strs):
|
|
return DataType.HEX
|
|
|
|
# ── 4. URL detection ───────────────────────────────────────────────
|
|
if all(re.match(URL_PATTERN, v) for v in sample_strs):
|
|
return DataType.URL
|
|
|
|
# ── 5. PORT detection → ENUM ───────────────────────────────────────
|
|
if all(_is_port_value(v) for v in sample_strs):
|
|
return DataType.ENUM
|
|
|
|
# ── 6. BOOL_ENUM detection ─────────────────────────────────────────
|
|
if all(v in BOOL_ENUM_VALUES for v in sample_strs):
|
|
return DataType.BOOL_ENUM
|
|
|
|
# ── 7. LAT_LON detection (value-based) ──────────────────────────────
|
|
non_blank = [v for v in sample_strs if v not in ('', ' ', '+', '-')]
|
|
if non_blank:
|
|
try:
|
|
floats = [float(v) for v in non_blank]
|
|
max_abs = max(abs(f) for f in floats)
|
|
if (
|
|
all(-180 <= f <= 180 for f in floats)
|
|
and any('.' in v for v in non_blank)
|
|
and max_abs > 10
|
|
):
|
|
return DataType.LAT_LON
|
|
except ValueError:
|
|
pass
|
|
|
|
# ── 8. ENUM detection (few unique values) ──────────────────────────
|
|
unique_count = non_null.n_unique()
|
|
if unique_count < 20:
|
|
return DataType.ENUM
|
|
|
|
# ── 9. FLOAT detection (>80 % can float) ───────────────────────────
|
|
float_count = sum(1 for v in sample_strs if _can_float(v))
|
|
if float_count / len(sample_strs) > 0.8:
|
|
# Guard: values with leading "+"/"-" (e.g., "+5", "-abc") that
|
|
# happen to parse as float should not force a FLOAT classification
|
|
# when non-float values exist in the column. This catches string
|
|
# columns where a "+" prefix is used as a text marker (phone codes,
|
|
# IDs, error codes) rather than a numeric sign.
|
|
if any(v.startswith('+') or v.startswith('-') for v in sample_strs):
|
|
if float_count < len(sample_strs):
|
|
return DataType.STRING
|
|
return DataType.FLOAT
|
|
|
|
return DataType.STRING
|
|
|
|
|
|
def _valid_ip_octets(ip_str: str) -> bool:
|
|
"""Check that each octet in an IPv4 string is 0-255."""
|
|
parts = ip_str.split('.')
|
|
if len(parts) != 4:
|
|
return False
|
|
for p in parts:
|
|
try:
|
|
n = int(p)
|
|
if n < 0 or n > 255:
|
|
return False
|
|
except ValueError:
|
|
return False
|
|
return True
|
|
|
|
|
|
def classify_column(
|
|
name: str,
|
|
series: pl.Series,
|
|
config_type: str | None = None,
|
|
) -> DataType:
|
|
"""Classify a single column based on its *name* and *series* values.
|
|
|
|
Priority
|
|
--------
|
|
1. **config_type** — user override from ``config.yaml``.
|
|
If provided (e.g. ``'ipv4'``), returns the corresponding
|
|
:class:`DataType` immediately.
|
|
2. **Value-based** pattern matching — samples up to 100 non-null
|
|
values and tests against IPv4 → HEX → URL → PORT → BOOL_ENUM
|
|
→ LAT_LON → ENUM → FLOAT in order.
|
|
3. **Name-based** heuristics — only consulted when value-based
|
|
detection returned ``STRING`` (i.e. no value pattern matched).
|
|
|
|
Parameters
|
|
----------
|
|
name:
|
|
Column name (used for name-based heuristics).
|
|
series:
|
|
Polars :class:`Series` containing the column's data.
|
|
config_type:
|
|
Optional user override (case-insensitive, matched to
|
|
:data:`_CONFIG_TYPE_MAP`).
|
|
|
|
Returns
|
|
-------
|
|
DataType
|
|
"""
|
|
# 1. Config override (highest priority)
|
|
if config_type is not None:
|
|
mapped = _CONFIG_TYPE_MAP.get(config_type.lower().strip())
|
|
if mapped is not None:
|
|
return mapped
|
|
|
|
# 2. Value-based sampling
|
|
result = _values_to_type(series)
|
|
|
|
# 3. Name-based heuristics (only when values returned STRING)
|
|
if result == DataType.STRING:
|
|
name_match = _name_to_type(name)
|
|
if name_match is not None:
|
|
return name_match
|
|
|
|
return result
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Schema classification (all columns)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def classify_schema(
|
|
lf: pl.LazyFrame,
|
|
config_overrides: dict[str, str] | None = None,
|
|
) -> dict[str, DataType]:
|
|
"""Classify all columns in a :class:`LazyFrame`.
|
|
|
|
Parameters
|
|
----------
|
|
lf:
|
|
A Polars :class:`LazyFrame` whose columns to classify.
|
|
config_overrides:
|
|
Optional mapping of ``{col_name: type_string}`` from a
|
|
``config.yaml`` ``columns:`` section. These take highest
|
|
priority (see :func:`classify_column`).
|
|
|
|
Returns
|
|
-------
|
|
dict[str, DataType]
|
|
Mapping from column name to inferred :class:`DataType`.
|
|
"""
|
|
config_overrides = config_overrides or {}
|
|
df_sample = lf.collect(streaming=True)
|
|
result: dict[str, DataType] = {}
|
|
for col_name in df_sample.columns:
|
|
config_type = config_overrides.get(col_name)
|
|
result[col_name] = classify_column(
|
|
name=col_name,
|
|
series=df_sample[col_name],
|
|
config_type=config_type,
|
|
)
|
|
return result
|