如何在Polars/PyArrow中高效实现IP归属网络查询？

Polars 内置高效匹配函数 #

Polars 从0.19.0版本开始提供了ipv4_lookup内置函数，完全对标KQL的ipv4-lookup功能，无需逐行循环，直接通过向量化操作完成大表匹配，性能比apply+ipaddress提升数个数量级。

代码示例 #

import polars as pl

# 客户端IP表
client_df = pl.DataFrame({
    "ClientIP": ["99.87.29.96", "10.0.0.5", "192.168.1.100"]
})

# 网段表（需包含网段列，可自定义列名）
network_df = pl.DataFrame({
    "network": ["99.96.0.0/13", "10.0.0.0/8", "192.168.0.0/16"],
    "network_desc": ["公共服务网段", "内部办公网段", "家庭局域网网段"]
})

# 执行匹配：自动匹配最长前缀网段，支持左连接/内连接等
result_df = client_df.join(
    network_df,
    how="left",
    on=pl.ipv4_lookup("ClientIP", network_df["network"])
)

print(result_df)

该函数会自动将IP和网段转换为底层高效的数值类型，批量完成匹配逻辑，同时默认优先匹配最长前缀的网段（若存在多个包含该IP的网段），完全适配大表场景。

PyArrow 实现方案 #

PyArrow暂无直接的ipv4-lookup内置函数，但可通过IP转整数+位运算的方式实现向量化匹配，避免Python级别的逐行循环，大幅提升性能。

核心思路 #

1. 将IPv4地址转换为32位整数（利用pyarrow.compute.ipv4_to_int32）
2. 对网段拆分出网络地址和前缀长度，计算出对应的网络整数和掩码整数
3. 通过按位与运算，判断IP整数与掩码的结果是否等于网段的网络整数

代码示例 #

import pyarrow as pa
import pyarrow.compute as pc

# 1. 处理客户端IP表：转换为整数
client_table = pa.table({"ClientIP": ["99.87.29.96", "10.0.0.5", "192.168.1.100"]})
client_table = client_table.append_column(
    "ip_int", pc.ipv4_to_int32(client_table["ClientIP"])
)

# 2. 处理网段表：计算网络整数、掩码
network_table = pa.table({
    "network": ["99.96.0.0/13", "10.0.0.0/8", "192.168.0.0/16"],
    "network_desc": ["公共服务网段", "内部办公网段", "家庭局域网网段"]
})

# 拆分网段为地址和前缀长度
split_net = pc.split_pattern(network_table["network"], "/")
network_table = network_table.append_column("net_addr", split_net.field(0))
network_table = network_table.append_column("prefix_len", pc.cast(split_net.field(1), pa.int32()))

# 计算网段的网络整数和掩码
network_table = network_table.append_column(
    "net_int", pc.ipv4_to_int32(network_table["net_addr"])
)
# 计算32位掩码：(0xFFFFFFFF << (32 - 前缀长度)) & 0xFFFFFFFF
mask = pc.bit_shift_left(pc.scalar(0xFFFFFFFF, pa.uint32()), pc.subtract(32, network_table["prefix_len"]))
mask = pc.bit_and(mask, pc.scalar(0xFFFFFFFF, pa.uint32()))
network_table = network_table.append_column("mask", mask)

# 3. 批量匹配：这里以笛卡尔积后过滤为例（大表建议结合分区或Polars协同处理）
cross_table = pc.cross_join(client_table, network_table)
matched = pc.equal(
    pc.bit_and(cross_table["ip_int"], cross_table["mask"]),
    cross_table["net_int"]
)
result_table = cross_table.filter(matched)

# 可选：去重，保留最长前缀的匹配（按prefix_len降序后去重）
result_table = pc.sort(result_table, sort_keys=[("prefix_len", "descending")])
result_table = pc.drop_duplicates(result_table, subset=["ClientIP"])

print(result_table.to_pandas())

通用优化建议 #

• 预计算网段元数据：提前计算网段的网络整数、掩码，避免重复计算
• 优先匹配最长前缀：将网段按前缀长度从长到短排序，确保精确匹配优先（Polars的ipv4_lookup已默认实现）
• 避免全量笛卡尔积：大表场景下优先使用Polars的内置匹配函数，或结合PyArrow的分区操作减少计算量

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