zsv 1.3.0

data/ext/zsv/vendor/zsv-1.3.0/app/count.c

@@ -0,0 +1,404 @@
+/*
+ * Copyright (C) 2021 Liquidaty and the zsv/lib contributors
+ * All rights reserved
+ *
+ * This file is part of zsv/lib, distributed under the license defined at
+ * https://opensource.org/licenses/MIT
+ */
+
+#include <stdio.h>
+#include <string.h>
+#include <stdlib.h>
+#include <pthread.h>
+#include <sys/types.h> // off_t
+
+#define ZSV_COMMAND count
+#include "zsv_command.h"
+#include <zsv/utils/file.h>
+#include <zsv/utils/os.h> // zsv_get_number_of_cores
+#include "utils/chunk.h"
+
+#define ZSV_COUNT_PARALLEL_MIN_BYTES (1024 * 1024 * 2)
+
+struct zsv_chunk_count_data {
+  unsigned int id;
+  size_t start_offset;
+  size_t end_offset;
+
+  size_t actual_next_row_start;
+  size_t row_count;
+  int status;
+
+  const char *input_path;
+  struct zsv_opts *opts_template;
+
+  int skip;
+};
+
+struct zsv_count_parallel_data {
+  unsigned int chunk_count;
+  struct zsv_chunk_count_data *chunks;
+  pthread_t *threads;
+};
+
+struct data {
+  zsv_parser parser;
+  size_t rows;
+
+  struct zsv_opts *opts;
+  const char *input_path;
+  unsigned int num_chunks;
+
+  int run_in_parallel;
+  int cancelled;
+#ifndef ZSV_NO_PARALLEL
+  struct zsv_count_parallel_data *pdata;
+  size_t end_offset_limit; // where this chunk (chunk 0) should stop
+  size_t next_row_start;   // where chunk 0 actually ended
+#endif
+};
+
+#ifndef ZSV_NO_PARALLEL
+static void *process_chunk_internal(struct zsv_chunk_count_data *cdata);
+
+static struct zsv_count_parallel_data *parallel_data_new(unsigned int count) {
+  struct zsv_count_parallel_data *pd = calloc(1, sizeof(*pd));
+  if (!pd)
+    return NULL;
+  pd->chunk_count = count;
+  pd->chunks = calloc(count, sizeof(*pd->chunks));
+  pd->threads = calloc(count, sizeof(*pd->threads));
+  if (!pd->chunks || !pd->threads) {
+    free(pd->chunks);
+    free(pd->threads);
+    free(pd);
+    return NULL;
+  }
+  return pd;
+}
+
+static void parallel_data_delete(struct zsv_count_parallel_data *pd) {
+  if (pd) {
+    free(pd->chunks);
+    free(pd->threads);
+    free(pd);
+  }
+}
+
+#endif
+
+/* serial (non-parallelized) row handlers */
+static void row_verbose(void *ctx) {
+  struct data *data = ctx;
+  data->rows++;
+  if (data->rows % 1000000 == 0)
+    fprintf(stderr, "Processed %zu data rows\n", data->rows / 1000000);
+}
+
+static void row_simple(void *ctx) {
+  ((struct data *)ctx)->rows++;
+}
+
+#ifndef ZSV_NO_PARALLEL
+/* parallelized row handers */
+static void row_parallel_done(void *ctx) {
+  struct data *data = ctx;
+  // Find start of the next row
+  data->next_row_start = zsv_cum_scanned_length(data->parser) - zsv_row_length_raw_bytes(data->parser);
+  zsv_abort(data->parser);
+  data->cancelled = 1;
+}
+
+static void row_parallel(void *ctx) {
+  struct data *data = ctx;
+  data->rows++;
+
+  if (UNLIKELY((off_t)zsv_cum_scanned_length(data->parser) >= data->end_offset_limit)) {
+    // We crossed the boundary. We must finish this row, then stop.
+    // Switch handler to 'done' to catch the exact end of this row.
+    zsv_set_row_handler(data->parser, row_parallel_done);
+  }
+}
+
+struct worker_ctx {
+  struct zsv_chunk_count_data *cdata;
+  zsv_parser parser;
+  size_t limit_len;
+  int cancelled;
+};
+
+static void worker_row_done(void *ctx) {
+  struct worker_ctx *wctx = ctx;
+  // Calculate absolute offset of the *next* row start
+  size_t scanned = zsv_cum_scanned_length(wctx->parser);
+  wctx->cdata->actual_next_row_start = wctx->cdata->start_offset + scanned - zsv_row_length_raw_bytes(wctx->parser);
+  zsv_abort(wctx->parser);
+  wctx->cancelled = 1;
+}
+
+static void worker_row(void *ctx) {
+  struct worker_ctx *wctx = ctx;
+  wctx->cdata->row_count++;
+
+  if (UNLIKELY((off_t)zsv_cum_scanned_length(wctx->parser) >= wctx->limit_len)) {
+    zsv_set_row_handler(wctx->parser, worker_row_done);
+  }
+}
+
+static void *process_chunk_thread(void *arg) {
+  struct zsv_chunk_count_data *cdata = arg;
+  return process_chunk_internal(cdata);
+}
+
+static void *process_chunk_internal(struct zsv_chunk_count_data *cdata) {
+  cdata->row_count = 0;
+  cdata->status = 0;
+
+  if (cdata->start_offset >= cdata->end_offset) {
+    cdata->actual_next_row_start = cdata->start_offset;
+    cdata->skip = 1;
+    return NULL;
+  }
+
+  struct zsv_opts opts = *cdata->opts_template;
+  struct worker_ctx wctx = {0};
+  wctx.cdata = cdata;
+  wctx.limit_len = cdata->end_offset - cdata->start_offset;
+
+  FILE *f = fopen(cdata->input_path, "rb");
+  if (!f) {
+    cdata->status = zsv_status_error;
+    return NULL;
+  }
+
+  if (fseeko(f, cdata->start_offset, SEEK_SET) != 0) {
+    fclose(f);
+    cdata->status = zsv_status_error;
+    return NULL;
+  }
+
+  opts.stream = f;
+  opts.ctx = &wctx;
+  opts.row_handler = worker_row;
+
+  wctx.parser = zsv_new(&opts);
+  if (wctx.parser == NULL) {
+    fclose(f);
+    cdata->status = zsv_status_error;
+    return NULL;
+  }
+
+  enum zsv_status status = zsv_status_ok;
+  while (status == zsv_status_ok && !wctx.cancelled) {
+    status = zsv_parse_more(wctx.parser);
+  }
+
+  // if finished naturally (eof)
+  if (!wctx.cancelled) {
+    cdata->actual_next_row_start = cdata->start_offset + zsv_cum_scanned_length(wctx.parser);
+  }
+
+  zsv_finish(wctx.parser);
+  zsv_delete(wctx.parser);
+  fclose(f);
+  return NULL;
+}
+#endif
+
+static void header_handler(void *ctx) {
+  struct data *data = ctx;
+#ifndef ZSV_NO_PARALLEL
+  if (data->input_path && data->num_chunks > 1) {
+    size_t header_end = zsv_cum_scanned_length(data->parser);
+    struct zsv_chunk_position *offsets =
+      zsv_guess_file_chunks(data->input_path, data->num_chunks, ZSV_COUNT_PARALLEL_MIN_BYTES, header_end
+#ifndef ZSV_NO_ONLY_CRLF
+                            ,
+                            data->opts->only_crlf_rowend
+#endif
+      );
+
+    if (offsets) {
+      data->pdata = parallel_data_new(data->num_chunks);
+      if (!data->pdata) {
+        fprintf(stderr, "Out of memory!\n");
+        zsv_free_chunks(offsets);
+      } else {
+        data->run_in_parallel = 1;
+        if (data->opts->verbose) {
+          for (unsigned int i = 0; i < data->num_chunks; i++) {
+            fprintf(stderr, "Chunk %i: %zu - %zu\n", i + 1, offsets[i].start, offsets[i].end);
+          }
+        }
+
+        /* set up worker chunks (1..n) */
+        for (unsigned int i = 1; i < data->num_chunks; i++) {
+          struct zsv_chunk_count_data *c = &data->pdata->chunks[i];
+          c->id = i;
+          c->start_offset = offsets[i].start;
+          c->end_offset = offsets[i].end;
+          c->input_path = data->input_path;
+          c->opts_template = data->opts;
+
+          if (pthread_create(&data->pdata->threads[i], NULL, process_chunk_thread, c) != 0) {
+            fprintf(stderr, "Error creating thread %d\n", i);
+            data->run_in_parallel = 0;
+            break;
+          }
+        }
+
+        if (data->run_in_parallel) {
+          data->end_offset_limit = offsets[0].end;
+          zsv_set_row_handler(data->parser, row_parallel);
+          data->run_in_parallel = 1;
+        }
+      }
+      zsv_free_chunks(offsets);
+    }
+  }
+#endif
+
+  if (!data->run_in_parallel) { // single-threaded serial run
+    data->run_in_parallel = 0;
+    zsv_set_row_handler(data->parser, data->opts->verbose ? row_verbose : row_simple);
+  }
+}
+
+static int count_usage(void) {
+  static const char *usage = "Usage: count [options]\n"
+                             "\n"
+                             "Options:\n"
+                             "  -h,--help             : show usage\n"
+                             "  -i,--input <filename> : use specified file input\n"
+#ifndef ZSV_NO_PARALLEL
+                             "  -j,--jobs <n>         : number of jobs (parallel threads)\n"
+                             "  --parallel            : use all available cores\n"
+#endif
+    ;
+  printf("%s\n", usage);
+  return 0;
+}
+
+int ZSV_MAIN_FUNC(ZSV_COMMAND)(int argc, const char *argv[], struct zsv_opts *optsp,
+                               struct zsv_prop_handler *custom_prop_handler) {
+  struct data data = {0};
+  struct zsv_opts opts = *optsp;
+  data.opts = &opts;
+
+  int err = 0;
+  for (int i = 1; !err && i < argc; i++) {
+    const char *arg = argv[i];
+    if (!strcmp(arg, "-h") || !strcmp(arg, "--help")) {
+      count_usage();
+      goto count_done;
+    }
+    if (!strcmp(arg, "-i") || !strcmp(arg, "--input") || *arg != '-') {
+      err = 1;
+      if ((!strcmp(arg, "-i") || !strcmp(arg, "--input")) && ++i >= argc)
+        fprintf(stderr, "%s option requires a filename\n", arg);
+      else {
+        if (opts.stream)
+          fprintf(stderr, "Input may not be specified more than once\n");
+        else if (!(opts.stream = fopen(argv[i], "rb")))
+          fprintf(stderr, "Unable to open for reading: %s\n", argv[i]);
+        else {
+          data.input_path = argv[i];
+          err = 0;
+        }
+      }
+#ifndef ZSV_NO_PARALLEL
+    } else if (!strcmp(arg, "-j") || !strcmp(arg, "--jobs")) {
+      if (++i >= argc)
+        err = 1;
+      else
+        data.num_chunks = atoi(argv[i]);
+    } else if (!strcmp(arg, "--parallel")) {
+      data.num_chunks = zsv_get_number_of_cores();
+      if (data.num_chunks < 2) {
+        fprintf(stderr, "Warning: --parallel specified but only one core found; using -j 4 instead");
+        data.num_chunks = 4;
+      }
+#endif
+    } else {
+      fprintf(stderr, "Unrecognized option: %s\n", arg);
+      err = 1;
+    }
+  }
+
+#ifdef NO_STDIN
+  if (!opts.stream || opts.stream == stdin) {
+    fprintf(stderr, "Please specify an input file\n");
+    err = 1;
+  }
+#endif
+#ifndef ZSV_NO_PARALLEL
+  if (data.num_chunks > 1) {
+    enum zsv_chunk_status chstat = zsv_chunkable(data.input_path, &opts);
+    if (chstat != zsv_chunk_status_ok) {
+      fprintf(stderr, "%s\n", zsv_chunk_status_str(chstat));
+      err = 1;
+    }
+  }
+#endif
+  if (!err) {
+    opts.row_handler = header_handler;
+    opts.ctx = &data;
+
+    if (zsv_new_with_properties(&opts, custom_prop_handler, data.input_path, &data.parser) != zsv_status_ok) {
+      fprintf(stderr, "Unable to initialize parser\n");
+      err = 1;
+    } else {
+      enum zsv_status status;
+
+      /* Main Parse Loop */
+      while (!data.cancelled && (status = zsv_parse_more(data.parser)) == zsv_status_ok)
+        ;
+      zsv_finish(data.parser);
+
+#ifndef ZSV_NO_PARALLEL
+      if (data.run_in_parallel) {
+        if (!data.next_row_start)
+          // not likely to get here but just in case
+          data.next_row_start = zsv_cum_scanned_length(data.parser);
+
+        size_t total_rows = data.rows;
+        // aggregate results
+        for (unsigned int i = 1; i < data.num_chunks; i++) {
+          pthread_join(data.pdata->threads[i], NULL);
+
+          struct zsv_chunk_count_data *prev_chunk = (i == 1) ? NULL : &data.pdata->chunks[i - 1];
+          struct zsv_chunk_count_data *curr_chunk = &data.pdata->chunks[i];
+
+          // determine where the previous chunk actually ended
+          size_t prev_end = (i == 1) ? data.next_row_start : prev_chunk->actual_next_row_start;
+          // check overlap
+          if (prev_end > curr_chunk->start_offset) {
+            if (data.opts->verbose) {
+              fprintf(stderr, "Overlap detected at chunk %u (expected %zu, got %zu). Reprocessing.\n", i,
+                      curr_chunk->start_offset, prev_end);
+            }
+            // reprocess synchronously
+            curr_chunk->start_offset = prev_end;
+            process_chunk_internal(curr_chunk);
+          }
+
+          total_rows += curr_chunk->row_count;
+        }
+
+        printf("%zu\n", total_rows);
+        parallel_data_delete(data.pdata);
+
+      } else
+#endif
+        // result from running serially
+        printf("%zu\n", data.rows);
+      zsv_delete(data.parser);
+    }
+  }
+
+count_done:
+  if (opts.stream && opts.stream != stdin)
+    fclose(opts.stream);
+
+  return err;
+}