Merge pull request #785 from wireservice/performance

Avoid exponential cleanup algorithm, closes #575

Author: James McKinney

csvkit/cleanup.py

@@ -20,23 +20,6 @@ def join_rows(rows, joiner=' '):
     return fixed_row
 
 
-def extract_joinable_row_errors(errs):
-    joinable = []
-
-    for err in reversed(errs):
-        if type(err) is not LengthMismatchError:
-            break
-
-        if joinable and err.line_number != joinable[-1].line_number - 1:
-            break
-
-        joinable.append(err)
-
-    joinable.reverse()
-
-    return joinable
-
-
 class RowChecker(object):
     """
     Iterate over rows of a CSV producing cleaned rows and storing error rows.
@@ -45,7 +28,6 @@ class RowChecker(object):
     def __init__(self, reader):
         self.reader = reader
         self.column_names = next(reader)
-
         self.errors = []
         self.rows_joined = 0
         self.joins = 0
@@ -54,39 +36,52 @@ def checked_rows(self):
         """
         A generator which yields rows which are ready to write to output.
         """
+        length = len(self.column_names)
         line_number = self.reader.line_num
+        joinable_row_errors = []
 
         for row in self.reader:
             try:
-                if len(row) != len(self.column_names):
-                    raise LengthMismatchError(line_number, row, len(self.column_names))
+                if len(row) != length:
+                    raise LengthMismatchError(line_number, row, length)
 
                 yield row
+
+                # Don't join rows across valid rows.
+                joinable_row_errors = []
             except LengthMismatchError as e:
                 self.errors.append(e)
 
-                joinable_row_errors = extract_joinable_row_errors(self.errors)
+                # Don't join with long rows.
+                if len(row) > length:
+                    joinable_row_errors = []
+                else:
+                    joinable_row_errors.append(e)
 
-                while joinable_row_errors:
-                    fixed_row = join_rows([err.row for err in joinable_row_errors], joiner=' ')
+                    while joinable_row_errors:
+                        fixed_row = join_rows([error.row for error in joinable_row_errors], joiner=' ')
 
-                    if len(fixed_row) < len(self.column_names):
-                        break
+                        if len(fixed_row) < length:
+                            break
 
-                    if len(fixed_row) == len(self.column_names):
-                        self.rows_joined += len(joinable_row_errors)
-                        self.joins += 1
+                        if len(fixed_row) == length:
+                            self.rows_joined += len(joinable_row_errors)
+                            self.joins += 1
 
-                        yield fixed_row
+                            yield fixed_row
 
-                        for fixed in joinable_row_errors:
-                            self.errors.remove(fixed)
+                            for fixed in joinable_row_errors:
+                                joinable_row_errors.remove(fixed)
+                                self.errors.remove(fixed)
 
-                        break
+                            break
 
-                    joinable_row_errors = joinable_row_errors[1:]  # keep trying in case we're too long because of a straggler
+                        joinable_row_errors = joinable_row_errors[1:]  # keep trying in case we're too long because of a straggler
 
             except CSVTestException as e:
                 self.errors.append(e)
 
+                # Don't join rows across other errors.
+                joinable_row_errors = []
+
             line_number = self.reader.line_num

tests/test_cleanup.py

@@ -5,7 +5,7 @@
 except ImportError:
     import unittest
 
-from csvkit.cleanup import extract_joinable_row_errors, join_rows
+from csvkit.cleanup import join_rows
 from csvkit.exceptions import CSVTestException, LengthMismatchError
 
 
@@ -25,45 +25,6 @@ def test_fix_rows(self):
         self.assertEqual(" ".join([start[0][-1], start[1][0], start[2][0], start[3][0]]), fixed[2])
         self.assertEqual(start[3][1], fixed[3])
 
-    def test_extract_joinable_row_errors(self):
-        e1 = LengthMismatchError(1, ['foo', 'bar', 'baz'], 10)
-        e2 = LengthMismatchError(2, ['foo', 'bar', 'baz'], 10)
-        e3 = LengthMismatchError(3, ['foo', 'bar', 'baz'], 10)
-        errs = [e1, e2, e3]
-        joinable = extract_joinable_row_errors(errs)
-        self.assertEqual(3, len(joinable))
-        for e, j in zip(errs, joinable):
-            self.assertTrue(e is j)
-
-    def test_extract_joinable_row_errors_2(self):
-        e1 = LengthMismatchError(1, ['foo', 'bar', 'baz'], 10)
-        e2 = CSVTestException(2, ['foo', 'bar', 'baz'], "A throwaway message.")
-        e3 = LengthMismatchError(3, ['foo', 'bar', 'baz'], 10)
-        errs = [e1, e2, e3]
-        joinable = extract_joinable_row_errors(errs)
-        self.assertEqual(1, len(joinable))
-        self.assertTrue(next(iter(joinable)) is e3)
-
-    def test_extract_joinable_row_errors_3(self):
-        e1 = CSVTestException(1, ['foo', 'bar', 'baz'], "A throwaway message.")
-        e2 = LengthMismatchError(2, ['foo', 'bar', 'baz'], 10)
-        e3 = LengthMismatchError(3, ['foo', 'bar', 'baz'], 10)
-        errs = [e1, e2, e3]
-        joinable = extract_joinable_row_errors(errs)
-        self.assertEqual(2, len(joinable))
-        joinable = list(joinable)
-        self.assertTrue(joinable[0] is e2)
-        self.assertTrue(joinable[1] is e3)
-
-    def test_extract_joinable_row_errors_4(self):
-        e1 = CSVTestException(1, ['foo', 'bar', 'baz'], "A throwaway message.")
-        e2 = LengthMismatchError(2, ['foo', 'bar', 'baz'], 10)
-        e3 = LengthMismatchError(4, ['foo', 'bar', 'baz'], 10)
-        errs = [e1, e2, e3]
-        joinable = extract_joinable_row_errors(errs)
-        self.assertEqual(1, len(joinable))
-        self.assertTrue(next(iter(joinable)) is e3)
-
     def test_real_world_join_fail(self):
         start = [['168772', '1102', '$0.23 TO $0.72', 'HOUR', '1.5%'],
                  ['GROSS', '1.5% '],