easy-cs-rec-custommodel 0.8.6__py2.py3-none-any.whl

easy_rec/python/input/kafka_dataset.py

@@ -0,0 +1,144 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed 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.
+# ==============================================================================
+"""Kafka Dataset."""
+
+import logging
+import traceback
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+
+try:
+  from easy_rec.python.ops import gen_kafka_ops
+except ImportError:
+  logging.warning('failed to import gen_kafka_ops: %s' % traceback.format_exc())
+
+
+class KafkaDataset(dataset_ops.Dataset):
+  """A Kafka Dataset that consumes the message."""
+
+  def __init__(self,
+               topics,
+               servers='localhost',
+               group='',
+               eof=False,
+               timeout=1000,
+               config_global=None,
+               config_topic=None,
+               message_key=False,
+               message_offset=False):
+    """Create a KafkaReader.
+
+    Args:
+      topics: A `tf.string` tensor containing one or more subscriptions,
+              in the format of [topic:partition:offset:length],
+              by default length is -1 for unlimited.
+      servers: A list of bootstrap servers.
+      group: The consumer group id.
+      eof: If True, the kafka reader will stop on EOF.
+      timeout: The timeout value for the Kafka Consumer to wait
+               (in millisecond).
+      config_global: A `tf.string` tensor containing global configuration
+                     properties in [Key=Value] format,
+                     eg. ["enable.auto.commit=false",
+                          "heartbeat.interval.ms=2000"],
+                     please refer to 'Global configuration properties'
+                     in librdkafka doc.
+      config_topic: A `tf.string` tensor containing topic configuration
+                    properties in [Key=Value] format,
+                    eg. ["auto.offset.reset=earliest"],
+                    please refer to 'Topic configuration properties'
+                    in librdkafka doc.
+      message_key: If True, the kafka will output both message value and key.
+      message_offset: If True, the kafka will output both message value and offset.
+    """
+    self._topics = ops.convert_to_tensor(
+        topics, dtype=dtypes.string, name='topics')
+    self._servers = ops.convert_to_tensor(
+        servers, dtype=dtypes.string, name='servers')
+    self._group = ops.convert_to_tensor(
+        group, dtype=dtypes.string, name='group')
+    self._eof = ops.convert_to_tensor(eof, dtype=dtypes.bool, name='eof')
+    self._timeout = ops.convert_to_tensor(
+        timeout, dtype=dtypes.int64, name='timeout')
+    config_global = config_global if config_global else []
+    self._config_global = ops.convert_to_tensor(
+        config_global, dtype=dtypes.string, name='config_global')
+    config_topic = config_topic if config_topic else []
+    self._config_topic = ops.convert_to_tensor(
+        config_topic, dtype=dtypes.string, name='config_topic')
+    self._message_key = message_key
+    self._message_offset = message_offset
+    super(KafkaDataset, self).__init__()
+
+  def _inputs(self):
+    return []
+
+  def _as_variant_tensor(self):
+    return gen_kafka_ops.io_kafka_dataset_v2(
+        self._topics,
+        self._servers,
+        self._group,
+        self._eof,
+        self._timeout,
+        self._config_global,
+        self._config_topic,
+        self._message_key,
+        self._message_offset,
+    )
+
+  @property
+  def output_classes(self):
+    if self._message_key ^ self._message_offset:
+      return (ops.Tensor, ops.Tensor)
+    elif self._message_key and self._message_offset:
+      return (ops.Tensor, ops.Tensor, ops.Tensor)
+    return (ops.Tensor)
+
+  @property
+  def output_shapes(self):
+    if self._message_key ^ self._message_offset:
+      return ((tensor_shape.TensorShape([]), tensor_shape.TensorShape([])))
+    elif self._message_key and self._message_offset:
+      return ((tensor_shape.TensorShape([]), tensor_shape.TensorShape([]),
+               tensor_shape.TensorShape([])))
+    return ((tensor_shape.TensorShape([])))
+
+  @property
+  def output_types(self):
+    if self._message_key ^ self._message_offset:
+      return ((dtypes.string, dtypes.string))
+    elif self._message_key and self._message_offset:
+      return ((dtypes.string, dtypes.string, dtypes.string))
+    return ((dtypes.string))
+
+
+def write_kafka_v2(message, topic, servers='localhost', name=None):
+  """Write kafka.
+
+  Args:
+    message: A `Tensor` of type `string`. 0-D.
+    topic: A `tf.string` tensor containing one subscription,
+      in the format of topic:partition.
+    servers: A list of bootstrap servers.
+    name: A name for the operation (optional).
+
+  Returns:
+    A `Tensor` of type `string`. 0-D.
+  """
+  return gen_kafka_ops.io_write_kafka_v2(
+      message=message, topic=topic, servers=servers, name=name)

easy_rec/python/input/kafka_input.py

@@ -0,0 +1,235 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import json
+import logging
+import traceback
+
+import six
+import tensorflow as tf
+
+from easy_rec.python.input.input import Input
+from easy_rec.python.input.kafka_dataset import KafkaDataset
+from easy_rec.python.utils.config_util import parse_time
+
+if tf.__version__.startswith('1.'):
+  from tensorflow.python.platform import gfile
+else:
+  import tensorflow.io.gfile as gfile
+
+try:
+  from kafka import KafkaConsumer, TopicPartition
+except ImportError:
+  logging.warning(
+      'kafka-python is not installed[%s]. You can install it by: pip install kafka-python'
+      % traceback.format_exc())
+
+if tf.__version__ >= '2.0':
+  ignore_errors = tf.data.experimental.ignore_errors()
+  tf = tf.compat.v1
+else:
+  ignore_errors = tf.contrib.data.ignore_errors()
+
+
+class KafkaInput(Input):
+
+  DATA_OFFSET = 'DATA_OFFSET'
+
+  def __init__(self,
+               data_config,
+               feature_config,
+               kafka_config,
+               task_index=0,
+               task_num=1,
+               check_mode=False,
+               pipeline_config=None):
+    super(KafkaInput,
+          self).__init__(data_config, feature_config, '', task_index, task_num,
+                         check_mode, pipeline_config)
+    self._kafka = kafka_config
+    self._offset_dict = {}
+    if self._kafka is not None:
+      consumer = KafkaConsumer(
+          group_id='kafka_dataset_consumer',
+          bootstrap_servers=[self._kafka.server],
+          api_version_auto_timeout_ms=60000)  # in miliseconds
+      partitions = consumer.partitions_for_topic(self._kafka.topic)
+      self._num_partition = len(partitions)
+      logging.info('all partitions[%d]: %s' % (self._num_partition, partitions))
+
+      # determine kafka offsets for each partition
+      offset_type = self._kafka.WhichOneof('offset')
+      if offset_type is not None:
+        if offset_type == 'offset_time':
+          ts = parse_time(self._kafka.offset_time)
+          input_map = {
+              TopicPartition(partition=part_id, topic=self._kafka.topic):
+              ts * 1000 for part_id in partitions
+          }
+          part_offsets = consumer.offsets_for_times(input_map)
+          # part_offsets is a dictionary:
+          # {
+          #    TopicPartition(topic=u'kafka_data_20220408', partition=0):
+          #       OffsetAndTimestamp(offset=2, timestamp=1650611437895)
+          # }
+          for part in part_offsets:
+            self._offset_dict[part.partition] = part_offsets[part].offset
+            logging.info(
+                'Find offset by time, topic[%s], partition[%d], timestamp[%ss], offset[%d], offset_timestamp[%dms]'
+                % (self._kafka.topic, part.partition, ts,
+                   part_offsets[part].offset, part_offsets[part].timestamp))
+        elif offset_type == 'offset_info':
+          offset_dict = json.loads(self._kafka.offset_info)
+          for part in offset_dict:
+            part_id = int(part)
+            self._offset_dict[part_id] = offset_dict[part]
+        else:
+          assert 'invalid offset_type: %s' % offset_type
+    self._task_offset_dict = {}
+
+  def _preprocess(self, field_dict):
+    output_dict = super(KafkaInput, self)._preprocess(field_dict)
+
+    # append offset fields
+    if Input.DATA_OFFSET in field_dict:
+      output_dict[Input.DATA_OFFSET] = field_dict[Input.DATA_OFFSET]
+
+    # for _get_features to include DATA_OFFSET
+    if Input.DATA_OFFSET not in self._appended_fields:
+      self._appended_fields.append(Input.DATA_OFFSET)
+
+    return output_dict
+
+  def _parse_csv(self, line, message_key, message_offset):
+    record_defaults = [
+        self.get_type_defaults(t, v)
+        for t, v in zip(self._input_field_types, self._input_field_defaults)
+    ]
+
+    fields = tf.decode_csv(
+        line,
+        use_quote_delim=False,
+        field_delim=self._data_config.separator,
+        record_defaults=record_defaults,
+        name='decode_csv')
+
+    inputs = {self._input_fields[x]: fields[x] for x in self._effective_fids}
+
+    for x in self._label_fids:
+      inputs[self._input_fields[x]] = fields[x]
+
+    # record current offset
+    def _parse_offset(message_offset):
+      for kv in message_offset:
+        if six.PY3:
+          kv = kv.decode('utf-8')
+        k, v = kv.split(':')
+        k = int(k)
+        v = int(v)
+        if k not in self._task_offset_dict or v > self._task_offset_dict[k]:
+          self._task_offset_dict[k] = v
+      return json.dumps(self._task_offset_dict)
+
+    inputs[Input.DATA_OFFSET] = tf.py_func(_parse_offset, [message_offset],
+                                           tf.string)
+    return inputs
+
+  def restore(self, checkpoint_path):
+    if checkpoint_path is None:
+      return
+
+    offset_path = checkpoint_path + '.offset'
+    if not gfile.Exists(offset_path):
+      return
+
+    logging.info('will restore kafka offset from  %s' % offset_path)
+    with gfile.GFile(offset_path, 'r') as fin:
+      offset_dict = json.load(fin)
+      self._offset_dict = {}
+      for k in offset_dict:
+        v = offset_dict[k]
+        k = int(k)
+        if k not in self._offset_dict or v > self._offset_dict[k]:
+          self._offset_dict[k] = v
+
+  def _get_topics(self):
+    task_num = self._task_num
+    task_index = self._task_index
+    if self._data_config.chief_redundant and self._mode == tf.estimator.ModeKeys.TRAIN:
+      task_index = max(task_index - 1, 0)
+      task_num = max(task_num - 1, 1)
+
+    topics = []
+    self._task_offset_dict = {}
+    for part_id in range(self._num_partition):
+      if (part_id % task_num) == task_index:
+        offset = self._offset_dict.get(part_id, 0)
+        topics.append('%s:%d:%d' % (self._kafka.topic, part_id, offset))
+        self._task_offset_dict[part_id] = offset
+    logging.info('assigned topic partitions: %s' % (','.join(topics)))
+    assert len(
+        topics) > 0, 'no partitions are assigned for this task(%d/%d)' % (
+            self._task_index, self._task_num)
+    return topics
+
+  def _build(self, mode, params):
+    num_parallel_calls = self._data_config.num_parallel_calls
+    task_topics = self._get_topics()
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      assert self._kafka is not None, 'kafka_train_input is not set.'
+      train_kafka = self._kafka
+      logging.info(
+          'train kafka server: %s topic: %s task_num: %d task_index: %d topics: %s'
+          % (train_kafka.server, train_kafka.topic, self._task_num,
+             self._task_index, task_topics))
+
+      dataset = KafkaDataset(
+          task_topics,
+          servers=train_kafka.server,
+          group=train_kafka.group,
+          eof=False,
+          config_global=list(self._kafka.config_global),
+          config_topic=list(self._kafka.config_topic),
+          message_key=True,
+          message_offset=True)
+
+      if self._data_config.shuffle:
+        dataset = dataset.shuffle(
+            self._data_config.shuffle_buffer_size,
+            seed=2020,
+            reshuffle_each_iteration=True)
+    else:
+      eval_kafka = self._kafka
+      assert self._kafka is not None, 'kafka_eval_input is not set.'
+
+      logging.info(
+          'eval kafka server: %s topic: %s task_num: %d task_index: %d topics: %s'
+          % (eval_kafka.server, eval_kafka.topic, self._task_num,
+             self._task_index, task_topics))
+
+      dataset = KafkaDataset(
+          task_topics,
+          servers=self._kafka.server,
+          group=eval_kafka.group,
+          eof=False,
+          config_global=list(self._kafka.config_global),
+          config_topic=list(self._kafka.config_topic),
+          message_key=True,
+          message_offset=True)
+
+    dataset = dataset.batch(self._data_config.batch_size)
+    dataset = dataset.map(
+        self._parse_csv, num_parallel_calls=num_parallel_calls)
+    if self._data_config.ignore_error:
+      dataset = dataset.apply(ignore_errors)
+    dataset = dataset.prefetch(buffer_size=self._prefetch_size)
+    dataset = dataset.map(
+        map_func=self._preprocess, num_parallel_calls=num_parallel_calls)
+
+    dataset = dataset.prefetch(buffer_size=self._prefetch_size)
+
+    if mode != tf.estimator.ModeKeys.PREDICT:
+      dataset = dataset.map(lambda x:
+                            (self._get_features(x), self._get_labels(x)))
+    else:
+      dataset = dataset.map(lambda x: (self._get_features(x)))
+    return dataset

easy_rec/python/input/load_parquet.py

@@ -0,0 +1,317 @@
+import logging
+import multiprocessing
+import queue
+
+import numpy as np
+import pandas as pd
+
+
+def start_data_proc(task_index,
+                    task_num,
+                    num_proc,
+                    file_que,
+                    data_que,
+                    proc_start_que,
+                    proc_stop_que,
+                    batch_size,
+                    label_fields,
+                    sparse_fea_names,
+                    dense_fea_names,
+                    dense_fea_cfgs,
+                    reserve_fields,
+                    drop_remainder,
+                    need_pack=True):
+  mp_ctxt = multiprocessing.get_context('spawn')
+  proc_arr = []
+  for proc_id in range(num_proc):
+    proc = mp_ctxt.Process(
+        target=load_data_proc,
+        args=(proc_id, file_que, data_que, proc_start_que, proc_stop_que,
+              batch_size, label_fields, sparse_fea_names, dense_fea_names,
+              dense_fea_cfgs, reserve_fields, drop_remainder, task_index,
+              task_num, need_pack),
+        name='task_%d_data_proc_%d' % (task_index, proc_id))
+    proc.daemon = True
+    proc.start()
+    proc_arr.append(proc)
+  return proc_arr
+
+
+def _should_stop(proc_stop_que):
+  try:
+    proc_stop_que.get(block=False)
+    logging.info('data_proc stop signal received')
+    proc_stop_que.close()
+    return True
+  except queue.Empty:
+    return False
+  except ValueError:
+    return True
+  except AssertionError:
+    return True
+
+
+def _add_to_que(data_dict, data_que, proc_stop_que):
+  while True:
+    try:
+      data_que.put(data_dict, timeout=5)
+      return True
+    except queue.Full:
+      logging.warning('data_que is full')
+      if _should_stop(proc_stop_que):
+        return False
+    except ValueError:
+      logging.warning('data_que is closed')
+      return False
+    except AssertionError:
+      logging.warning('data_que is closed')
+      return False
+
+
+def _get_one_file(file_que, proc_stop_que):
+  while True:
+    try:
+      input_file = file_que.get(timeout=1)
+      return input_file
+    except queue.Empty:
+      pass
+  return None
+
+
+def _pack_sparse_feas(data_dict, sparse_fea_names):
+  fea_val_arr = []
+  fea_len_arr = []
+  for fea_name in sparse_fea_names:
+    fea_len_arr.append(data_dict[fea_name][0])
+    fea_val_arr.append(data_dict[fea_name][1])
+    del data_dict[fea_name]
+  fea_lens = np.concatenate(fea_len_arr, axis=0)
+  fea_vals = np.concatenate(fea_val_arr, axis=0)
+  data_dict['sparse_fea'] = (fea_lens, fea_vals)
+
+
+def _pack_dense_feas(data_dict, dense_fea_names, dense_fea_cfgs):
+  fea_val_arr = []
+  for fea_name, fea_cfg in zip(dense_fea_names, dense_fea_cfgs):
+    fea_val_arr.append(data_dict[fea_name].reshape([-1, fea_cfg.raw_input_dim]))
+    del data_dict[fea_name]
+  fea_vals = np.concatenate(fea_val_arr, axis=1)
+  data_dict['dense_fea'] = fea_vals
+
+
+def _reshape_dense_feas(data_dict, dense_fea_names, dense_fea_cfgs):
+  for fea_name, fea_cfg in zip(dense_fea_names, dense_fea_cfgs):
+    data_dict[fea_name] = data_dict[fea_name].reshape(
+        [-1, fea_cfg.raw_input_dim])
+
+
+def _load_dense(input_data, field_names, sid, eid, dense_dict):
+  for k in field_names:
+    if isinstance(input_data[k][0], np.ndarray):
+      np_dtype = type(input_data[k][sid][0])
+      dense_dict[k] = np.array([x[0] for x in input_data[k][sid:eid]],
+                               dtype=np_dtype)
+    else:
+      dense_dict[k] = input_data[k][sid:eid].to_numpy()
+
+
+def _load_and_pad_dense(input_data, field_names, sid, dense_dict,
+                        part_dense_dict, part_dense_dict_n, batch_size):
+  for k in field_names:
+    if isinstance(input_data[k][0], np.ndarray):
+      np_dtype = type(input_data[k][sid][0])
+      tmp_lbls = np.array([x[0] for x in input_data[k][sid:]], dtype=np_dtype)
+    else:
+      tmp_lbls = input_data[k][sid:].to_numpy()
+    if part_dense_dict is not None and k in part_dense_dict:
+      tmp_lbls = np.concatenate([part_dense_dict[k], tmp_lbls], axis=0)
+      if len(tmp_lbls) > batch_size:
+        dense_dict[k] = tmp_lbls[:batch_size]
+        part_dense_dict_n[k] = tmp_lbls[batch_size:]
+      elif len(tmp_lbls) == batch_size:
+        dense_dict[k] = tmp_lbls
+      else:
+        part_dense_dict_n[k] = tmp_lbls
+    else:
+      part_dense_dict_n[k] = tmp_lbls
+
+
+def load_data_proc(proc_id, file_que, data_que, proc_start_que, proc_stop_que,
+                   batch_size, label_fields, sparse_fea_names, dense_fea_names,
+                   dense_fea_cfgs, reserve_fields, drop_remainder, task_index,
+                   task_num, need_pack):
+  logging.info('data proc %d start, proc_start_que=%s' %
+               (proc_id, proc_start_que.qsize()))
+  proc_start_que.get()
+  effective_fields = sparse_fea_names + dense_fea_names
+  all_fields = effective_fields
+  if label_fields is not None:
+    all_fields = all_fields + label_fields
+  if reserve_fields is not None:
+    for tmp in reserve_fields:
+      if tmp not in all_fields:
+        all_fields.append(tmp)
+  logging.info('data proc %d start, file_que.qsize=%d' %
+               (proc_id, file_que.qsize()))
+  num_files = 0
+  part_data_dict = {}
+
+  is_good = True
+  total_batch_cnt = 0
+  total_sample_cnt = 0
+  while is_good:
+    if _should_stop(proc_stop_que):
+      is_good = False
+      break
+    input_file = _get_one_file(file_que, proc_stop_que)
+    if input_file is None:
+      break
+    num_files += 1
+    input_data = pd.read_parquet(input_file, columns=all_fields)
+    data_len = len(input_data[all_fields[0]])
+    total_sample_cnt += data_len
+    batch_num = int(data_len / batch_size)
+    res_num = data_len % batch_size
+
+    sid = 0
+    for batch_id in range(batch_num):
+      eid = sid + batch_size
+      data_dict = {}
+
+      if label_fields is not None and len(label_fields) > 0:
+        _load_dense(input_data, label_fields, sid, eid, data_dict)
+
+      if reserve_fields is not None and len(reserve_fields) > 0:
+        data_dict['reserve'] = {}
+        _load_dense(input_data, reserve_fields, sid, eid, data_dict['reserve'])
+
+      if len(sparse_fea_names) > 0:
+        for k in sparse_fea_names:
+          val = input_data[k][sid:eid]
+          if isinstance(input_data[k][sid], np.ndarray):
+            all_lens = np.array([len(x) for x in val], dtype=np.int32)
+            all_vals = np.concatenate(val.to_numpy())
+          else:
+            all_lens = np.ones([len(val)], dtype=np.int32)
+            all_vals = val.to_numpy()
+          assert np.sum(all_lens) == len(
+              all_vals), 'len(all_vals)=%d np.sum(all_lens)=%d' % (
+                  len(all_vals), np.sum(all_lens))
+          data_dict[k] = (all_lens, all_vals)
+
+      if len(dense_fea_names) > 0:
+        _load_dense(input_data, dense_fea_names, sid, eid, data_dict)
+
+      if need_pack:
+        if len(sparse_fea_names) > 0:
+          _pack_sparse_feas(data_dict, sparse_fea_names)
+        if len(dense_fea_names) > 0:
+          _pack_dense_feas(data_dict, dense_fea_names, dense_fea_cfgs)
+      else:
+        if len(dense_fea_names) > 0:
+          _reshape_dense_feas(data_dict, dense_fea_names, dense_fea_cfgs)
+      # logging.info('task_index=%d sid=%d eid=%d total_len=%d' % (task_index, sid, eid,
+      #      len(data_dict['sparse_fea'][1])))
+      if not _add_to_que(data_dict, data_que, proc_stop_que):
+        logging.info('add to que failed')
+        is_good = False
+        break
+      total_batch_cnt += 1
+      sid += batch_size
+
+    if res_num > 0 and is_good:
+      data_dict = {}
+      part_data_dict_n = {}
+
+      if label_fields is not None and len(label_fields) > 0:
+        _load_and_pad_dense(input_data, label_fields, sid, data_dict,
+                            part_data_dict, part_data_dict_n, batch_size)
+
+      if reserve_fields is not None and len(reserve_fields) > 0:
+        data_dict['reserve'] = {}
+        part_data_dict_n['reserve'] = {}
+        _load_and_pad_dense(input_data, label_fields, sid, data_dict['reserve'],
+                            part_data_dict['reserve'],
+                            part_data_dict_n['reserve'], batch_size)
+
+      if len(dense_fea_names) > 0:
+        _load_and_pad_dense(input_data, dense_fea_names, sid, data_dict,
+                            part_data_dict, part_data_dict_n, batch_size)
+
+      if len(sparse_fea_names) > 0:
+        for k in sparse_fea_names:
+          val = input_data[k][sid:]
+
+          if isinstance(input_data[k][sid], np.ndarray):
+            all_lens = np.array([len(x) for x in val], dtype=np.int32)
+            all_vals = np.concatenate(val.to_numpy())
+          else:
+            all_lens = np.ones([len(val)], dtype=np.int32)
+            all_vals = val.to_numpy()
+
+          if part_data_dict is not None and k in part_data_dict:
+            tmp_lens = np.concatenate([part_data_dict[k][0], all_lens], axis=0)
+            tmp_vals = np.concatenate([part_data_dict[k][1], all_vals], axis=0)
+            if len(tmp_lens) > batch_size:
+              tmp_res_lens = tmp_lens[batch_size:]
+              tmp_lens = tmp_lens[:batch_size]
+              tmp_num_elems = np.sum(tmp_lens)
+              tmp_res_vals = tmp_vals[tmp_num_elems:]
+              tmp_vals = tmp_vals[:tmp_num_elems]
+              part_data_dict_n[k] = (tmp_res_lens, tmp_res_vals)
+              data_dict[k] = (tmp_lens, tmp_vals)
+            elif len(tmp_lens) == batch_size:
+              data_dict[k] = (tmp_lens, tmp_vals)
+            else:
+              part_data_dict_n[k] = (tmp_lens, tmp_vals)
+          else:
+            part_data_dict_n[k] = (all_lens, all_vals)
+
+      if effective_fields[0] in data_dict:
+        if need_pack:
+          if len(sparse_fea_names) > 0:
+            _pack_sparse_feas(data_dict, sparse_fea_names)
+          if len(dense_fea_names) > 0:
+            _pack_dense_feas(data_dict, dense_fea_names, dense_fea_cfgs)
+        else:
+          if len(dense_fea_names) > 0:
+            _reshape_dense_feas(data_dict, dense_fea_names, dense_fea_cfgs)
+        if not _add_to_que(data_dict, data_que, proc_stop_que):
+          logging.info('add to que failed')
+          is_good = False
+          break
+        total_batch_cnt += 1
+      part_data_dict = part_data_dict_n
+  if len(part_data_dict) > 0 and is_good:
+    batch_len = len(part_data_dict[effective_fields[0]][0])
+    if not drop_remainder:
+      if need_pack:
+        if len(sparse_fea_names) > 0:
+          _pack_sparse_feas(part_data_dict, sparse_fea_names)
+        if len(dense_fea_names) > 0:
+          _pack_dense_feas(part_data_dict, dense_fea_names, dense_fea_cfgs)
+      else:
+        if len(dense_fea_names) > 0:
+          _reshape_dense_feas(part_data_dict, dense_fea_names, dense_fea_cfgs)
+      logging.info('remainder batch: %s sample_num=%d' %
+                   (','.join(part_data_dict.keys()), batch_len))
+      _add_to_que(part_data_dict, data_que, proc_stop_que)
+      total_batch_cnt += 1
+    else:
+      logging.warning('drop remain %d samples as drop_remainder is set' %
+                      batch_len)
+  if is_good:
+    is_good = _add_to_que(None, data_que, proc_stop_que)
+  logging.info(
+      'data_proc_id[%d]: is_good = %s, total_batch_cnt=%d, total_sample_cnt=%d'
+      % (proc_id, is_good, total_batch_cnt, total_sample_cnt))
+  data_que.close(wait_send_finish=is_good)
+
+  while not is_good:
+    try:
+      if file_que.get(timeout=1) is None:
+        break
+    except queue.Empty:
+      pass
+  file_que.close()
+  logging.info('data proc %d done, file_num=%d' % (proc_id, num_files))