# Buggy code
def _batching(func):
    @wraps(func)
    def arg_wrapper(*args, **kwargs):
        # priority: decorator > class_attribute
        # by default data is in args[1:] (self needs to be taken into account)
        b_size = batch_size or getattr(args[0], "batch_size", None)
        # no batching if b_size is None
        if b_size is None:
            return func(*args, **kwargs)

        args = list(args)
        default_logger.debug(
            f"batching enabled for {func.__qualname__} batch_size={b_size} "
            f"num_batch={num_batch} axis={split_over_axis}"
        )

        # assume all datas have the same length
        full_data_size = _get_size(args[slice_on], split_over_axis)
        total_size = _get_total_size(full_data_size, b_size, num_batch)
        final_result = []
        data_iterators = [
            batch_iterator(args[slice_on + i][:total_size], b_size, split_over_axis)
            for i in range(0, num_data)
        ]

        for batch in data_iterators[0]:
            args[slice_on] = batch
            for idx in range(1, num_data):
                args[slice_on + idx] = next(data_iterators[idx])

            r = func(*args, **kwargs)

            if r is not None:
                final_result.append(r)

        return _merge_results_after_batching(final_result, merge_over_axis)

    return arg_wrapper


# Traceback (TypeError)
Traceback (most recent call last):
File "/home/cristian/envs/lyrics/lib/python3.7/site-packages/jina/peapods/pea.py", line 292, in msg_callback
self.zmqlet.send_message(self._callback(msg))
File "/home/cristian/envs/lyrics/lib/python3.7/site-packages/jina/peapods/pea.py", line 266, in _callback
self.pre_hook(msg).handle(msg).post_hook(msg)
File "/home/cristian/envs/lyrics/lib/python3.7/site-packages/jina/peapods/pea.py", line 160, in handle
self.executor(self.request_type)
File "/home/cristian/envs/lyrics/lib/python3.7/site-packages/jina/executors/__init__.py", line 570, in __call__
d()
File "/home/cristian/envs/lyrics/lib/python3.7/site-packages/jina/drivers/__init__.py", line 251, in __call__
self._traverse_apply(self.req.docs, *args, **kwargs)
File "/home/cristian/envs/lyrics/lib/python3.7/site-packages/jina/drivers/__init__.py", line 268, in _traverse_apply
**kwargs,
File "/home/cristian/envs/lyrics/lib/python3.7/site-packages/jina/drivers/__init__.py", line 281, in _traverse_rec
doc.chunks, doc, 'chunks', path[1:], *args, **kwargs
File "/home/cristian/envs/lyrics/lib/python3.7/site-packages/jina/drivers/__init__.py", line 284, in _traverse_rec
self._apply_all(docs, parent_doc, parent_edge_type, *args, **kwargs)
File "/home/cristian/envs/lyrics/lib/python3.7/site-packages/jina/drivers/search.py", line 124, in _apply_all
idx, dist = self.exec_fn(embed_vecs, top_k=int(self.top_k))
File "/home/cristian/envs/lyrics/lib/python3.7/site-packages/jina/executors/indexers/vector.py", line 279, in query
dist = self._euclidean(_keys, self.query_handler)
File "/home/cristian/envs/lyrics/lib/python3.7/site-packages/jina/executors/decorators.py", line 241, in arg_wrapper
for b in batch_iterator(data[:total_size], b_size, split_over_axis, yield_slice=yield_slice):
TypeError: 'NoneType' object is not subscriptable

# Fixed code
def _batching(func):
    @wraps(func)
    def arg_wrapper(*args, **kwargs):
        data = args[slice_on]
        # priority: decorator > class_attribute
        # by default data is in args[1:] (self needs to be taken into account)
        b_size = batch_size or getattr(args[0], "batch_size", None)
        # no batching if b_size is None
        if b_size is None or data is None:
            return func(*args, **kwargs)

        args = list(args)
        default_logger.debug(
            f"batching enabled for {func.__qualname__} batch_size={b_size} "
            f"num_batch={num_batch} axis={split_over_axis}"
        )

        # assume all datas have the same length
        full_data_size = _get_size(args[slice_on], split_over_axis)
        total_size = _get_total_size(full_data_size, b_size, num_batch)
        final_result = []
        data_iterators = [
            batch_iterator(args[slice_on + i][:total_size], b_size, split_over_axis)
            for i in range(0, num_data)
        ]

        for batch in data_iterators[0]:
            args[slice_on] = batch
            for idx in range(1, num_data):
                args[slice_on + idx] = next(data_iterators[idx])

            r = func(*args, **kwargs)

            if r is not None:
                final_result.append(r)

        return _merge_results_after_batching(final_result, merge_over_axis)

    return arg_wrapper