Postgresへ大量挿入する場合のパフォーマンスを比較してみた

const BUFFER_SIZE: usize = 4096;
struct CopyDataSink<C: DerefMut<Target = PgConnection>> {
    encode_buf: PgArgumentBuffer,
    data_buf: Vec<u8>,
    copy_in: PgCopyIn<C>,
}

type BoxError = Box<dyn std::error::Error + 'static + Send + Sync>;

impl<C: DerefMut<Target = PgConnection>> CopyDataSink<C> {
    fn new(copy_in: PgCopyIn<C>) -> Self {
        let mut data_buf = Vec::with_capacity(BUFFER_SIZE * 2);
        const COPY_SIGNATURE: &[u8] = &[
            b'P', b'G', b'C', b'O', b'P', b'Y', b'\n', // "PGCOPY\n"
            0xFF,  // \377 (8進数) = 0xFF (16進数)
            b'\r', b'\n', // "\r\n"
            0x00,  // \0
        ];

        assert_eq!(COPY_SIGNATURE.len(), 11);
        data_buf.extend_from_slice(COPY_SIGNATURE); // 署名
        data_buf.extend_from_slice(&0_i32.to_be_bytes()); // フラグフィールド
        data_buf.extend_from_slice(&0_i32.to_be_bytes()); // ヘッダ拡張領域長

        CopyDataSink {
            encode_buf: Default::default(),
            data_buf,
            copy_in,
        }
    }

    async fn send(&mut self) -> Result<(), BoxError> {
        let _copy_in = self.copy_in.send(self.data_buf.as_slice()).await?;

        self.data_buf.clear();
        Ok(())
    }

    async fn finish(mut self) -> Result<u64, BoxError> {
        const COPY_TRAILER: &[u8] = &(-1_i16).to_be_bytes();

        self.data_buf.extend(COPY_TRAILER);
        self.send().await?;
        self.copy_in.finish().await.map_err(|e| e.into())
    }

    fn insert_row(&mut self) {
        let num_col = self.copy_in.num_columns() as i16;
        self.data_buf.extend(num_col.to_be_bytes());
    }

    async fn add<'q, T>(&mut self, value: T) -> Result<(), BoxError>
    where
        T: sqlx::Encode<'q, Postgres> + sqlx::Type<Postgres>,
    {
        let is_null = value.encode_by_ref(&mut self.encode_buf)?;

        match is_null {
            sqlx::encode::IsNull::Yes => {
                self.data_buf.extend((-1_i32).to_be_bytes());
            }
            sqlx::encode::IsNull::No => {
                self.data_buf
                    .extend((self.encode_buf.len() as i32).to_be_bytes());
                self.data_buf.extend_from_slice(self.encode_buf.as_slice());
            }
        }

        self.encode_buf.clear();

        if self.data_buf.len() > BUFFER_SIZE {
            self.send().await?;
        }

        Ok(())
    }
}

async fn buffered_copy_in<F, C>(generator: impl Iterator<Item = Author>, coverter: C, mut f: F)
where
    C: Fn(Author) -> String,
    F: AsyncFnMut(bytes::Bytes),
{
    let mut buf = bytes::BytesMut::with_capacity(BUFFER_SIZE * 2);
    for author in generator {
        let data = coverter(author);
        buf.extend_from_slice(data.as_bytes());

        if buf.len() >= BUFFER_SIZE {
            let bytes = buf.split().freeze();
            f(bytes).await;
        }
    }
    if !buf.is_empty() {
        let bytes = buf.split().freeze();
        f(bytes).await;
    }
}

    let chunk_size = 10000;
    assert!(count % chunk_size == 0);

    let mut query = String::from("INSERT INTO authors (id, name, bio) VALUES ");
    let parameters = (0..chunk_size)
        .map(|i| format!("(${}, ${}, ${})", i * 3 + 1, i * 3 + 2, i * 3 + 3))
        .collect::<Vec<_>>()
        .join(", ");
    query.push_str(&parameters);