当前位置:网站首页>使用pymongo保存数据到MongoDB的工具类
使用pymongo保存数据到MongoDB的工具类
2022-08-08 15:34:00 【呆萌的代Ma】
pymongo 的增删改查请参考:Python连接MongoDB,使用pymongo进行增删改查
在此基础上,引入多线程(pymongo多线程安全),打包成一个工具类,可以直接把一个dataframe保存/更新/覆盖的工具类:
工具类
from pymongo import MongoClient
import pandas as pd
import numpy as np
class MongoHelper(object):
def __init__(self, database, collection=None, host="127.0.0.1", port=27017):
mongo_conn = MongoClient(host=host, port=port)
self.mongo_db = mongo_conn.get_database(database) # 数据库连接(pymongo自带连接池)
self.coll = None
if collection: # 集合
self.coll = self.mongo_db.get_collection(collection)
def set_collection(self, collection):
"""指定链接的集合"""
self.coll = self.mongo_db.get_collection(collection)
def get_coll(self):
return self.coll
def _save_or_update_mongodb(self, record, dict_value, drop=False):
"""根据检查_id,如果存在就覆盖,如果不存在就新增 :param record: 数据库记录 :param dict_value: 待保存的数据 :param drop: 如果有记录,是否先删除后保存 """
if not record: # 不存在,直接插入
self.coll.insert_one(dict_value)
elif drop is True: # 先删除,然后插入,相当于覆盖原始数据
self.coll.delete_one(record)
self.coll.insert_one(dict_value)
else: # 更新
self.coll.update_one(record, {
"$set": dict_value,
})
def get_all_data(self, record_list=["_id"]):
"""获取全部数据,record_list中为需要的数据列名称"""
for _record in self.coll.find({
}, record_list):
yield _record
def save_new_dataframe_to_mongo(self, dataframe: pd.DataFrame):
"""保存dataframe格式的新数据到MongoDB,会自动分配_id,也可以指定_id"""
coll = self.get_coll()
coll.insert_many(dataframe.to_dict(orient="record"))
def update_dataframe(self, dataframe: pd.DataFrame, record_by: list = ["_id"], drop=False):
"""多线程插入dataframe"""
import threading
coll = self.get_coll()
# 剔除dataframe中不存在的列名
record_by = [_r for _r in record_by if _r in dataframe.columns]
thread_list = []
for df_record in dataframe.to_dict(orient="record"):
find_dict = {
}
for _r_by in record_by:
find_dict[_r_by] = df_record[_r_by]
record = coll.find_one(find_dict, {
})
thread = threading.Thread(target=self._save_or_update_mongodb, args=(record, df_record, drop))
thread.start()
thread_list.append(thread)
for _thr in thread_list: # 等待线程全部执行完毕
_thr.join()
def update_dict(self, new_dict_value: dict, record_by: list = ["_id"], drop=False):
"""更新一个dict的数据"""
coll = self.get_coll()
record_by = [_r for _r in record_by if _r in new_dict_value.keys()]
find_dict = {
}
for _r_by in record_by:
find_dict[_r_by] = new_dict_value[_r_by]
record = coll.find_one(find_dict, {
})
self._save_or_update_mongodb(record, new_dict_value, drop=drop)
示例代码
from pymongo import MongoClient
import pandas as pd
import numpy as np
class MongoHelper(object):
def __init__(self, database, collection=None, host="127.0.0.1", port=27017):
mongo_conn = MongoClient(host=host, port=port)
self.mongo_db = mongo_conn.get_database(database) # 数据库连接(pymongo自带连接池)
self.coll = None
if collection: # 集合
self.coll = self.mongo_db.get_collection(collection)
def set_collection(self, collection):
"""指定链接的集合"""
self.coll = self.mongo_db.get_collection(collection)
def get_coll(self):
return self.coll
def _save_or_update_mongodb(self, record, dict_value, drop=False):
"""根据检查_id,如果存在就覆盖,如果不存在就新增 :param record: 数据库记录 :param dict_value: 待保存的数据 :param drop: 如果有记录,是否先删除后保存 """
if not record: # 不存在,直接插入
self.coll.insert_one(dict_value)
elif drop is True: # 先删除,然后插入,相当于覆盖原始数据
self.coll.delete_one(record)
self.coll.insert_one(dict_value)
else: # 更新
self.coll.update_one(record, {
"$set": dict_value,
})
def get_all_data(self, record_list=["_id"]):
"""获取全部数据,record_list中为需要的数据列名称"""
for _record in self.coll.find({
}, record_list):
yield _record
def save_new_dataframe_to_mongo(self, dataframe: pd.DataFrame):
"""保存dataframe格式的新数据到MongoDB,会自动分配_id,也可以指定_id"""
coll = self.get_coll()
coll.insert_many(dataframe.to_dict(orient="record"))
def update_dataframe(self, dataframe: pd.DataFrame, record_by: list = ["_id"], drop=False):
"""多线程插入dataframe"""
import threading
coll = self.get_coll()
# 剔除dataframe中不存在的列名
record_by = [_r for _r in record_by if _r in dataframe.columns]
thread_list = []
for df_record in dataframe.to_dict(orient="record"):
find_dict = {
}
for _r_by in record_by:
find_dict[_r_by] = df_record[_r_by]
record = coll.find_one(find_dict, {
})
thread = threading.Thread(target=self._save_or_update_mongodb, args=(record, df_record, drop))
thread.start()
thread_list.append(thread)
for _thr in thread_list: # 等待线程全部执行完毕
_thr.join()
def update_dict(self, new_dict_value: dict, record_by: list = ["_id"], drop=False):
"""更新一个dict的数据"""
coll = self.get_coll()
record_by = [_r for _r in record_by if _r in new_dict_value.keys()]
find_dict = {
}
for _r_by in record_by:
find_dict[_r_by] = new_dict_value[_r_by]
record = coll.find_one(find_dict, {
})
self._save_or_update_mongodb(record, new_dict_value, drop=drop)
def main():
mongo_helper = MongoHelper(database='test_db', collection="test_coll")
df1 = pd.DataFrame(np.random.randn(100, 2), columns=['a', 'b'])
df1["_id"] = [str(i) for i in range(df1.shape[0])]
mongo_helper.save_new_dataframe_to_mongo(df1)
# 改变原始数据
df2 = pd.DataFrame(np.random.randn(100, 2), columns=['a', 'c'])
df2["_id"] = [str(i) for i in range(df2.shape[0])]
mongo_helper.update_dataframe(df2)
if __name__ == '__main__':
main()
如果是
边栏推荐
猜你喜欢
随机推荐
Power BI简介
[Unity entry plan] Unity instance - how to protect data members through encapsulation in C#
一文读懂字节跳动“埋点验证平台”
腾讯又一长达 8 年的服务下架。。。
Create a 2D array
一文搞懂│XSS攻击、SQL注入、CSRF攻击、DDOS攻击、DNS劫持
Mx_yolov3环境配置+模型测试训练
湖北钠斯网络数字藏品交易系统
5G NR RRC连接控制
A16z:为什么 NFT 创作者要选择 cc0?
Notes on the development of kindergarten enrollment registration system based on WeChat applet
CS231n:6 训练神经网络(二)
兆骑科创创业赛事活动举办平台,投融资对接,线上直播路演
如何制作网页
循环神经网络RNN入门介绍
Interview questions 17.05. Letters and numbers
18、学习MySQL ALTER命令
Chat with wine and chat, build an asynchronous non-blocking (aioredis) real-time (websocket) communication chat system based on Vue3.0+Tornado6.1+Redis publish-subscribe (pubsub) mode
【kali-权限提升】(4.2.5)社会工程学工具包:PowerShell攻击向量(防报毒)
leetcode/delete the nth node from the bottom of the linked list









