🤗 datasets与SQL数据库集成：结构化数据处理方案
项目地址: https://gitcode.com/gh_mirrors/da/datasets 你是否还在为机器学习项目中的结构化数据处理而烦恼？当面对SQL数据库中的海量数据时，如何高效地将其转换为模型训练可用的格式？本文将介绍如何使用🤗 datasets库与SQL数据库无缝集成，解决数据加载、处理和转换的全流程问题。读完本文后，你将能够：
- 从SQLite、PostgreSQL等数据库直接加载数据
- 使用SQL查询筛选和预处理数据
- 将处理后的数据集写回数据库
- 利用datasets的并行处理能力加速数据操作
集成方案概述
🤗 datasets提供了完整的SQL数据库集成方案，通过Dataset.from_sql方法可以直接从数据库表或SQL查询结果创建数据集，同时支持将处理后的数据集写回数据库。这种集成方式避免了繁琐的中间文件转换，提高了数据处理效率。
核心组件
- SQLDatasetReader：负责从SQL数据库读取数据，位于src/datasets/io/sql.py
- SQLDatasetWriter：负责将数据集写回SQL数据库，支持批量写入和并行处理
- SQL包装器：支持多种数据库引擎，包括SQLite、PostgreSQL等
快速开始：SQLite数据库集成
SQLite是一个轻量级数据库，非常适合快速原型开发和小型项目。下面我们以SQLite为例，展示如何从数据库加载数据并进行基本操作。
1. 准备数据库
首先创建一个SQLite数据库并插入示例数据：
import sqlite3
import pandas as pd
# 创建数据库连接
conn = sqlite3.connect("us_health_data.db")
# 从CSV文件加载数据到DataFrame
df = pd.read_csv("https://raw.githubusercontent.com/nytimes/health-data/master/us-states.csv")
# 将DataFrame写入SQLite数据库
df.to_sql("states", conn, if_exists="replace", index=False)
conn.close()
2. 从数据库加载数据集
使用Dataset.from_sql方法从SQLite数据库加载数据：
from datasets import Dataset
# 数据库连接URI
uri = "sqlite:///us_health_data.db"
# 从表加载数据
dataset = Dataset.from_sql("states", uri)
# 查看数据集信息
print(dataset)
# Dataset({
# features: ['date', 'state', 'fips', 'cases', 'deaths'],
# num_rows: 54382
# })
3. 使用SQL查询过滤数据
除了加载整个表，还可以直接使用SQL查询加载筛选后的数据：
# 从SQL查询加载数据
query = 'SELECT * FROM states WHERE state="California" AND cases > 10000;'
dataset = Dataset.from_sql(query, uri)
print(dataset)
# Dataset({
# features: ['date', 'state', 'fips', 'cases', 'deaths'],
# num_rows: 1019
# })
高级功能：数据处理与写回
数据过滤与转换
加载数据后，可以使用datasets提供的各种方法进行数据处理：
# 过滤数据
california_dataset = dataset.filter(lambda x: x["state"] == "California")
# 数据转换
def add_case_rate(example):
example["case_rate"] = example["cases"] / example["deaths"] if example["deaths"] > 0 else 0
return example
processed_dataset = california_dataset.map(add_case_rate)
将数据集写回数据库
处理完成后，可以使用SqlDatasetWriter将数据集写回数据库：
from datasets.io.sql import SqlDatasetWriter
# 创建写回器
writer = SqlDatasetWriter(
dataset=processed_dataset,
name="processed_health_data",
con=uri,
batch_size=1000,
num_proc=2 # 并行写入
)
# 执行写回操作
written_rows = writer.write()
print(f"成功写入 {written_rows} 行数据")
多数据库支持
PostgreSQL集成
虽然SQLite适合轻量级应用，但在生产环境中，你可能需要使用更强大的数据库如PostgreSQL。连接方式类似，只需更改连接URI：
# PostgreSQL连接URI
pg_uri = "postgresql://user:password@localhost:5432/mydatabase"
# 从PostgreSQL加载数据
dataset = Dataset.from_sql("SELECT * FROM large_dataset LIMIT 1000", pg_uri)
数据库连接配置
不同数据库的连接URI格式不同，以下是常见数据库的连接格式：
| 数据库 | 连接URI格式 |
|---|---|
| SQLite | sqlite:///database.db |
| PostgreSQL | postgresql://user:password@host:port/dbname |
| MySQL | mysql+pymysql://user:password@host:port/dbname |
| SQL Server | mssql+pyodbc://user:password@dsn_name |
更多数据库连接方式请参考官方文档。
性能优化技巧
批量处理大表
对于大型数据库表，可以使用批量处理减少内存占用：
# 批量读取数据
def batch_query(offset, batch_size=1000):
query = f"SELECT * FROM large_table LIMIT {batch_size} OFFSET {offset};"
return Dataset.from_sql(query, uri)
# 处理多个批次
offset = 0
while True:
batch = batch_query(offset)
if len(batch) == 0:
break
# 处理批次数据
process_batch(batch)
offset += 1000
并行数据处理
利用datasets的并行处理能力加速数据转换：
# 使用多进程处理数据
processed_dataset = dataset.map(
add_case_rate,
num_proc=4, # 使用4个进程
batched=True, # 批量处理
batch_size=1000
)
常见问题解决
数据类型不匹配
如果数据库中的数据类型与datasets推断的类型不匹配，可以显式指定特征类型：
from datasets import Features, Value
# 显式指定特征类型
features = Features({
"date": Value("string"),
"state": Value("string"),
"fips": Value("int32"),
"cases": Value("int64"),
"deaths": Value("int64")
})
# 使用指定的特征类型加载数据
dataset = Dataset.from_sql("states", uri, features=features)
数据库连接超时
对于远程数据库，可以通过设置连接参数解决超时问题：
# PostgreSQL连接带超时参数
pg_uri = "postgresql://user:password@localhost:5432/mydatabase?connect_timeout=10"
总结与展望
🤗 datasets与SQL数据库的集成提供了一种高效、灵活的结构化数据处理方案，特别适合机器学习项目中的数据准备工作。通过直接连接数据库，你可以避免繁琐的数据导出/导入步骤，同时利用datasets强大的数据处理能力和与其他机器学习框架的无缝集成。
未来，datasets将进一步增强数据库集成功能，包括对更多数据库类型的支持、更高效的查询优化和实时数据流处理能力。
相关资源：
希望本文能帮助你更好地利用🤗 datasets处理SQL数据库中的结构化数据。如有任何问题或建议，欢迎在项目仓库提交issue或PR。
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
