shuffle.py

"""
Shuffle lines in a [big] file
shuffle.py input_file.csv output_file.csv [<max. lines in memory>] [<random seed>]

"""

import sys
import random

input_file = sys.argv[1]
output_file = sys.argv[2]

try:
	lines_in_memory = int( sys.argv[3] )
except IndexError:
	lines_in_memory = 25000
	
print "caching %s lines at a time..." % ( lines_in_memory )
	
try:
	random_seed = sys.argv[4]
	random.seed( random_seed )
	print "random seed: %s" % ( random_seed )
except IndexError:
	pass
	
# first count

print "counting lines..."

f = open( input_file )

count =  0
for line in f:
	count += 1
	
	if count % 100000 == 0:
		print count
	
print count
		
# then shuffle		

print "shuffling..."

o_f = open( output_file, 'wb' )
	
order = range( count )
random.shuffle( order )

epoch = 0
	
while order:

	current_lines = {}
	current_lines_count = 0

	current_chunk = order[:lines_in_memory]
	current_chunk_dict = { x: 1 for x in current_chunk }		# faster "in"
	current_chunk_length = len( current_chunk )
	
	order = order[lines_in_memory:]
	
	f.seek( 0 )
	count = 0
		
	for line in f:
		if count in current_chunk_dict:
			current_lines[count] = line
			current_lines_count += 1
			if current_lines_count == current_chunk_length:
				break
		count += 1	
		if count % 100000 == 0:
			print count		
	
	print "writing..."
	
	for l in current_chunk:
		o_f.write( current_lines[l] )
	
	lines_saved = current_chunk_length + epoch * lines_in_memory
	epoch += 1
	print "pass %s complete (%s lines saved)" % ( epoch, lines_saved )