BUG: fix HDFStore iterator to handle a where properly (GH8014)

Author: jreback

doc/source/v0.15.0.txt

@@ -454,7 +454,7 @@ Bug Fixes
 - Bug in pickle deserialization that failed for pre-0.14.1 containers with dup items trying to avoid ambiguity
   when matching block and manager items, when there's only one block there's no ambiguity (:issue:`7794`)
 
-
+- Bug in HDFStore iteration when passing a where (:issue:`8014`)
 
 - Bug in repeated timeseries line and area plot may result in ``ValueError`` or incorrect kind (:issue:`7733`)
 

pandas/io/pytables.py

@@ -662,15 +662,22 @@ def select(self, key, where=None, start=None, stop=None, columns=None,
         s = self._create_storer(group)
         s.infer_axes()
 
-        # what we are actually going to do for a chunk
         def func(_start, _stop):
-            return s.read(where=where, start=_start, stop=_stop,
+            return s.read(start=_start, stop=_stop,
+                          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")
+
+            # 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)
@@ -779,18 +786,26 @@ 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]
 
-        def func(_start, _stop):
-            if where is not None:
-                c = s.read_coordinates(where=where, start=_start, stop=_stop, **kwargs)
-            else:
-                c = None
+        # for a not-none where, select the coordinates and chunk on those
+        if where is not None:
+            c = s.read_coordinates(where=where, **kwargs)
 
-            objs = [t.read(where=c, start=_start, stop=_stop,
-                           columns=columns, **kwargs) for t in tbls]
+            def func(_start, _stop):
+                objs = [t.read(where=c[_start:_stop], 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]
+
+                # 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,
@@ -1317,20 +1332,20 @@ def __init__(self, store, func, nrows, start=None, stop=None,
         if chunksize is None:
             chunksize = 100000
 
-        self.chunksize = chunksize
+        self.chunksize = int(chunksize)
         self.auto_close = auto_close
 
     def __iter__(self):
         current = self.start
         while current < self.stop:
-            stop = current + self.chunksize
-            v = self.func(current, stop)
-            current = stop
 
-            if v is None:
+            stop = min(current + self.chunksize, self.stop)
+            value = self.func(current, stop)
+            if value is None:
                 continue
+            current = current + min(self.chunksize,len(value))
 
-            yield v
+            yield value
 
         self.close()
 

pandas/io/tests/test_pytables.py

@@ -3376,6 +3376,72 @@ def test_select_iterator(self):
             #result = concat(results)
             #tm.assert_frame_equal(expected, result)
 
+    def test_select_iterator_8014(self):
+
+        # single table
+        with ensure_clean_store(self.path) as store:
+
+            chunksize=1e4
+            expected = tm.makeTimeDataFrame(100064, 'S')
+            _maybe_remove(store, 'df')
+            store.append('df',expected)
+
+            beg_dt = expected.index[0]
+            end_dt = expected.index[-1]
+
+            #
+            # w/o iterator
+            #
+
+            # select w/o iteration and no where clause works
+            result = store.select('df')
+            tm.assert_frame_equal(expected, result)
+
+            # select w/o iterator and where clause, single term, begin
+            # of range, works
+            where = "index >= '%s'" % beg_dt
+            result = store.select('df',where=where)
+            tm.assert_frame_equal(expected, result)
+
+            # select w/o iterator and where clause, single term, end
+            # of range, works
+            where = "index <= '%s'" % end_dt
+            result = store.select('df',where=where)
+            tm.assert_frame_equal(expected, result)
+
+            # select w/o iterator and where clause, inclusive range,
+            # works
+            where = "index >= '%s' & index <= '%s'" % (beg_dt, end_dt)
+            result = store.select('df',where=where)
+            tm.assert_frame_equal(expected, result)
+
+            #
+            # with iterator
+            #
+
+            # select w/iterator and no where clause works
+            results = [ s for s in store.select('df',chunksize=chunksize) ]
+            result = concat(results)
+            tm.assert_frame_equal(expected, result)
+
+            # select w/iterator and where clause, single term, begin of range
+            where = "index >= '%s'" % beg_dt
+            results = [ s for s in store.select('df',where=where,chunksize=chunksize) ]
+            result = concat(results)
+            tm.assert_frame_equal(expected, result)
+
+            # select w/iterator and where clause, single term, end of range
+            where = "index <= '%s'" % end_dt
+            results = [ s for s in store.select('df',where=where,chunksize=chunksize) ]
+            result = concat(results)
+            tm.assert_frame_equal(expected, result)
+
+            # select w/iterator and where clause, inclusive range
+            where = "index >= '%s' & index <= '%s'" % (beg_dt, end_dt)
+            results = [ s for s in store.select('df',where=where,chunksize=chunksize) ]
+            result = concat(results)
+            tm.assert_frame_equal(expected, result)
+
     def test_retain_index_attributes(self):
 
         # GH 3499, losing frequency info on index recreation