提高 SQLAlchemy 更新效率

Question

我有两个表，用户 (~200.000) 和推文 (~2.000.000) 我需要更新所有用户，包括推文数量、收藏（他们的推文）、回复、转推。 这是在脚本中：

@classmethod
def get_user_tweet_counts(cls, user_id):
    return (db_session
        .query(
            func.sum(Tweet.favorite_count).label('favorite_count'),
            func.sum(Tweet.retweet_count).label('retweet_count'),
            func.sum(Tweet.reply_count).label('reply_count'),
            func.count(Tweet.id).label('tweet_count'))
        .filter(Tweet.user_id == user_id)
        .group_by(Tweet.user_id).first())  # This will always be one result, should I query differently?

db_session：

engine = create_engine('postgresql://tweetsql:tweetsql@127.0.0.1/tweetsql')
db_session = scoped_session(sessionmaker(autocommit=False,
                                         autoflush=True,
                                         bind=engine))
Base = declarative_base()
Base.query = db_session.query_property()

10 分钟循环：

for user in all_users:
    update_count += 1
    aggregation_result = Tweet.get_user_tweet_counts(user.id)
    user.total_tweet_favourites = aggregation_result[0] or 0
    user.total_tweet_retweets = aggregation_result[1] or 0
    user.total_tweet_replies = aggregation_result[2] or 0
    user.tweet_count = aggregation_result[3] or 0
User.save()  # this just calls db_session.commit()
# We only commit the session once to speed things up

User 和 Tweet 声明如下：

User(Base)，Tweet(Base)（来自 db_session sn-p）。

当它运行时，python 达到 80% 的 cpu 和 ~600mb 的内存。我怎样才能使它变得更好？ Tweet 在 user_id 和它自己的 id 上有索引。

Answer 1

Here 是 SQLAlchemy 的作者的一个很好的答案。基本上，如果您需要扩展到大量行，则需要绕过 ORM。

在您的特定情况下，您可以编写单个查询来使用 SQL 聚合实现相同的结果：

UPDATE users SET
  total_tweet_favourites = aggregated.total_tweet_favourites,
  total_tweet_retweets = aggregated.total_tweet_retweets,
  total_tweet_replies = aggregated.total_tweet_replies,
  tweet_count = aggregated.tweet_count
FROM (
  SELECT
    users.id AS id,
    SUM(tweets.favorite_count) AS total_tweet_favourites,
    SUM(tweets.retweet_count) AS total_tweet_retweets,
    SUM(tweets.reply_count) AS total_tweet_replies,
    COUNT(tweets.id) AS tweet_count
  FROM users JOIN tweets ON tweets.user_id = users.id
  GROUP BY users.id
) aggregated
WHERE users.id = aggregated.id;

将其转换为 SQLAlchemy：

aggregated = session \
    .query(
        User.id.label("id"),
        func.sum(Tweet.favorite_count).label("total_tweet_favourites"),
        func.sum(Tweet.retweet_count).label("total_tweet_retweets"),
        func.sum(Tweet.reply_count).label("total_tweet_replies"),
        func.count(Tweet.id).label("tweet_count")) \
    .select_from(User) \
    .join(Tweet) \
    .group_by(User.id) \
    .subquery() \
    .alias("aggregated")
query = User.__table__ \
    .update() \
    .values(
        total_tweet_favourites=aggregated.c.total_tweet_favourites,
        total_tweet_retweets=aggregated.c.total_tweet_retweets,
        total_tweet_replies=aggregated.c.total_tweet_replies,
        tweet_count=aggregated.c.tweet_count) \
    .where(User.__table__.c.id == aggregated.c.id)
session.execute(query)