Merge branch 'hdf_iterator' of https://github.com/jreback/pandas into hdf_iterator

* 'hdf_iterator' of https://github.com/jreback/pandas:
  BUG: fix HDFStore iterator to handle a where properly (GH8014)

Author: Brendan Boerner

pandas/io/pytables.py

@@ -662,28 +662,18 @@ def select(self, key, where=None, start=None, stop=None, columns=None,
         s = self._create_storer(group)
         s.infer_axes()
 
-        def func(_start, _stop):
+        # function to call on iteration
+        def func(_start, _stop, _where):
             return s.read(start=_start, stop=_stop,
-                          where=where,
+                          where=_where,
                           columns=columns, **kwargs)
 
-        if iterator or chunksize is not None:
-            if not s.is_table:
-                raise TypeError(
-                    "can only use an iterator or chunksize on a table")
+        # create the iterator
+        it = TableIterator(self, s, func, where=where, nrows=s.nrows, start=start,
+                           stop=stop, iterator=iterator, chunksize=chunksize,
+                           auto_close=auto_close)
 
-            # read the coordinates & iterate
-            if where is not None:
-                c = s.read_coordinates(where=where, **kwargs)
-                def func(_start, _stop):
-                    return s.read(where=c[_start:_stop], columns=columns, **kwargs)
-
-            return TableIterator(self, func, nrows=s.nrows, start=start,
-                                 stop=stop, chunksize=chunksize,
-                                 auto_close=auto_close)
-
-        return TableIterator(self, func, nrows=s.nrows, start=start, stop=stop,
-                             auto_close=auto_close).get_values()
+        return it.get_result()
 
     def select_as_coordinates(
             self, key, where=None, start=None, stop=None, **kwargs):
@@ -786,34 +776,22 @@ def select_as_multiple(self, keys, where=None, selector=None, columns=None,
         # axis is the concentation axes
         axis = list(set([t.non_index_axes[0][0] for t in tbls]))[0]
 
-        # for a not-none where, select the coordinates and chunk on those
-        if where is not None:
-            c = s.read_coordinates(where=where, **kwargs)
+        def func(_start, _stop, _where):
 
-            def func(_start, _stop):
-                objs = [t.read(where=c[_start:_stop], columns=columns, **kwargs) for t in tbls]
+            # retrieve the objs, _where is always passed as a set of coordinates here
+            objs = [t.read(where=_where, columns=columns, **kwargs) for t in tbls]
 
-                # concat and return
-                return concat(objs, axis=axis,
-                              verify_integrity=False).consolidate()
+            # concat and return
+            return concat(objs, axis=axis,
+                          verify_integrity=False).consolidate()
 
-        else:
-
-            def func(_start, _stop):
-                objs = [t.read(start=_start, stop=_stop,
-                               columns=columns, **kwargs) for t in tbls]
+        # create the iterator
+        it = TableIterator(self, s, func, where=where, nrows=nrows, start=start,
+                           stop=stop, iterator=iterator, chunksize=chunksize,
+                           auto_close=auto_close)
 
-                # concat and return
-                return concat(objs, axis=axis,
-                              verify_integrity=False).consolidate()
-
-        if iterator or chunksize is not None:
-            return TableIterator(self, func, nrows=nrows, start=start,
-                                 stop=stop, chunksize=chunksize,
-                                 auto_close=auto_close)
+        return it.get_result(coordinates=True)
 
-        return TableIterator(self, func, nrows=nrows, start=start, stop=stop,
-                             auto_close=auto_close).get_values()
 
     def put(self, key, value, format=None, append=False, **kwargs):
         """
@@ -1308,42 +1286,54 @@ class TableIterator(object):
         ----------
 
         store : the reference store
-        func  : the function to get results
+        s     : the refered storer
+        func  : the function to execute the query
+        where : the where of the query
         nrows : the rows to iterate on
         start : the passed start value (default is None)
         stop  : the passed stop value (default is None)
-        chunksize : the passed chunking valeu (default is 50000)
+        iterator : boolean, whether to use the default iterator
+        chunksize : the passed chunking value (default is 50000)
         auto_close : boolean, automatically close the store at the end of
             iteration, default is False
         kwargs : the passed kwargs
         """
 
-    def __init__(self, store, func, nrows, start=None, stop=None,
-                 chunksize=None, auto_close=False):
+    def __init__(self, store, s, func, where, nrows, start=None, stop=None,
+                 iterator=False, chunksize=None, auto_close=False):
         self.store = store
-        self.func = func
+        self.s     = s
+        self.func  = func
+        self.where = where
         self.nrows = nrows or 0
         self.start = start or 0
 
         if stop is None:
             stop = self.nrows
         self.stop = min(self.nrows, stop)
 
-        if chunksize is None:
-            chunksize = 100000
+        self.coordinates = None
+        if iterator or chunksize is not None:
+            if chunksize is None:
+                chunksize = 100000
+            self.chunksize = int(chunksize)
+        else:
+            self.chunksize = None
 
-        self.chunksize = int(chunksize)
         self.auto_close = auto_close
 
     def __iter__(self):
+
+        # iterate
         current = self.start
         while current < self.stop:
 
             stop = min(current + self.chunksize, self.stop)
-            value = self.func(current, stop)
+            value = self.func(None, None, self.coordinates[current:stop])
             if value is None:
                 continue
-            current = current + min(self.chunksize,len(value))
+
+            current = current + self.chunksize
 
             yield value
 
@@ -1353,12 +1343,29 @@ def close(self):
         if self.auto_close:
             self.store.close()
 
-    def get_values(self):
-        results = self.func(self.start, self.stop)
+    def get_result(self, coordinates=False):
+
+        #  return the actual iterator
+        if self.chunksize is not None:
+            if not self.s.is_table:
+                raise TypeError(
+                    "can only use an iterator or chunksize on a table")
+
+            self.coordinates = self.s.read_coordinates(where=self.where)
+
+            return self
+
+        # if specified read via coordinates (necessary for multiple selections
+        if coordinates:
+            where = self.s.read_coordinates(where=self.where)
+        else:
+            where = self.where
+
+        # directly return the result
+        results = self.func(self.start, self.stop, where)
         self.close()
         return results
 
-
 class IndexCol(StringMixin):
 
     """ an index column description class