[parser] adding chunked shuffle as a C function (writes each line to one of n random files, runs shuf on each file and concatenates the result). Adding a version which allows specifying a specific chunk size, and using a 2GB limit for address parser training. Allowing gshuf again for Mac as it seems the only problem there was not having enough memory when testing on a Mac laptop. The new limited-memory version should be fast enough.
This commit is contained in:
@@ -3,11 +3,14 @@
|
||||
#include <config.h>
|
||||
#include "string_utils.h"
|
||||
|
||||
// Run shuf/gshuf on a file in-place if the shuf command is available.
|
||||
bool shuffle_file(char *filename) {
|
||||
char *shuffle_command = NULL;
|
||||
|
||||
#if defined(HAVE_SHUF)
|
||||
shuffle_command = "shuf";
|
||||
#elif defined(HAVE_GSHUF)
|
||||
shuffle_command = "gshuf";
|
||||
#else
|
||||
return false;
|
||||
#endif
|
||||
@@ -22,3 +25,97 @@ bool shuffle_file(char *filename) {
|
||||
|
||||
return ret == EXIT_SUCCESS;
|
||||
}
|
||||
|
||||
// Assign each line of the file randomly to n chunks and shuffle each file sequentially in-memory
|
||||
// This approach will produce a random permutation of the lines using limited memory
|
||||
bool shuffle_file_chunked(char *filename, size_t parts) {
|
||||
char *shuffle_command = NULL;
|
||||
|
||||
// Linux
|
||||
#if defined(HAVE_SHUF)
|
||||
shuffle_command = "shuf";
|
||||
// Mac
|
||||
#elif defined(HAVE_GSHUF)
|
||||
shuffle_command = "gshuf";
|
||||
#else
|
||||
return false;
|
||||
#endif
|
||||
|
||||
if (filename == NULL) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Make sure the input file exists
|
||||
FILE *f = fopen(filename, "r");
|
||||
if (f == NULL) {
|
||||
return false;
|
||||
}
|
||||
fclose(f);
|
||||
|
||||
// This is an in-place shuffle to keep the API simple
|
||||
char *outfile = filename;
|
||||
|
||||
char_array *command = char_array_new();
|
||||
|
||||
// Split the file randomly into $parts files
|
||||
// Need to be assigned randomly, not just every nth line or it's not really a random permutation
|
||||
char_array_cat_printf(command, "awk -v parts=%zu -v filename=%s 'BEGIN{srand();} { print > filename\".\"int(rand() * parts) }'", parts, filename);
|
||||
|
||||
int ret = system(char_array_get_string(command));
|
||||
if (ret != EXIT_SUCCESS) {
|
||||
goto exit_char_array_allocated;
|
||||
}
|
||||
|
||||
// Run shuf sequentially on each of the $parts files
|
||||
// This should be sequential, not parallelized as the goal is
|
||||
// to limit memory usage when shuffling large files
|
||||
for (size_t i = 0; i < parts; i++) {
|
||||
char_array_clear(command);
|
||||
char_array_cat_printf(command, "%s %s.%zu %s %s.tmp", shuffle_command, filename, i, i > 0 ? ">>" : ">", outfile);
|
||||
ret = system(char_array_get_string(command));
|
||||
if (ret != EXIT_SUCCESS) {
|
||||
goto exit_char_array_allocated;
|
||||
}
|
||||
|
||||
// Delete the file temp file
|
||||
char_array_clear(command);
|
||||
char_array_cat_printf(command, "rm %s.%zu", filename, i);
|
||||
ret = system(char_array_get_string(command));
|
||||
if (ret != EXIT_SUCCESS) {
|
||||
goto exit_char_array_allocated;
|
||||
}
|
||||
}
|
||||
|
||||
char_array_clear(command);
|
||||
char_array_cat_printf(command, "mv %s.tmp %s", outfile, outfile);
|
||||
ret = system(char_array_get_string(command));
|
||||
if (ret != EXIT_SUCCESS) {
|
||||
goto exit_char_array_allocated;
|
||||
}
|
||||
|
||||
exit_char_array_allocated:
|
||||
char_array_destroy(command);
|
||||
return ret == EXIT_SUCCESS;
|
||||
}
|
||||
|
||||
// Shuffle a file in-place, specifying a rough upper bound on system memory
|
||||
bool shuffle_file_chunked_size(char *filename, size_t chunk_size) {
|
||||
FILE *f = fopen(filename, "r");
|
||||
|
||||
if (f == NULL) return false;
|
||||
|
||||
fseek(f, 0L, SEEK_END);
|
||||
size_t size = ftell(f);
|
||||
fclose(f);
|
||||
|
||||
size_t parts = size / chunk_size + 1;
|
||||
// If the file is smaller than the chunk size, do a
|
||||
// simple in-memory shuffle of the whole file
|
||||
if (parts == 1) {
|
||||
return shuffle_file(filename);
|
||||
}
|
||||
|
||||
return shuffle_file_chunked(filename, parts);
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user