<commit_before>use 5.10.0;
use strict;
use warnings;

# ABSTRACT: Create an annotation from a query
# TODO: 1) Support sort
####### 2) Don't initialize db
####### 3) figure out why even with field => 'pos' we get very slow perf

# 1) igure out if we want to change that.
#    - It is ignored because the elasticsearch cluster is too slow on a per-shard basis
# 2) Document it

package SeqFromQuery;
our $VERSION = '0.001';

use namespace::autoclean;
use DDP;
use lib './lib';
use Mouse 2;

use MCE::Loop;

use Search::Elasticsearch;

use Seq::Output::Delimiters;

use Cpanel::JSON::XS qw/decode_json encode_json/;
use YAML::XS qw/LoadFile/;
use Statistics::Distributions qw(chisqrdistr udistr);

use Math::Round qw/nhimult round/;

# Defines basic things needed in builder and annotator, like logPath,
# Also initializes the database with database_dir unnecessarily
extends 'Seq::Base';

# Defines most of the properties that can be configured at run time
# Needed because there are variations of Seq.pm, ilke SeqFromQuery.pm
with 'Seq::Definition';

# An archive, containing an "annotation" file
has inputQueryBody => (is => 'ro', isa => 'HashRef', required => 1);

# Post-processing to run, before commiting the annotation
has pipeline => (is => 'ro', isa => 'ArrayRef[HashRef]');

# Probably the user id
has indexName => (is => 'ro', required => 1);

# The index type; probably the job id
has indexType => (is => 'ro', required => 1);

has assembly => (is => 'ro', isa => 'Str', required => 1);

has configPath => (is => 'ro', isa => 'Str', default => 'config/');

# has commitEvery => (is => 'ro', default => 5000);

# The user may have given some header fields already
# If so, this is a re-indexing job, and we will want to append the header fields
has fieldNames => (is => 'ro', isa => 'ArrayRef', required => 1);

has indexConfig => (is => 'ro', isa => 'HashRef', required => 1);

# Elasticsearch connection configuration
has connection => (is => 'ro', isa => 'HashRef', required => 1);

has batchSize => (is => 'ro', isa => 'Num', default => 5000);

has shards => (is => 'ro', isa => 'Num', lazy => 1, default => sub {
  my $self = shift;

  return $self->indexConfig->{index_settings}->{index}->{number_of_shards};
});

has maxShards => (is => 'ro', isa => 'Num');

my $prettyCoder = Cpanel::JSON::XS->new->ascii->pretty->allow_nonref;;

around BUILDARGS => sub {
  my ($orig, $class, $href) = @_;

  my %data = %$href;

  if($data{indexConfig}) {
    if(!ref $data{indexConfig}) {
      $data{indexConfig} = decode_json($data{indexConfig});
    }

    return $class->$orig(\%data);
  }

  my $cf = path('config')->child($data{assembly} . '.mapping.yml')->stringify();

  $data{indexConfig} = LoadFile($cf);

  return $class->$orig(\%data);
};

sub BUILD {
  my $self = shift;

  # Makes or fails silently if exists
  $self->outDir->mkpath();

  if(!$self->shards) {
    $self->log('fatal', "Cannot read number of shards in index");
  }
}

sub annotate {
  my $self = shift;

  $self->log( 'info', 'Beginning saving annotation from query' );

  $self->log( 'info', 'Input query is: ' . $prettyCoder->encode($self->inputQueryBody) );

  my $d = Seq::Output::Delimiters->new();

  my %delims = (
    'pos' => $d->positionDelimiter,
    'value' => $d->valueDelimiter,
    'overlap' => $d->overlapDelimiter,
    'miss' => $d->emptyFieldChar,
    'fieldSep' => $d->fieldSeparator,
  );

  ################## Make the full output path ######################
  # The output path always respects the $self->output_file_base attribute path;
  my ($err, $outFh, $statsFh) = $self->_getFileHandles();

  if ($err) {
    $self->_errorWithCleanup($err);
    return ($err, undef);
  }

  # TODO: Support sort
  $self->_cleanQuery();

  my ($parentsAref, $childrenAref) = $self->_getHeader();

  my @childrenOrOnly = @$childrenAref;
  my @parentNames = @$parentsAref;

  my @fieldNames = @{$self->fieldNames};
  my $outputHeader = join($delims{fieldSep}, @fieldNames);

  say $outFh $outputHeader;

  if($statsFh) {
    say $statsFh $outputHeader;
  }

  ($err, my $discordantIdx) = $self->_getDiscordantIdx();

  if($err) {
    $self->log('fatal', "Couldn't find discordant index");
  }

  # TODO: figure out why even with field => 'pos' we get very slow perf
  # when having slices > shards with very large (1000+ terms) queries

  my ($nSlices, $batchSize, $timeout) = $self->_getSearchParams();

  my $slice = {
    field => 'pos',
    max => $nSlices,

  };

  my $messageFreq = (2e4 / 4) * $self->max_threads;

  my $hasSort = exists $self->inputQueryBody->{sort};
  my $progressFunc = $self->_makeLogProgress($hasSort, $outFh, $statsFh, $messageFreq);

  my @filterFunctions;
  if($self->pipeline) {
    my ($err, $filters) = $self->_makePipeline();

    if($err) {
      return ($err, undef);
    }

    if(defined $filters) {
      @filterFunctions = @$filters;
    }
  }

  MCE::Loop::init {
    max_workers => $nSlices,
    chunk_size => 1,
    gather => $progressFunc,
  };

  $self->log('info', "Beginning to create annotation from the query");

  # We do parallel fetch using sliced scroll queries.
  # It is far too expensive to do a single-threaded scroll
  # and pass the entire structure to MCE;
  # Will get "recursion limit" errors in Storable due to the size of the structure
  # for bulk queries
  mce_loop {
    my ($mce, $chunkRef, $chunkId) = @_;

    my $id = $_;

    my $es = Search::Elasticsearch->new($self->connection);

    $slice->{id} = $id;
    $self->inputQueryBody->{slice} = $slice;

    my $scroll = $es->scroll_helper(
      scroll      => $timeout,
      size        => $batchSize,
      body        => $self->inputQueryBody,
      index       => $self->indexName,
      type        => $self->indexType,
    );

    while(my @docs = $scroll->next($batchSize)) {
      my @sourceData;
      # $#sourceData = $#docs;

      my $skipped = 0;
      DOCS: for my $doc (@docs) {
        if(@filterFunctions) {
          for my $f (@filterFunctions) {
            if($f->($doc->{_source})) {
              $skipped++;

              next DOCS;
            }
          }
        }

        my @rowData;
        # Initialize all values to undef
        # Output.pm requires a sparse array for any missing values
        # To preserve output order
        $#rowData = $#fieldNames;

        for my $y (0 .. $#fieldNames) {
          $rowData[$y] = _populateArrayPathFromHash($childrenOrOnly[$y], $doc->{_source}{$parentNames[$y]});
        }

        if($rowData[$discordantIdx][0][0] eq 'false') {
          $rowData[$discordantIdx][0][0] = 0;
        } elsif($rowData[$discordantIdx][0][0] eq 'true') {
          $rowData[$discordantIdx][0][0] = 1;
        }

        push @sourceData, \@rowData;
      }

      my $outputString = _makeOutputString(\@sourceData, \%delims);

      $mce->gather(scalar @docs - $skipped, $outputString, $id);
    }
  } (0 .. $nSlices - 1);

  # Flush
  $progressFunc->(0, undef, undef, 1);

  MCE::Loop::finish();

  ################ Finished writing file. If statistics, print those ##########
  # First sync output to ensure everything needed is written
  # then close all file handles and move files to output dir
  $err = $self->safeSystem('sync')
         ||
         $self->safeClose($outFh)
         ||
         ($statsFh && $self->safeClose($statsFh))
         ||
         $self->_moveFilesToOutputDir();

  if($err) {
    $self->_errorWithCleanup($err);

    return ($err, undef);
  }

  return ($err, $self->outputFilesInfo);
}

sub _makePipeline {
  my $self = shift;

  state $funcs = {
    binomMaf => \&makeBinomFilter,
    hwe => \&makeHweFilter,
  };

  if(!$self->pipeline) {
    return (undef, undef);
  }

  my @funcs;

  for my $step (@{$self->pipeline}) {
    if(!exists $funcs->{$step->{key}}) {
      $self->log('warn', "Couldn't find function for $step->{key}");
      next;
    }

    my $makeFunc = $funcs->{$step->{key}};

    # Skipped for some reason
    if(!$makeFunc) {
      next;
    }

    push @funcs, $makeFunc->($self, $step);
  }

  return (undef, \@funcs);
}

sub _getSearchParams {
  my $self = shift;

  my $numTerms = $self->_getNumTerms($self->inputQueryBody);

  # -1 simply means we can't approximate the size
  my $nSlices;
  if($numTerms < 0 || $numTerms > 10000) {
    $nSlices = $self->shards;
  } else {
    $nSlices = $self->_getSlices();
  }

  if($nSlices < 1) {
    $nSlices = 1;
  }

  my $batchSize = $self->batchSize;<commit_msg>max size must be 2 or more<commit_after>use 5.10.0;
use strict;
use warnings;

# ABSTRACT: Create an annotation from a query
# TODO: 1) Support sort
####### 2) Don't initialize db
####### 3) figure out why even with field => 'pos' we get very slow perf

# 1) igure out if we want to change that.
#    - It is ignored because the elasticsearch cluster is too slow on a per-shard basis
# 2) Document it

package SeqFromQuery;
our $VERSION = '0.001';

use namespace::autoclean;
use DDP;
use lib './lib';
use Mouse 2;

use MCE::Loop;

use Search::Elasticsearch;

use Seq::Output::Delimiters;

use Cpanel::JSON::XS qw/decode_json encode_json/;
use YAML::XS qw/LoadFile/;
use Statistics::Distributions qw(chisqrdistr udistr);

use Math::Round qw/nhimult round/;

# Defines basic things needed in builder and annotator, like logPath,
# Also initializes the database with database_dir unnecessarily
extends 'Seq::Base';

# Defines most of the properties that can be configured at run time
# Needed because there are variations of Seq.pm, ilke SeqFromQuery.pm
with 'Seq::Definition';

# An archive, containing an "annotation" file
has inputQueryBody => (is => 'ro', isa => 'HashRef', required => 1);

# Post-processing to run, before commiting the annotation
has pipeline => (is => 'ro', isa => 'ArrayRef[HashRef]');

# Probably the user id
has indexName => (is => 'ro', required => 1);

# The index type; probably the job id
has indexType => (is => 'ro', required => 1);

has assembly => (is => 'ro', isa => 'Str', required => 1);

has configPath => (is => 'ro', isa => 'Str', default => 'config/');

# has commitEvery => (is => 'ro', default => 5000);

# The user may have given some header fields already
# If so, this is a re-indexing job, and we will want to append the header fields
has fieldNames => (is => 'ro', isa => 'ArrayRef', required => 1);

has indexConfig => (is => 'ro', isa => 'HashRef', required => 1);

# Elasticsearch connection configuration
has connection => (is => 'ro', isa => 'HashRef', required => 1);

has batchSize => (is => 'ro', isa => 'Num', default => 5000);

has shards => (is => 'ro', isa => 'Num', lazy => 1, default => sub {
  my $self = shift;

  return $self->indexConfig->{index_settings}->{index}->{number_of_shards};
});

has maxShards => (is => 'ro', isa => 'Num');

my $prettyCoder = Cpanel::JSON::XS->new->ascii->pretty->allow_nonref;;

around BUILDARGS => sub {
  my ($orig, $class, $href) = @_;

  my %data = %$href;

  if($data{indexConfig}) {
    if(!ref $data{indexConfig}) {
      $data{indexConfig} = decode_json($data{indexConfig});
    }

    return $class->$orig(\%data);
  }

  my $cf = path('config')->child($data{assembly} . '.mapping.yml')->stringify();

  $data{indexConfig} = LoadFile($cf);

  return $class->$orig(\%data);
};

sub BUILD {
  my $self = shift;

  # Makes or fails silently if exists
  $self->outDir->mkpath();

  if(!$self->shards) {
    $self->log('fatal', "Cannot read number of shards in index");
  }
}

sub annotate {
  my $self = shift;

  $self->log( 'info', 'Beginning saving annotation from query' );

  $self->log( 'info', 'Input query is: ' . $prettyCoder->encode($self->inputQueryBody) );

  my $d = Seq::Output::Delimiters->new();

  my %delims = (
    'pos' => $d->positionDelimiter,
    'value' => $d->valueDelimiter,
    'overlap' => $d->overlapDelimiter,
    'miss' => $d->emptyFieldChar,
    'fieldSep' => $d->fieldSeparator,
  );

  ################## Make the full output path ######################
  # The output path always respects the $self->output_file_base attribute path;
  my ($err, $outFh, $statsFh) = $self->_getFileHandles();

  if ($err) {
    $self->_errorWithCleanup($err);
    return ($err, undef);
  }

  # TODO: Support sort
  $self->_cleanQuery();

  my ($parentsAref, $childrenAref) = $self->_getHeader();

  my @childrenOrOnly = @$childrenAref;
  my @parentNames = @$parentsAref;

  my @fieldNames = @{$self->fieldNames};
  my $outputHeader = join($delims{fieldSep}, @fieldNames);

  say $outFh $outputHeader;

  if($statsFh) {
    say $statsFh $outputHeader;
  }

  ($err, my $discordantIdx) = $self->_getDiscordantIdx();

  if($err) {
    $self->log('fatal', "Couldn't find discordant index");
  }

  # TODO: figure out why even with field => 'pos' we get very slow perf
  # when having slices > shards with very large (1000+ terms) queries

  my ($nSlices, $batchSize, $timeout) = $self->_getSearchParams();

  my $slice = {
    field => 'pos',
    max => $nSlices,

  };

  my $messageFreq = (2e4 / 4) * $self->max_threads;

  my $hasSort = exists $self->inputQueryBody->{sort};
  my $progressFunc = $self->_makeLogProgress($hasSort, $outFh, $statsFh, $messageFreq);

  my @filterFunctions;
  if($self->pipeline) {
    my ($err, $filters) = $self->_makePipeline();

    if($err) {
      return ($err, undef);
    }

    if(defined $filters) {
      @filterFunctions = @$filters;
    }
  }

  MCE::Loop::init {
    max_workers => $nSlices,
    chunk_size => 1,
    gather => $progressFunc,
  };

  $self->log('info', "Beginning to create annotation from the query");

  # We do parallel fetch using sliced scroll queries.
  # It is far too expensive to do a single-threaded scroll
  # and pass the entire structure to MCE;
  # Will get "recursion limit" errors in Storable due to the size of the structure
  # for bulk queries
  mce_loop {
    my ($mce, $chunkRef, $chunkId) = @_;

    my $id = $_;

    my $es = Search::Elasticsearch->new($self->connection);

    $slice->{id} = $id;
    $self->inputQueryBody->{slice} = $slice;

    my $scroll = $es->scroll_helper(
      scroll      => $timeout,
      size        => $batchSize,
      body        => $self->inputQueryBody,
      index       => $self->indexName,
      type        => $self->indexType,
    );

    while(my @docs = $scroll->next($batchSize)) {
      my @sourceData;
      # $#sourceData = $#docs;

      my $skipped = 0;
      DOCS: for my $doc (@docs) {
        if(@filterFunctions) {
          for my $f (@filterFunctions) {
            if($f->($doc->{_source})) {
              $skipped++;

              next DOCS;
            }
          }
        }

        my @rowData;
        # Initialize all values to undef
        # Output.pm requires a sparse array for any missing values
        # To preserve output order
        $#rowData = $#fieldNames;

        for my $y (0 .. $#fieldNames) {
          $rowData[$y] = _populateArrayPathFromHash($childrenOrOnly[$y], $doc->{_source}{$parentNames[$y]});
        }

        if($rowData[$discordantIdx][0][0] eq 'false') {
          $rowData[$discordantIdx][0][0] = 0;
        } elsif($rowData[$discordantIdx][0][0] eq 'true') {
          $rowData[$discordantIdx][0][0] = 1;
        }

        push @sourceData, \@rowData;
      }

      my $outputString = _makeOutputString(\@sourceData, \%delims);

      $mce->gather(scalar @docs - $skipped, $outputString, $id);
    }
  } (0 .. $nSlices - 1);

  # Flush
  $progressFunc->(0, undef, undef, 1);

  MCE::Loop::finish();

  ################ Finished writing file. If statistics, print those ##########
  # First sync output to ensure everything needed is written
  # then close all file handles and move files to output dir
  $err = $self->safeSystem('sync')
         ||
         $self->safeClose($outFh)
         ||
         ($statsFh && $self->safeClose($statsFh))
         ||
         $self->_moveFilesToOutputDir();

  if($err) {
    $self->_errorWithCleanup($err);

    return ($err, undef);
  }

  return ($err, $self->outputFilesInfo);
}

sub _makePipeline {
  my $self = shift;

  state $funcs = {
    binomMaf => \&makeBinomFilter,
    hwe => \&makeHweFilter,
  };

  if(!$self->pipeline) {
    return (undef, undef);
  }

  my @funcs;

  for my $step (@{$self->pipeline}) {
    if(!exists $funcs->{$step->{key}}) {
      $self->log('warn', "Couldn't find function for $step->{key}");
      next;
    }

    my $makeFunc = $funcs->{$step->{key}};

    # Skipped for some reason
    if(!$makeFunc) {
      next;
    }

    push @funcs, $makeFunc->($self, $step);
  }

  return (undef, \@funcs);
}

sub _getSearchParams {
  my $self = shift;

  my $numTerms = $self->_getNumTerms($self->inputQueryBody);

  # -1 simply means we can't approximate the size
  my $nSlices;
  if($numTerms < 0 || $numTerms > 10000) {
    $nSlices = $self->shards;
  } else {
    $nSlices = $self->_getSlices();
  }

  if($nSlices < 2) {
    $nSlices = 2;
  }

  my $batchSize = $self->batchSize;
{'fraction_non_alphanumeric': 0.17369697686801158, 'fraction_numerical': 0.007513459149263445, 'mean_word_length': 3.160472557223726, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 297, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6675795', 'n_tokens_mistral': 5950, 'n_tokens_neox': 5673, 'n_words': 1587}