contrib/storage-kafka/src/main/java/org/apache/drill/exec/store/kafka/KafkaGroupScan.java - drill - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  * http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 package org.apache.drill.exec.store.kafka;

 import java.time.Duration;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;

 import org.apache.drill.exec.ExecConstants;
 import com.google.common.collect.Sets;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.drill.common.PlanStringBuilder;
 import org.apache.drill.common.exceptions.ExecutionSetupException;
 import org.apache.drill.common.exceptions.UserException;
 import org.apache.drill.common.expression.SchemaPath;
 import org.apache.drill.exec.physical.EndpointAffinity;
 import org.apache.drill.exec.physical.base.AbstractGroupScan;
 import org.apache.drill.exec.physical.base.GroupScan;
 import org.apache.drill.exec.physical.base.PhysicalOperator;
 import org.apache.drill.exec.physical.base.ScanStats;
 import org.apache.drill.exec.physical.base.ScanStats.GroupScanProperty;
 import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
 import org.apache.drill.exec.store.StoragePluginRegistry;
 import org.apache.drill.exec.store.schedule.AffinityCreator;
 import org.apache.drill.exec.store.schedule.AssignmentCreator;
 import org.apache.drill.exec.store.schedule.CompleteWork;
 import org.apache.drill.exec.store.schedule.EndpointByteMap;
 import org.apache.drill.exec.store.schedule.EndpointByteMapImpl;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import com.fasterxml.jackson.annotation.JacksonInject;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.annotation.JsonTypeName;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ListMultimap;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;

 @JsonTypeName("kafka-scan")
 public class KafkaGroupScan extends AbstractGroupScan {
   private static final Logger logger = LoggerFactory.getLogger(KafkaGroupScan.class);

   // Assuming default average topic message size as 1KB, which will be used to
   // compute the stats and work assignments
   private static final long MSG_SIZE = 1024;

   private final KafkaStoragePlugin kafkaStoragePlugin;
   private final KafkaScanSpec kafkaScanSpec;

   private final List<SchemaPath> columns;
   private final int records;
   private ListMultimap<Integer, PartitionScanWork> assignments;
   private List<EndpointAffinity> affinities;

   private Map<TopicPartition, PartitionScanWork> partitionWorkMap;

   @JsonCreator
   public KafkaGroupScan(@JsonProperty("userName") String userName,
                         @JsonProperty("kafkaStoragePluginConfig") KafkaStoragePluginConfig kafkaStoragePluginConfig,
                         @JsonProperty("columns") List<SchemaPath> columns,
                         @JsonProperty("records") int records,
                         @JsonProperty("kafkaScanSpec") KafkaScanSpec scanSpec,
                         @JacksonInject StoragePluginRegistry pluginRegistry) throws ExecutionSetupException {
     this(userName,
         pluginRegistry.resolve(kafkaStoragePluginConfig, KafkaStoragePlugin.class),
         columns,
         records,
         scanSpec);
   }

   public KafkaGroupScan(KafkaStoragePlugin kafkaStoragePlugin, KafkaScanSpec kafkaScanSpec, List<SchemaPath> columns, int records) {
     super(StringUtils.EMPTY);
     this.kafkaStoragePlugin = kafkaStoragePlugin;
     this.columns = columns;
     this.records = records;
     this.kafkaScanSpec = kafkaScanSpec;
     init();
   }

   public KafkaGroupScan(String userName,
                         KafkaStoragePlugin kafkaStoragePlugin,
                         List<SchemaPath> columns,
                         int records,
                         KafkaScanSpec kafkaScanSpec) {
     super(userName);
     this.kafkaStoragePlugin = kafkaStoragePlugin;
     this.columns = columns;
     this.records = records;
     this.kafkaScanSpec = kafkaScanSpec;
     init();
   }

   public KafkaGroupScan(KafkaGroupScan that) {
     super(that);
     this.kafkaStoragePlugin = that.kafkaStoragePlugin;
     this.columns = that.columns;
     this.records = that.records;
     this.kafkaScanSpec = that.kafkaScanSpec;
     this.assignments = that.assignments;
     this.partitionWorkMap = that.partitionWorkMap;
   }

   public KafkaGroupScan(KafkaGroupScan that, List<SchemaPath> columns) {
     super(that);
     this.kafkaStoragePlugin = that.kafkaStoragePlugin;
     this.columns = columns;
     this.records = that.records;
     this.kafkaScanSpec = that.kafkaScanSpec;
     this.assignments = that.assignments;
     this.partitionWorkMap = that.partitionWorkMap;
   }

   public KafkaGroupScan(KafkaGroupScan that, int records) {
     super(that);
     this.kafkaStoragePlugin = that.kafkaStoragePlugin;
     this.columns = that.columns;
     this.records = records;
     this.kafkaScanSpec = that.kafkaScanSpec;
     this.assignments = that.assignments;
     this.partitionWorkMap = that.partitionWorkMap;
   }

   public static class PartitionScanWork implements CompleteWork {

     private final EndpointByteMapImpl byteMap;
     private final KafkaPartitionScanSpec partitionScanSpec;

     public PartitionScanWork(EndpointByteMap byteMap, KafkaPartitionScanSpec partitionScanSpec) {
       this.byteMap = (EndpointByteMapImpl)byteMap;
       this.partitionScanSpec = partitionScanSpec;
     }

     @Override
     public int compareTo(CompleteWork o) {
       return Long.compare(getTotalBytes(), o.getTotalBytes());
     }

     @Override
     public long getTotalBytes() {
       return (partitionScanSpec.getEndOffset() - partitionScanSpec.getStartOffset()) * MSG_SIZE;
     }

     @Override
     public EndpointByteMap getByteMap() {
       return byteMap;
     }

     public KafkaPartitionScanSpec getPartitionScanSpec() {
       return partitionScanSpec;
     }
   }

   /**
    * Computes work per topic partition, based on start and end offset of each
    * corresponding topicPartition
    */
   private void init() {
     partitionWorkMap = Maps.newHashMap();
     Collection<DrillbitEndpoint> endpoints = kafkaStoragePlugin.getContext().getBits();
     Map<String, DrillbitEndpoint> endpointMap = endpoints.stream()
       .collect(Collectors.toMap(
         DrillbitEndpoint::getAddress,
         Function.identity(),
         (o, n) -> n));

     Map<TopicPartition, Long> startOffsetsMap = Maps.newHashMap();
     Map<TopicPartition, Long> endOffsetsMap = Maps.newHashMap();
     List<PartitionInfo> topicPartitions;
     String topicName = kafkaScanSpec.getTopicName();

     KafkaConsumer<?, ?> kafkaConsumer = null;
     try {
       kafkaConsumer = new KafkaConsumer<>(kafkaStoragePlugin.getConfig().getKafkaConsumerProps(),
         new ByteArrayDeserializer(), new ByteArrayDeserializer());
       if (!kafkaConsumer.listTopics().containsKey(topicName)) {
         throw UserException.dataReadError()
             .message("Table '%s' does not exist", topicName)
             .build(logger);
       }
       kafkaConsumer.subscribe(Collections.singletonList(topicName));
       // based on KafkaConsumer JavaDoc, seekToBeginning/seekToEnd functions
       // evaluates lazily, seeking to the first/last offset in all partitions only
       // when poll(long) or
       // position(TopicPartition) are called
       kafkaConsumer.poll(Duration.ofSeconds(5));
       Set<TopicPartition> assignments = waitForConsumerAssignment(kafkaConsumer);
       topicPartitions = kafkaConsumer.partitionsFor(topicName);

       // fetch start offsets for each topicPartition
       kafkaConsumer.seekToBeginning(assignments);
       for (TopicPartition topicPartition : assignments) {
         startOffsetsMap.put(topicPartition, kafkaConsumer.position(topicPartition));
       }

       // fetch end offsets for each topicPartition
       kafkaConsumer.seekToEnd(assignments);
       for (TopicPartition topicPartition : assignments) {
         endOffsetsMap.put(topicPartition, kafkaConsumer.position(topicPartition));
       }
     } catch (Exception e) {
       throw UserException.dataReadError(e)
         .message("Failed to fetch start/end offsets of the topic %s", topicName)
         .addContext(e.getMessage())
         .build(logger);
     } finally {
       kafkaStoragePlugin.registerToClose(kafkaConsumer);
     }

     // computes work for each end point
     for (PartitionInfo partitionInfo : topicPartitions) {
       TopicPartition topicPartition = new TopicPartition(partitionInfo.topic(), partitionInfo.partition());
       long lastCommittedOffset = startOffsetsMap.get(topicPartition);
       long latestOffset = endOffsetsMap.get(topicPartition);
       logger.debug("Latest offset of {} is {}", topicPartition, latestOffset);
       logger.debug("Last committed offset of {} is {}", topicPartition, lastCommittedOffset);
       KafkaPartitionScanSpec partitionScanSpec = new KafkaPartitionScanSpec(topicPartition.topic(), topicPartition.partition(), lastCommittedOffset, latestOffset);
       PartitionScanWork work = new PartitionScanWork(new EndpointByteMapImpl(), partitionScanSpec);
       Node[] inSyncReplicas = partitionInfo.inSyncReplicas();
       for (Node isr : inSyncReplicas) {
         String host = isr.host();
         DrillbitEndpoint ep = endpointMap.get(host);
         if (ep != null) {
           work.getByteMap().add(ep, work.getTotalBytes());
         }
       }
       partitionWorkMap.put(topicPartition, work);
     }
   }


   /** Workaround for Kafka > 2.0 version due to KIP-505.
    * It can be replaced with Kafka implementation once it will be introduced.
    * @param consumer Kafka consumer whom need to get assignments
    * @return
    * @throws InterruptedException
    */
   private Set<TopicPartition> waitForConsumerAssignment(Consumer consumer) throws InterruptedException {
     Set<TopicPartition> assignments = consumer.assignment();

     long waitingForAssigmentTimeout = kafkaStoragePlugin.getContext().getOptionManager().getLong(ExecConstants.KAFKA_POLL_TIMEOUT);
     long timeout = 0;

     while (assignments.isEmpty() && timeout < waitingForAssigmentTimeout) {
       Thread.sleep(500);
       timeout += 500;
       assignments = consumer.assignment();
     }

     if (timeout >= waitingForAssigmentTimeout) {
       throw UserException.dataReadError()
         .message("Consumer assignment wasn't completed within the timeout %s", waitingForAssigmentTimeout)
         .build(logger);
     }

     return assignments;
   }

   @Override
   public void applyAssignments(List<DrillbitEndpoint> incomingEndpoints) {
     assignments = AssignmentCreator.getMappings(incomingEndpoints, Lists.newArrayList(partitionWorkMap.values()));
   }

   @Override
   public GroupScan applyLimit(int maxRecords) {
     if (maxRecords > records) { // pass the limit value into sub-scan
       return new KafkaGroupScan(this, maxRecords);
     } else { // stop the transform
       return null;
     }
   }

   @Override
   public boolean supportsLimitPushdown() {
     return true;
   }

   @Override
   public KafkaSubScan getSpecificScan(int minorFragmentId) {
     List<PartitionScanWork> workList = assignments.get(minorFragmentId);

     List<KafkaPartitionScanSpec> scanSpecList = workList.stream()
       .map(PartitionScanWork::getPartitionScanSpec)
       .collect(Collectors.toList());

     return new KafkaSubScan(getUserName(), kafkaStoragePlugin, columns, records, scanSpecList);
   }

   @Override
   public int getMaxParallelizationWidth() {
     return partitionWorkMap.values().size();
   }

   @Override
   public ScanStats getScanStats() {
     long messageCount = 0;
     for (PartitionScanWork work : partitionWorkMap.values()) {
       messageCount += (work.getPartitionScanSpec().getEndOffset() - work.getPartitionScanSpec().getStartOffset());
     }
     return new ScanStats(GroupScanProperty.EXACT_ROW_COUNT, messageCount, 1, messageCount * MSG_SIZE);
   }

   @Override
   public String getDigest() {
     return toString();
   }

   @Override
   public PhysicalOperator getNewWithChildren(List<PhysicalOperator> children) {
     Preconditions.checkArgument(children.isEmpty());
     return new KafkaGroupScan(this);
   }

   @Override
   public List<EndpointAffinity> getOperatorAffinity() {
     if (affinities == null) {
       affinities = AffinityCreator.getAffinityMap(Lists.newArrayList(partitionWorkMap.values()));
     }
     return affinities;
   }

   @Override
   @JsonIgnore
   public boolean canPushdownProjects(List<SchemaPath> columns) {
     return true;
   }

   @Override
   public GroupScan clone(List<SchemaPath> columns) {
     return new KafkaGroupScan(this, columns);
   }

   public GroupScan cloneWithNewSpec(List<KafkaPartitionScanSpec> partitionScanSpecList) {
     KafkaGroupScan clone = new KafkaGroupScan(this);
     HashSet<TopicPartition> partitionsInSpec = Sets.newHashSet();

     for (KafkaPartitionScanSpec scanSpec : partitionScanSpecList) {
       TopicPartition tp = new TopicPartition(scanSpec.getTopicName(), scanSpec.getPartitionId());
       partitionsInSpec.add(tp);

       PartitionScanWork newScanWork = new PartitionScanWork(partitionWorkMap.get(tp).getByteMap(), scanSpec);
       clone.partitionWorkMap.put(tp, newScanWork);
     }

     //Remove unnecessary partitions from partitionWorkMap
     clone.partitionWorkMap.keySet().removeIf(tp -> !partitionsInSpec.contains(tp));
     return clone;
   }

   @JsonProperty
   public KafkaStoragePluginConfig getKafkaStoragePluginConfig() {
     return kafkaStoragePlugin.getConfig();
   }

   @Override
   @JsonProperty
   public List<SchemaPath> getColumns() {
     return columns;
   }

   @JsonProperty
   public int getRecords() {
     return records;
   }

   @JsonProperty
   public KafkaScanSpec getKafkaScanSpec() {
     return kafkaScanSpec;
   }

   @JsonIgnore
   public KafkaStoragePlugin getStoragePlugin() {
     return kafkaStoragePlugin;
   }

   @Override
   public String toString() {
     return new PlanStringBuilder("")
         .field("scanSpec", kafkaScanSpec)
         .field("columns", columns)
         .field("records", records)
         .toString();
   }

   @JsonIgnore
   public List<KafkaPartitionScanSpec> getPartitionScanSpecList() {
     return partitionWorkMap.values().stream()
       .map(work -> work.partitionScanSpec.clone())
       .collect(Collectors.toList());
   }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/
	package org.apache.drill.exec.store.kafka;

	import java.time.Duration;
	import java.util.Collection;
	import java.util.Collections;
	import java.util.HashSet;
	import java.util.List;
	import java.util.Map;
	import java.util.Set;
	import java.util.function.Function;
	import java.util.stream.Collectors;

	import org.apache.drill.exec.ExecConstants;
	import com.google.common.collect.Sets;
	import org.apache.commons.lang3.StringUtils;
	import org.apache.drill.common.PlanStringBuilder;
	import org.apache.drill.common.exceptions.ExecutionSetupException;
	import org.apache.drill.common.exceptions.UserException;
	import org.apache.drill.common.expression.SchemaPath;
	import org.apache.drill.exec.physical.EndpointAffinity;
	import org.apache.drill.exec.physical.base.AbstractGroupScan;
	import org.apache.drill.exec.physical.base.GroupScan;
	import org.apache.drill.exec.physical.base.PhysicalOperator;
	import org.apache.drill.exec.physical.base.ScanStats;
	import org.apache.drill.exec.physical.base.ScanStats.GroupScanProperty;
	import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
	import org.apache.drill.exec.store.StoragePluginRegistry;
	import org.apache.drill.exec.store.schedule.AffinityCreator;
	import org.apache.drill.exec.store.schedule.AssignmentCreator;
	import org.apache.drill.exec.store.schedule.CompleteWork;
	import org.apache.drill.exec.store.schedule.EndpointByteMap;
	import org.apache.drill.exec.store.schedule.EndpointByteMapImpl;
	import org.apache.kafka.clients.consumer.Consumer;
	import org.apache.kafka.clients.consumer.KafkaConsumer;
	import org.apache.kafka.common.Node;
	import org.apache.kafka.common.PartitionInfo;
	import org.apache.kafka.common.TopicPartition;
	import org.apache.kafka.common.serialization.ByteArrayDeserializer;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import com.fasterxml.jackson.annotation.JacksonInject;
	import com.fasterxml.jackson.annotation.JsonCreator;
	import com.fasterxml.jackson.annotation.JsonIgnore;
	import com.fasterxml.jackson.annotation.JsonProperty;
	import com.fasterxml.jackson.annotation.JsonTypeName;
	import com.google.common.base.Preconditions;
	import com.google.common.collect.ListMultimap;
	import com.google.common.collect.Lists;
	import com.google.common.collect.Maps;

	@JsonTypeName("kafka-scan")
	public class KafkaGroupScan extends AbstractGroupScan {
	private static final Logger logger = LoggerFactory.getLogger(KafkaGroupScan.class);

	// Assuming default average topic message size as 1KB, which will be used to
	// compute the stats and work assignments
	private static final long MSG_SIZE = 1024;

	private final KafkaStoragePlugin kafkaStoragePlugin;
	private final KafkaScanSpec kafkaScanSpec;

	private final List<SchemaPath> columns;
	private final int records;
	private ListMultimap<Integer, PartitionScanWork> assignments;
	private List<EndpointAffinity> affinities;

	private Map<TopicPartition, PartitionScanWork> partitionWorkMap;

	@JsonCreator
	public KafkaGroupScan(@JsonProperty("userName") String userName,
	@JsonProperty("kafkaStoragePluginConfig") KafkaStoragePluginConfig kafkaStoragePluginConfig,
	@JsonProperty("columns") List<SchemaPath> columns,
	@JsonProperty("records") int records,
	@JsonProperty("kafkaScanSpec") KafkaScanSpec scanSpec,
	@JacksonInject StoragePluginRegistry pluginRegistry) throws ExecutionSetupException {
	this(userName,
	pluginRegistry.resolve(kafkaStoragePluginConfig, KafkaStoragePlugin.class),
	columns,
	records,
	scanSpec);
	}

	public KafkaGroupScan(KafkaStoragePlugin kafkaStoragePlugin, KafkaScanSpec kafkaScanSpec, List<SchemaPath> columns, int records) {
	super(StringUtils.EMPTY);
	this.kafkaStoragePlugin = kafkaStoragePlugin;
	this.columns = columns;
	this.records = records;
	this.kafkaScanSpec = kafkaScanSpec;
	init();
	}

	public KafkaGroupScan(String userName,
	KafkaStoragePlugin kafkaStoragePlugin,
	List<SchemaPath> columns,
	int records,
	KafkaScanSpec kafkaScanSpec) {
	super(userName);
	this.kafkaStoragePlugin = kafkaStoragePlugin;
	this.columns = columns;
	this.records = records;
	this.kafkaScanSpec = kafkaScanSpec;
	init();
	}

	public KafkaGroupScan(KafkaGroupScan that) {
	super(that);
	this.kafkaStoragePlugin = that.kafkaStoragePlugin;
	this.columns = that.columns;
	this.records = that.records;
	this.kafkaScanSpec = that.kafkaScanSpec;
	this.assignments = that.assignments;
	this.partitionWorkMap = that.partitionWorkMap;
	}

	public KafkaGroupScan(KafkaGroupScan that, List<SchemaPath> columns) {
	super(that);
	this.kafkaStoragePlugin = that.kafkaStoragePlugin;
	this.columns = columns;
	this.records = that.records;
	this.kafkaScanSpec = that.kafkaScanSpec;
	this.assignments = that.assignments;
	this.partitionWorkMap = that.partitionWorkMap;
	}

	public KafkaGroupScan(KafkaGroupScan that, int records) {
	super(that);
	this.kafkaStoragePlugin = that.kafkaStoragePlugin;
	this.columns = that.columns;
	this.records = records;
	this.kafkaScanSpec = that.kafkaScanSpec;
	this.assignments = that.assignments;
	this.partitionWorkMap = that.partitionWorkMap;
	}

	public static class PartitionScanWork implements CompleteWork {

	private final EndpointByteMapImpl byteMap;
	private final KafkaPartitionScanSpec partitionScanSpec;

	public PartitionScanWork(EndpointByteMap byteMap, KafkaPartitionScanSpec partitionScanSpec) {
	this.byteMap = (EndpointByteMapImpl)byteMap;
	this.partitionScanSpec = partitionScanSpec;
	}

	@Override
	public int compareTo(CompleteWork o) {
	return Long.compare(getTotalBytes(), o.getTotalBytes());
	}

	@Override
	public long getTotalBytes() {
	return (partitionScanSpec.getEndOffset() - partitionScanSpec.getStartOffset()) * MSG_SIZE;
	}

	@Override
	public EndpointByteMap getByteMap() {
	return byteMap;
	}

	public KafkaPartitionScanSpec getPartitionScanSpec() {
	return partitionScanSpec;
	}
	}

	/**
	* Computes work per topic partition, based on start and end offset of each
	* corresponding topicPartition
	*/
	private void init() {
	partitionWorkMap = Maps.newHashMap();
	Collection<DrillbitEndpoint> endpoints = kafkaStoragePlugin.getContext().getBits();
	Map<String, DrillbitEndpoint> endpointMap = endpoints.stream()
	.collect(Collectors.toMap(
	DrillbitEndpoint::getAddress,
	Function.identity(),
	(o, n) -> n));

	Map<TopicPartition, Long> startOffsetsMap = Maps.newHashMap();
	Map<TopicPartition, Long> endOffsetsMap = Maps.newHashMap();
	List<PartitionInfo> topicPartitions;
	String topicName = kafkaScanSpec.getTopicName();

	KafkaConsumer<?, ?> kafkaConsumer = null;
	try {
	kafkaConsumer = new KafkaConsumer<>(kafkaStoragePlugin.getConfig().getKafkaConsumerProps(),
	new ByteArrayDeserializer(), new ByteArrayDeserializer());
	if (!kafkaConsumer.listTopics().containsKey(topicName)) {
	throw UserException.dataReadError()
	.message("Table '%s' does not exist", topicName)
	.build(logger);
	}
	kafkaConsumer.subscribe(Collections.singletonList(topicName));
	// based on KafkaConsumer JavaDoc, seekToBeginning/seekToEnd functions
	// evaluates lazily, seeking to the first/last offset in all partitions only
	// when poll(long) or
	// position(TopicPartition) are called
	kafkaConsumer.poll(Duration.ofSeconds(5));
	Set<TopicPartition> assignments = waitForConsumerAssignment(kafkaConsumer);
	topicPartitions = kafkaConsumer.partitionsFor(topicName);

	// fetch start offsets for each topicPartition
	kafkaConsumer.seekToBeginning(assignments);
	for (TopicPartition topicPartition : assignments) {
	startOffsetsMap.put(topicPartition, kafkaConsumer.position(topicPartition));
	}

	// fetch end offsets for each topicPartition
	kafkaConsumer.seekToEnd(assignments);
	for (TopicPartition topicPartition : assignments) {
	endOffsetsMap.put(topicPartition, kafkaConsumer.position(topicPartition));
	}
	} catch (Exception e) {
	throw UserException.dataReadError(e)
	.message("Failed to fetch start/end offsets of the topic %s", topicName)
	.addContext(e.getMessage())
	.build(logger);
	} finally {
	kafkaStoragePlugin.registerToClose(kafkaConsumer);
	}

	// computes work for each end point
	for (PartitionInfo partitionInfo : topicPartitions) {
	TopicPartition topicPartition = new TopicPartition(partitionInfo.topic(), partitionInfo.partition());
	long lastCommittedOffset = startOffsetsMap.get(topicPartition);
	long latestOffset = endOffsetsMap.get(topicPartition);
	logger.debug("Latest offset of {} is {}", topicPartition, latestOffset);
	logger.debug("Last committed offset of {} is {}", topicPartition, lastCommittedOffset);
	KafkaPartitionScanSpec partitionScanSpec = new KafkaPartitionScanSpec(topicPartition.topic(), topicPartition.partition(), lastCommittedOffset, latestOffset);
	PartitionScanWork work = new PartitionScanWork(new EndpointByteMapImpl(), partitionScanSpec);
	Node[] inSyncReplicas = partitionInfo.inSyncReplicas();
	for (Node isr : inSyncReplicas) {
	String host = isr.host();
	DrillbitEndpoint ep = endpointMap.get(host);
	if (ep != null) {
	work.getByteMap().add(ep, work.getTotalBytes());
	}
	}
	partitionWorkMap.put(topicPartition, work);
	}
	}


	/** Workaround for Kafka > 2.0 version due to KIP-505.
	* It can be replaced with Kafka implementation once it will be introduced.
	* @param consumer Kafka consumer whom need to get assignments
	* @return
	* @throws InterruptedException
	*/
	private Set<TopicPartition> waitForConsumerAssignment(Consumer consumer) throws InterruptedException {
	Set<TopicPartition> assignments = consumer.assignment();

	long waitingForAssigmentTimeout = kafkaStoragePlugin.getContext().getOptionManager().getLong(ExecConstants.KAFKA_POLL_TIMEOUT);
	long timeout = 0;

	while (assignments.isEmpty() && timeout < waitingForAssigmentTimeout) {
	Thread.sleep(500);
	timeout += 500;
	assignments = consumer.assignment();
	}

	if (timeout >= waitingForAssigmentTimeout) {
	throw UserException.dataReadError()
	.message("Consumer assignment wasn't completed within the timeout %s", waitingForAssigmentTimeout)
	.build(logger);
	}

	return assignments;
	}

	@Override
	public void applyAssignments(List<DrillbitEndpoint> incomingEndpoints) {
	assignments = AssignmentCreator.getMappings(incomingEndpoints, Lists.newArrayList(partitionWorkMap.values()));
	}

	@Override
	public GroupScan applyLimit(int maxRecords) {
	if (maxRecords > records) { // pass the limit value into sub-scan
	return new KafkaGroupScan(this, maxRecords);
	} else { // stop the transform
	return null;
	}
	}

	@Override
	public boolean supportsLimitPushdown() {
	return true;
	}

	@Override
	public KafkaSubScan getSpecificScan(int minorFragmentId) {
	List<PartitionScanWork> workList = assignments.get(minorFragmentId);

	List<KafkaPartitionScanSpec> scanSpecList = workList.stream()
	.map(PartitionScanWork::getPartitionScanSpec)
	.collect(Collectors.toList());

	return new KafkaSubScan(getUserName(), kafkaStoragePlugin, columns, records, scanSpecList);
	}

	@Override
	public int getMaxParallelizationWidth() {
	return partitionWorkMap.values().size();
	}

	@Override
	public ScanStats getScanStats() {
	long messageCount = 0;
	for (PartitionScanWork work : partitionWorkMap.values()) {
	messageCount += (work.getPartitionScanSpec().getEndOffset() - work.getPartitionScanSpec().getStartOffset());
	}
	return new ScanStats(GroupScanProperty.EXACT_ROW_COUNT, messageCount, 1, messageCount * MSG_SIZE);
	}

	@Override
	public String getDigest() {
	return toString();
	}

	@Override
	public PhysicalOperator getNewWithChildren(List<PhysicalOperator> children) {
	Preconditions.checkArgument(children.isEmpty());
	return new KafkaGroupScan(this);
	}

	@Override
	public List<EndpointAffinity> getOperatorAffinity() {
	if (affinities == null) {
	affinities = AffinityCreator.getAffinityMap(Lists.newArrayList(partitionWorkMap.values()));
	}
	return affinities;
	}

	@Override
	@JsonIgnore
	public boolean canPushdownProjects(List<SchemaPath> columns) {
	return true;
	}

	@Override
	public GroupScan clone(List<SchemaPath> columns) {
	return new KafkaGroupScan(this, columns);
	}

	public GroupScan cloneWithNewSpec(List<KafkaPartitionScanSpec> partitionScanSpecList) {
	KafkaGroupScan clone = new KafkaGroupScan(this);
	HashSet<TopicPartition> partitionsInSpec = Sets.newHashSet();

	for (KafkaPartitionScanSpec scanSpec : partitionScanSpecList) {
	TopicPartition tp = new TopicPartition(scanSpec.getTopicName(), scanSpec.getPartitionId());
	partitionsInSpec.add(tp);

	PartitionScanWork newScanWork = new PartitionScanWork(partitionWorkMap.get(tp).getByteMap(), scanSpec);
	clone.partitionWorkMap.put(tp, newScanWork);
	}

	//Remove unnecessary partitions from partitionWorkMap
	clone.partitionWorkMap.keySet().removeIf(tp -> !partitionsInSpec.contains(tp));
	return clone;
	}

	@JsonProperty
	public KafkaStoragePluginConfig getKafkaStoragePluginConfig() {
	return kafkaStoragePlugin.getConfig();
	}

	@Override
	@JsonProperty
	public List<SchemaPath> getColumns() {
	return columns;
	}

	@JsonProperty
	public int getRecords() {
	return records;
	}

	@JsonProperty
	public KafkaScanSpec getKafkaScanSpec() {
	return kafkaScanSpec;
	}

	@JsonIgnore
	public KafkaStoragePlugin getStoragePlugin() {
	return kafkaStoragePlugin;
	}

	@Override
	public String toString() {
	return new PlanStringBuilder("")
	.field("scanSpec", kafkaScanSpec)
	.field("columns", columns)
	.field("records", records)
	.toString();
	}

	@JsonIgnore
	public List<KafkaPartitionScanSpec> getPartitionScanSpecList() {
	return partitionWorkMap.values().stream()
	.map(work -> work.partitionScanSpec.clone())
	.collect(Collectors.toList());
	}
	}