Quando você tem um arquivo CSV muito grande e que gostaria de uma amostragem de linhas aleatórias, esse script pode te ajudar

Gist ID: 33b266be0830f76cd3200866788404df

import csv
import random
import sys
import os

"""
Quando você tem um arquivo CSV muito grande e que gostaria de uma amostragem de linhas aleatórias,
esse script pode te ajudar.

Gist: https://gist.github.com/gpupo/33b266be0830f76cd3200866788404df

"""

def gerar_amostra_csv(caminho_arquivo_original, caminho_arquivo_amostra, tamanho_amostra):
    """
    Gera uma amostra aleatória de um arquivo CSV.

    Args:
        caminho_arquivo_original (str): O caminho completo para o arquivo CSV original.
        caminho_arquivo_amostra (str): O caminho completo para o novo arquivo CSV de amostra a ser gerado.
        tamanho_amostra (int): O número de linhas que deseja incluir na amostra.
    """
    # Verificar se o arquivo original existe
    if not os.path.isfile(caminho_arquivo_original):
        print(f"O arquivo '{caminho_arquivo_original}' não existe.")
        return

    # Ler o arquivo CSV original e obter o número total de linhas
    with open(caminho_arquivo_original, 'r') as arquivo_original:
        leitor_csv = csv.reader(arquivo_original)
        linhas_originais = list(leitor_csv)
        numero_total_linhas = len(linhas_originais)

    # Verificar se o tamanho da amostra é válido
    if tamanho_amostra > numero_total_linhas:
        print(f"O tamanho da amostra ({tamanho_amostra}) é maior do que o número total de linhas no arquivo original ({numero_total_linhas}).")
        return

    # Ignorar a primeira linha (header)
    indices_amostra = random.sample(range(1, numero_total_linhas), tamanho_amostra)
    
    # Criar o arquivo de amostra e escrever as linhas selecionadas
    with open(caminho_arquivo_amostra, 'w', newline='') as arquivo_amostra:
        escritor_csv = csv.writer(arquivo_amostra)
        
        # Incluir a primeira linha (header) na amostra
        escritor_csv.writerow(linhas_originais[0])

        for indice in indices_amostra:
            escritor_csv.writerow(linhas_originais[indice])

    print(f"Amostra gerada com sucesso e salva em '{caminho_arquivo_amostra}'.")


def main():
    # Verificar se os argumentos foram fornecidos corretamente na linha de comando
    if len(sys.argv) != 4:
        print("Uso: python gerador_amostra_csv.py [caminho_arquivo_original] [caminho_arquivo_amostra] [tamanho_amostra]")
        return

    caminho_origem = sys.argv[1]
    caminho_destino = sys.argv[2]
    quantidade_linhas_amostra = int(sys.argv[3])

    gerar_amostra_csv(caminho_origem, caminho_destino, quantidade_linhas_amostra)


if __name__ == '__main__':
    main()