Golang中如何为SQL rows.Next迭代预分配切片？

Golang中如何为SQL rows.Next迭代预分配切片？ 你好，

我希望通过预分配切片（var posts []*PostModel）或其他方式来提升性能，但在不知道 rows 包含多少条记录的情况下，不确定该如何操作。因此，我不得不使用 posts = append(posts, post)。有人能给我一些建议吗？

谢谢

注意：我在所有 RDBMS 包中都遇到了同样的问题，例如 sql/database、MySQL、Postgres pg、pgx 等等。

func (s Storage) ListPostsByUser(ctx context.Context, args ListPostsByUserArgs) ([]*PostModel, error) {
	ctx, cancel := context.WithTimeout(ctx, s.timeout)
	defer cancel()

	qry := fmt.Sprintf(`
	SELECT
	id, user_id, text, created_at, deleted_at
	FROM posts
	WHERE user_id = $1
	ORDER BY %s
	LIMIT $2
	OFFSET $3
	`,
		args.OrderBy,
	)

	rows, err := s.Query(ctx, qry, args.UserID, args.Limit, args.Offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// START: Try to enhance starting from here --------------------------------
	var posts []*PostModel

	for rows.Next() {
		post := &PostModel{}

		err := rows.Scan(
			&post.ID,
			&post.UserID,
			&post.Text,
			&post.CreatedAt,
			&post.DeletedAt,
		)
		if err != nil {
			return nil, err
		}

		posts = append(posts, post)
	}
	// END: End of enhancement -------------------------------------------------

	if err := rows.Err(); err != nil {
		return nil, err
	}

	return posts, nil
}

更多关于Golang中如何为SQL rows.Next迭代预分配切片？的实战教程也可以访问 https://www.itying.com/category-94-b0.html

4 回复

我尝试了所有实现方案，但最终结果并不合理，无法继续推进，因为我最终需要处理/从切片中移除空对象，因为我并不总是能从数据库获取到 args.Limit 指定的数量。我认为最简洁、最安全的方法就是坚持使用我现有的方案。感谢两位的解答。

更多关于Golang中如何为SQL rows.Next迭代预分配切片？的实战系列教程也可以访问 https://www.itying.com/category-94-b0.html


我不确定是否理解了你的情况，但如果你正在执行分页查询，它应该从数据库中获取“limit”条记录，所以你的帖子切片最多只能有“limit”的大小。 因此，你可以这样做：

posts := make([]*PostModel, args.Limit)
index :=0
for rows.Next() {
    ...
    posts[index++] = post
}

这篇 Stack Overflow 文章给出了我认为你仅有的几个实际选项。

我认为没有任何内置的方法可以获取行数。另一个想法是你可以直接猜测。如果你知道至少会返回 X 条记录……就预先分配那么多空间。如果你猜多了，只需对你的切片进行子切片操作。如果你猜少了，就追加你缺少的数量（这虽然不理想，但比每次都追加要好！）。

也许你可以定期（在另一个线程中）运行一个计数数据库查询来得出你的猜测。

我还总是建议进行一些微基准测试或对你的应用程序进行性能分析！Append 操作相当快——除非有数百万条记录，否则查询本身很可能比追加操作花费的时间长得多。

在Golang中，可以通过rows.EstimatedRowCount()（pgx）或查询COUNT(*)来预分配切片。以下是两种实现方式：

方法1：使用COUNT(*)预查询（通用方法）

func (s Storage) ListPostsByUser(ctx context.Context, args ListPostsByUserArgs) ([]*PostModel, error) {
    ctx, cancel := context.WithTimeout(ctx, s.timeout)
    defer cancel()

    // 先查询总数
    countQuery := `
        SELECT COUNT(*) 
        FROM posts 
        WHERE user_id = $1
    `
    var total int
    err := s.QueryRow(ctx, countQuery, args.UserID).Scan(&total)
    if err != nil {
        return nil, err
    }

    // 根据实际需要返回的数量预分配切片
    expectedRows := args.Limit
    if total < args.Limit {
        expectedRows = total
    }
    
    // 预分配切片
    posts := make([]*PostModel, 0, expectedRows)

    qry := fmt.Sprintf(`
        SELECT id, user_id, text, created_at, deleted_at
        FROM posts
        WHERE user_id = $1
        ORDER BY %s
        LIMIT $2
        OFFSET $3
    `, args.OrderBy)

    rows, err := s.Query(ctx, qry, args.UserID, args.Limit, args.Offset)
    if err != nil {
        return nil, err
    }
    defer rows.Close()

    for rows.Next() {
        post := &PostModel{}
        err := rows.Scan(
            &post.ID,
            &post.UserID,
            &post.Text,
            &post.CreatedAt,
            &post.DeletedAt,
        )
        if err != nil {
            return nil, err
        }
        posts = append(posts, post)
    }

    if err := rows.Err(); err != nil {
        return nil, err
    }

    return posts, nil
}

方法2：使用pgx的EstimatedRowCount()（pgx特定）

import "github.com/jackc/pgx/v5"

func (s Storage) ListPostsByUserPgx(ctx context.Context, args ListPostsByUserArgs) ([]*PostModel, error) {
    ctx, cancel := context.WithTimeout(ctx, s.timeout)
    defer cancel()

    qry := fmt.Sprintf(`
        SELECT id, user_id, text, created_at, deleted_at
        FROM posts
        WHERE user_id = $1
        ORDER BY %s
        LIMIT $2
        OFFSET $3
    `, args.OrderBy)

    rows, err := s.Query(ctx, qry, args.UserID, args.Limit, args.Offset)
    if err != nil {
        return nil, err
    }
    defer rows.Close()

    // 使用pgx的EstimatedRowCount获取预估行数
    var initialCapacity int
    if pgxRows, ok := rows.(interface{ EstimatedRowCount() int }); ok {
        initialCapacity = pgxRows.EstimatedRowCount()
    }
    
    // 至少分配Limit大小的容量
    if initialCapacity < args.Limit {
        initialCapacity = args.Limit
    }
    
    posts := make([]*PostModel, 0, initialCapacity)

    for rows.Next() {
        post := &PostModel{}
        err := rows.Scan(
            &post.ID,
            &post.UserID,
            &post.Text,
            &post.CreatedAt,
            &post.DeletedAt,
        )
        if err != nil {
            return nil, err
        }
        posts = append(posts, post)
    }

    if err := rows.Err(); err != nil {
        return nil, err
    }

    return posts, nil
}

方法3：基于Limit的保守预分配

func (s Storage) ListPostsByUser(ctx context.Context, args ListPostsByUserArgs) ([]*PostModel, error) {
    ctx, cancel := context.WithTimeout(ctx, s.timeout)
    defer cancel()

    // 直接使用Limit作为容量（最坏情况）
    posts := make([]*PostModel, 0, args.Limit)

    qry := fmt.Sprintf(`
        SELECT id, user_id, text, created_at, deleted_at
        FROM posts
        WHERE user_id = $1
        ORDER BY %s
        LIMIT $2
        OFFSET $3
    `, args.OrderBy)

    rows, err := s.Query(ctx, qry, args.UserID, args.Limit, args.Offset)
    if err != nil {
        return nil, err
    }
    defer rows.Close()

    for rows.Next() {
        post := &PostModel{}
        err := rows.Scan(
            &post.ID,
            &post.UserID,
            &post.Text,
            &post.CreatedAt,
            &post.DeletedAt,
        )
        if err != nil {
            return nil, err
        }
        posts = append(posts, post)
    }

    if err := rows.Err(); err != nil {
        return nil, err
    }

    return posts, nil
}

方法1最准确但需要额外查询，方法3最简单且适用于所有数据库驱动。在大多数分页场景中，方法3以Limit作为预分配容量是最实用的选择。

回到顶部