114 lines
4.4 KiB
Python
Executable file
114 lines
4.4 KiB
Python
Executable file
#!/usr/bin/env python
|
|
import os
|
|
import pandas as pd
|
|
from tqdm import tqdm
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
|
|
# Define the input and output directories
|
|
INPUT_DIR = "scrapes"
|
|
OUTPUT_DIR = "sorted-and-deduplicated"
|
|
|
|
# Get the project root directory (parent directory of the bin directory)
|
|
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
|
|
INPUT_PATH = os.path.join(PROJECT_ROOT, INPUT_DIR)
|
|
OUTPUT_PATH = os.path.join(PROJECT_ROOT, OUTPUT_DIR)
|
|
CHUNK_SIZE = 10000 # Number of rows per chunk
|
|
|
|
def create_output_dirs():
|
|
if not os.path.exists(OUTPUT_PATH):
|
|
os.makedirs(OUTPUT_PATH)
|
|
for root, dirs, files in os.walk(INPUT_PATH):
|
|
for dir in dirs:
|
|
os.makedirs(os.path.join(OUTPUT_PATH, os.path.relpath(os.path.join(root, dir), INPUT_PATH)))
|
|
|
|
def process_csv_file(file_path):
|
|
output_data = []
|
|
initial_rows = 0
|
|
duplicates_removed = 0
|
|
try:
|
|
# First, try to read the file to determine the number of columns
|
|
df = pd.read_csv(file_path, nrows=5)
|
|
num_columns = len(df.columns)
|
|
|
|
# Now read the file in chunks with the correct number of columns
|
|
for chunk in pd.read_csv(file_path, chunksize=CHUNK_SIZE, on_bad_lines='warn'):
|
|
initial_rows += len(chunk)
|
|
before_dedup = len(chunk)
|
|
chunk.drop_duplicates(inplace=True)
|
|
after_dedup = len(chunk)
|
|
duplicates_removed += (before_dedup - after_dedup)
|
|
output_data.append(chunk)
|
|
except pd.errors.EmptyDataError:
|
|
print(f"Skipping empty file: {file_path}")
|
|
return None, 0, 0
|
|
except Exception as e:
|
|
print(f"Error processing file {file_path}: {str(e)}")
|
|
return None, 0, 0
|
|
|
|
if output_data:
|
|
result = pd.concat(output_data)
|
|
result.drop_duplicates(inplace=True)
|
|
result.sort_values(by=list(result.columns), inplace=True)
|
|
total_removed = initial_rows - len(result)
|
|
return result, 1, total_removed
|
|
return None, 0, 0
|
|
|
|
def save_csv_file(df, output_path):
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
df.to_csv(output_path, index=False)
|
|
|
|
def process_files_in_directory(root, files):
|
|
file_count = 0
|
|
duplicates_count = 0
|
|
for file in files:
|
|
if file.endswith(".csv"):
|
|
csv_path = os.path.join(root, file)
|
|
rel_path = os.path.relpath(csv_path, INPUT_PATH)
|
|
output_path = os.path.join(OUTPUT_PATH, rel_path)
|
|
try:
|
|
df, file_processed, duplicates_removed = process_csv_file(csv_path)
|
|
file_count += file_processed
|
|
duplicates_count += duplicates_removed
|
|
if df is not None:
|
|
save_csv_file(df, output_path)
|
|
except Exception as e:
|
|
print(f"Error processing file {csv_path}: {str(e)}")
|
|
return file_count, duplicates_count
|
|
|
|
def process_csv_files():
|
|
total_files = sum([len(files) for r, d, files in os.walk(INPUT_PATH)])
|
|
progress_bar = tqdm(total=total_files, desc="Processing CSV files", unit="file")
|
|
total_websites = 0
|
|
total_files_processed = 0
|
|
total_duplicates_removed = 0
|
|
with ProcessPoolExecutor() as executor:
|
|
futures = []
|
|
for root, dirs, files in os.walk(INPUT_PATH):
|
|
if files:
|
|
total_websites += 1
|
|
futures.append(executor.submit(process_files_in_directory, root, files))
|
|
for future in as_completed(futures):
|
|
try:
|
|
file_count, duplicates_count = future.result()
|
|
total_files_processed += file_count
|
|
total_duplicates_removed += duplicates_count
|
|
progress_bar.update(file_count)
|
|
except Exception as e:
|
|
print(f"Error in processing batch: {str(e)}")
|
|
progress_bar.close()
|
|
return total_websites, total_files_processed, total_duplicates_removed
|
|
|
|
def main():
|
|
print("Creating output directories...")
|
|
create_output_dirs()
|
|
print("Processing CSV files...")
|
|
total_websites, total_files_processed, total_duplicates_removed = process_csv_files()
|
|
print("Done processing CSV files.")
|
|
print(f"Processed CSV files are located in: {OUTPUT_PATH}")
|
|
print(f"Total websites processed: {total_websites}")
|
|
print(f"Total CSV files sorted and deduplicated: {total_files_processed}")
|
|
print(f"Total duplicates removed: {total_duplicates_removed}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|