<commit_before>/*
* Druid - a distributed column store.
* Copyright (C) 2012, 2013  Metamarkets Group Inc.
*
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU General Public License
* as published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
*/

package io.druid.indexing.common.task;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Ordering;
import com.google.common.collect.Sets;
import com.google.common.collect.TreeMultiset;
import com.google.common.primitives.Ints;
import com.metamx.common.logger.Logger;
import io.druid.data.input.Firehose;
import io.druid.data.input.FirehoseFactory;
import io.druid.data.input.InputRow;
import io.druid.data.input.impl.SpatialDimensionSchema;
import io.druid.granularity.QueryGranularity;
import io.druid.indexer.granularity.GranularitySpec;
import io.druid.indexing.common.TaskLock;
import io.druid.indexing.common.TaskStatus;
import io.druid.indexing.common.TaskToolbox;
import io.druid.indexing.common.index.YeOldePlumberSchool;
import io.druid.query.aggregation.AggregatorFactory;
import io.druid.segment.loading.DataSegmentPusher;
import io.druid.segment.realtime.FireDepartmentMetrics;
import io.druid.segment.realtime.Schema;
import io.druid.segment.realtime.plumber.Plumber;
import io.druid.segment.realtime.plumber.Sink;
import io.druid.timeline.DataSegment;
import io.druid.timeline.partition.NoneShardSpec;
import io.druid.timeline.partition.ShardSpec;
import io.druid.timeline.partition.SingleDimensionShardSpec;
import org.joda.time.DateTime;
import org.joda.time.Interval;

import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * Simple single-threaded indexing task. Meant to be easy to use for small and medium sized datasets. For "big data",
 * try launching multiple IndexTasks or using the {@link HadoopIndexTask}.
 */
public class IndexTask extends AbstractFixedIntervalTask
{
  private static final Logger log = new Logger(IndexTask.class);

  @JsonIgnore
  private final GranularitySpec granularitySpec;

  @JsonProperty
  private final List<SpatialDimensionSchema> spatialDimensions;

  @JsonIgnore
  private final AggregatorFactory[] aggregators;

  @JsonIgnore
  private final QueryGranularity indexGranularity;

  @JsonIgnore
  private final int targetPartitionSize;

  @JsonIgnore
  private final FirehoseFactory firehoseFactory;

  @JsonIgnore
  private final int rowFlushBoundary;

  @JsonCreator
  public IndexTask(
      @JsonProperty("id") String id,
      @JsonProperty("dataSource") String dataSource,
      @JsonProperty("granularitySpec") GranularitySpec granularitySpec,
      @JsonProperty("spatialDimensions") List<SpatialDimensionSchema> spatialDimensions,
      @JsonProperty("aggregators") AggregatorFactory[] aggregators,
      @JsonProperty("indexGranularity") QueryGranularity indexGranularity,
      @JsonProperty("targetPartitionSize") int targetPartitionSize,
      @JsonProperty("firehose") FirehoseFactory firehoseFactory,
      @JsonProperty("rowFlushBoundary") int rowFlushBoundary
  )
  {
    super(
        // _not_ the version, just something uniqueish
        id != null ? id : String.format("index_%s_%s", dataSource, new DateTime().toString()),
        dataSource,
        new Interval(
            granularitySpec.bucketIntervals().first().getStart(),
            granularitySpec.bucketIntervals().last().getEnd()
        )
    );

    this.granularitySpec = Preconditions.checkNotNull(granularitySpec, "granularitySpec");
    this.spatialDimensions = (spatialDimensions == null)
                             ? Lists.<SpatialDimensionSchema>newArrayList()
                             : spatialDimensions;
    this.aggregators = aggregators;
    this.indexGranularity = (indexGranularity == null) ? QueryGranularity.NONE : indexGranularity;
    this.targetPartitionSize = targetPartitionSize;
    this.firehoseFactory = Preconditions.checkNotNull(firehoseFactory, "firehoseFactory");
    this.rowFlushBoundary = rowFlushBoundary;
  }

  @Override
  public String getType()
  {
    return "index";
  }

  @Override
  public TaskStatus run(TaskToolbox toolbox) throws Exception
  {
    final TaskLock myLock = Iterables.getOnlyElement(getTaskLocks(toolbox));

    final Map<Interval, List<ShardSpec>> shardSpecMap = determinePartitions(targetPartitionSize);

    final Map<Interval, List<Schema>> schemass = Maps.transformEntries(
        shardSpecMap,
        new Maps.EntryTransformer<Interval, List<ShardSpec>, List<Schema>>()
        {
          @Override
          public List<Schema> transformEntry(
              Interval key, List<ShardSpec> shardSpecs
          )
          {
            return Lists.transform(
                shardSpecs,
                new Function<ShardSpec, Schema>()
                {
                  @Override
                  public Schema apply(final ShardSpec shardSpec)
                  {
                    return new Schema(getDataSource(), spatialDimensions, aggregators, indexGranularity, shardSpec);
                  }
                }
            );
          }
        }
    );

    final Set<DataSegment> segments = generateSegments(toolbox, schemass, myLock.getVersion());
    toolbox.pushSegments(segments);

    return TaskStatus.success(getId());
  }

  private Map<Interval, List<ShardSpec>> determinePartitions(
      final int targetPartitionSize
  ) throws IOException
  {
    Map<Interval, List<ShardSpec>> retVal = Maps.newLinkedHashMap();

    log.info("Determining partitions with targetPartitionSize[%d]", targetPartitionSize);

    // Blacklist dimensions that have multiple values per row
    final Set<String> unusableDimensions = Sets.newHashSet();

    // Load data
    try (Firehose firehose = firehoseFactory.connect()) {
      if (!firehose.hasMore()) {
        log.error("Unable to find any events to ingest! Check your firehose config!");
        return retVal;
      }
      InputRow inputRow = firehose.nextRow();

      for (Interval interval : granularitySpec.bucketIntervals()) {
        // Track values of all non-blacklisted dimensions
        final Map<String, TreeMultiset<String>> dimensionValueMultisets = Maps.newHashMap();

        boolean hasEventsInInterval = false;
        boolean done = false;
        while (!done && interval.contains(inputRow.getTimestampFromEpoch())) {
          hasEventsInInterval = true;

          // Extract dimensions from event
          for (final String dim : inputRow.getDimensions()) {
            final List<String> dimValues = inputRow.getDimension(dim);
            if (!unusableDimensions.contains(dim)) {
              if (dimValues.size() == 1) {
                // Track this value
                TreeMultiset<String> dimensionValueMultiset = dimensionValueMultisets.get(dim);
                if (dimensionValueMultiset == null) {
                  dimensionValueMultiset = TreeMultiset.create();
                  dimensionValueMultisets.put(dim, dimensionValueMultiset);
                }
                dimensionValueMultiset.add(dimValues.get(0));
              } else {
                // Only single-valued dimensions can be used for partitions
                unusableDimensions.add(dim);
                dimensionValueMultisets.remove(dim);
              }<commit_msg>fix small bug with unusable dims<commit_after>/*
* Druid - a distributed column store.
* Copyright (C) 2012, 2013  Metamarkets Group Inc.
*
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU General Public License
* as published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
*/

package io.druid.indexing.common.task;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Ordering;
import com.google.common.collect.Sets;
import com.google.common.collect.TreeMultiset;
import com.google.common.primitives.Ints;
import com.metamx.common.logger.Logger;
import io.druid.data.input.Firehose;
import io.druid.data.input.FirehoseFactory;
import io.druid.data.input.InputRow;
import io.druid.data.input.impl.SpatialDimensionSchema;
import io.druid.granularity.QueryGranularity;
import io.druid.indexer.granularity.GranularitySpec;
import io.druid.indexing.common.TaskLock;
import io.druid.indexing.common.TaskStatus;
import io.druid.indexing.common.TaskToolbox;
import io.druid.indexing.common.index.YeOldePlumberSchool;
import io.druid.query.aggregation.AggregatorFactory;
import io.druid.segment.loading.DataSegmentPusher;
import io.druid.segment.realtime.FireDepartmentMetrics;
import io.druid.segment.realtime.Schema;
import io.druid.segment.realtime.plumber.Plumber;
import io.druid.segment.realtime.plumber.Sink;
import io.druid.timeline.DataSegment;
import io.druid.timeline.partition.NoneShardSpec;
import io.druid.timeline.partition.ShardSpec;
import io.druid.timeline.partition.SingleDimensionShardSpec;
import org.joda.time.DateTime;
import org.joda.time.Interval;

import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * Simple single-threaded indexing task. Meant to be easy to use for small and medium sized datasets. For "big data",
 * try launching multiple IndexTasks or using the {@link HadoopIndexTask}.
 */
public class IndexTask extends AbstractFixedIntervalTask
{
  private static final Logger log = new Logger(IndexTask.class);

  @JsonIgnore
  private final GranularitySpec granularitySpec;

  @JsonProperty
  private final List<SpatialDimensionSchema> spatialDimensions;

  @JsonIgnore
  private final AggregatorFactory[] aggregators;

  @JsonIgnore
  private final QueryGranularity indexGranularity;

  @JsonIgnore
  private final int targetPartitionSize;

  @JsonIgnore
  private final FirehoseFactory firehoseFactory;

  @JsonIgnore
  private final int rowFlushBoundary;

  @JsonCreator
  public IndexTask(
      @JsonProperty("id") String id,
      @JsonProperty("dataSource") String dataSource,
      @JsonProperty("granularitySpec") GranularitySpec granularitySpec,
      @JsonProperty("spatialDimensions") List<SpatialDimensionSchema> spatialDimensions,
      @JsonProperty("aggregators") AggregatorFactory[] aggregators,
      @JsonProperty("indexGranularity") QueryGranularity indexGranularity,
      @JsonProperty("targetPartitionSize") int targetPartitionSize,
      @JsonProperty("firehose") FirehoseFactory firehoseFactory,
      @JsonProperty("rowFlushBoundary") int rowFlushBoundary
  )
  {
    super(
        // _not_ the version, just something uniqueish
        id != null ? id : String.format("index_%s_%s", dataSource, new DateTime().toString()),
        dataSource,
        new Interval(
            granularitySpec.bucketIntervals().first().getStart(),
            granularitySpec.bucketIntervals().last().getEnd()
        )
    );

    this.granularitySpec = Preconditions.checkNotNull(granularitySpec, "granularitySpec");
    this.spatialDimensions = (spatialDimensions == null)
                             ? Lists.<SpatialDimensionSchema>newArrayList()
                             : spatialDimensions;
    this.aggregators = aggregators;
    this.indexGranularity = (indexGranularity == null) ? QueryGranularity.NONE : indexGranularity;
    this.targetPartitionSize = targetPartitionSize;
    this.firehoseFactory = Preconditions.checkNotNull(firehoseFactory, "firehoseFactory");
    this.rowFlushBoundary = rowFlushBoundary;
  }

  @Override
  public String getType()
  {
    return "index";
  }

  @Override
  public TaskStatus run(TaskToolbox toolbox) throws Exception
  {
    final TaskLock myLock = Iterables.getOnlyElement(getTaskLocks(toolbox));

    final Map<Interval, List<ShardSpec>> shardSpecMap = determinePartitions(targetPartitionSize);

    final Map<Interval, List<Schema>> schemass = Maps.transformEntries(
        shardSpecMap,
        new Maps.EntryTransformer<Interval, List<ShardSpec>, List<Schema>>()
        {
          @Override
          public List<Schema> transformEntry(
              Interval key, List<ShardSpec> shardSpecs
          )
          {
            return Lists.transform(
                shardSpecs,
                new Function<ShardSpec, Schema>()
                {
                  @Override
                  public Schema apply(final ShardSpec shardSpec)
                  {
                    return new Schema(getDataSource(), spatialDimensions, aggregators, indexGranularity, shardSpec);
                  }
                }
            );
          }
        }
    );

    final Set<DataSegment> segments = generateSegments(toolbox, schemass, myLock.getVersion());
    toolbox.pushSegments(segments);

    return TaskStatus.success(getId());
  }

  private Map<Interval, List<ShardSpec>> determinePartitions(
      final int targetPartitionSize
  ) throws IOException
  {
    Map<Interval, List<ShardSpec>> retVal = Maps.newLinkedHashMap();

    log.info("Determining partitions with targetPartitionSize[%d]", targetPartitionSize);

    // Load data
    try (Firehose firehose = firehoseFactory.connect()) {
      if (!firehose.hasMore()) {
        log.error("Unable to find any events to ingest! Check your firehose config!");
        return retVal;
      }
      InputRow inputRow = firehose.nextRow();

      for (Interval interval : granularitySpec.bucketIntervals()) {
        // Blacklist dimensions that have multiple values per row
        final Set<String> unusableDimensions = Sets.newHashSet();

        // Track values of all non-blacklisted dimensions
        final Map<String, TreeMultiset<String>> dimensionValueMultisets = Maps.newHashMap();

        boolean hasEventsInInterval = false;
        boolean done = false;
        while (!done && interval.contains(inputRow.getTimestampFromEpoch())) {
          hasEventsInInterval = true;

          // Extract dimensions from event
          for (final String dim : inputRow.getDimensions()) {
            final List<String> dimValues = inputRow.getDimension(dim);
            if (!unusableDimensions.contains(dim)) {
              if (dimValues.size() == 1) {
                // Track this value
                TreeMultiset<String> dimensionValueMultiset = dimensionValueMultisets.get(dim);
                if (dimensionValueMultiset == null) {
                  dimensionValueMultiset = TreeMultiset.create();
                  dimensionValueMultisets.put(dim, dimensionValueMultiset);
                }
                dimensionValueMultiset.add(dimValues.get(0));
              } else {
                // Only single-valued dimensions can be used for partitions
                unusableDimensions.add(dim);
                dimensionValueMultisets.remove(dim);
              }
{'fraction_non_alphanumeric': 0.09009442473085945, 'fraction_numerical': 0.0026462981896914657, 'mean_word_length': 3.410610079575597, 'pattern_counts': {'":': 0, '<': 49, '<?xml version=': 0, '>': 49, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 6}, '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': '5938482', 'n_tokens_mistral': 4752, 'n_tokens_neox': 4592, 'n_words': 1103}