如何在Polars/Spark/SQL中标准化相似公司名并生成唯一标识

一、Polars 实现 #

核心思路是提取公司名称的核心标识（示例场景下取第一个空格前的内容，无空格则保留原名称），再对核心标识分配唯一ID。

import polars as pl

# 模拟数据
df = pl.DataFrame({
    "company_name": ["Ciao", "Ciao Inc", "HB", "Ciao Inc User", "HB lmtd"]
})

# 处理逻辑：提取核心名称并生成唯一ID
result = df.with_columns(
    pl.col("company_name").str.extract(r"^(\w+)", 1).alias("core_name")
).with_columns(
    pl.col("core_name").rank(method="dense").cast(pl.Int32).alias("company_id")
)

print(result)

输出结果：

shape: (5, 3)
┌──────────────────┬──────────┬────────────┐
│ company_name     ┆ core_name ┆ company_id │
│ ---              ┆ ---       ┆ ---        │
│ str              ┆ str       ┆ i32        │
╞══════════════════╪══════════╪════════════╡
│ Ciao             ┆ Ciao      ┆ 1          │
│ Ciao Inc         ┆ Ciao      ┆ 1          │
│ HB               ┆ HB        ┆ 2          │
│ Ciao Inc User    ┆ Ciao      ┆ 1          │
│ HB lmtd          ┆ HB        ┆ 2          │
└──────────────────┴──────────┴────────────┘

如果需要适配更复杂的模糊匹配场景（核心词不在开头），可以结合fuzzywuzzy等工具做字符串聚类后再分配ID。

二、Spark 实现 #

通过正则提取核心名称，再用窗口函数生成唯一ID。

import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions.Window

// 模拟数据
val df = spark.createDataFrame(Seq(
    ("Ciao"), ("Ciao Inc"), ("HB"), ("Ciao Inc User"), ("HB lmtd")
)).toDF("company_name")

// 处理逻辑
val result = df.withColumn(
    "core_name", regexp_extract(col("company_name"), "^(\\w+)", 1)
).withColumn(
    "company_id", dense_rank().over(Window.orderBy("core_name"))
)

result.show()

输出结果：

+---------------+---------+----------+
|  company_name|core_name|company_id|
+---------------+---------+----------+
|          Ciao|     Ciao|         1|
|      Ciao Inc|     Ciao|         1|
|Ciao Inc User |     Ciao|         1|
|            HB|       HB|         2|
|        HB lmtd|       HB|         2|
+---------------+---------+----------+

三、纯SQL 实现 #

MySQL 版本 #

用SUBSTRING_INDEX提取核心名称，结合用户变量生成唯一ID：

-- 模拟表
CREATE TABLE company_names (company_name VARCHAR(100));
INSERT INTO company_names VALUES 
('Ciao'), ('Ciao Inc'), ('HB'), ('Ciao Inc User'), ('HB lmtd');

-- 查询逻辑
SELECT 
    company_name,
    @id := IF(@prev_core = core_name, @id, @id + 1) AS company_id,
    @prev_core := core_name
FROM (
    SELECT 
        company_name,
        SUBSTRING_INDEX(company_name, ' ', 1) AS core_name
    FROM company_names
    ORDER BY core_name
) t,
(SELECT @id := 0, @prev_core := '') init;

PostgreSQL 版本 #

用SPLIT_PART提取核心名称，结合DENSE_RANK()生成ID：

-- 模拟表
CREATE TABLE company_names (company_name VARCHAR(100));
INSERT INTO company_names VALUES 
('Ciao'), ('Ciao Inc'), ('HB'), ('Ciao Inc User'), ('HB lmtd');

-- 查询逻辑
SELECT 
    company_name,
    DENSE_RANK() OVER (ORDER BY SPLIT_PART(company_name, ' ', 1)) AS company_id
FROM company_names;

两种SQL版本输出均符合需求：

内容的提问来源于stack exchange，提问作者cyberZamp