Browse Source

HUE-8384 [search] Improve importer format guessing performance

jdesjean 7 years ago
parent
commit
1b1d8479ba
1 changed files with 75 additions and 3 deletions
  1. 75 3
      desktop/libs/indexer/src/indexer/file_format.py

+ 75 - 3
desktop/libs/indexer/src/indexer/file_format.py

@@ -33,7 +33,7 @@ from indexer.indexers.morphline_operations import get_operator
 LOG = logging.getLogger(__name__)
 
 
-IMPORT_PEEK_SIZE = 1024 * 1024 * 5
+IMPORT_PEEK_SIZE = 1024 * 1024
 IMPORT_PEEK_NLINES = 20
 
 
@@ -337,9 +337,77 @@ class CSVFormat(FileFormat):
   def _guess_dialect(cls, sample):
     sniffer = csv.Sniffer()
     dialect = sniffer.sniff(sample)
-    has_header = sniffer.has_header(sample)
+    has_header = cls._hasHeader(sniffer, sample, dialect)
     return dialect, has_header
 
+  # Copied from python2.7/csv.py with small modification to 1st line
+  # Results in large performance gain from not having to reprocess the file if dialect is known.
+  @classmethod
+  def _hasHeader(self, sniffer, sample, dialect):
+    # ******Changed from********
+    # rdr = reader(StringIO(sample), self.sniff(sample))
+    from _csv import reader
+    rdr = reader(StringIO.StringIO(sample), dialect)
+
+    header = rdr.next()  # assume first row is header
+
+    columns = len(header)
+    columnTypes = {}
+    for i in range(columns): columnTypes[i] = None
+
+    checked = 0
+    for row in rdr:
+      # arbitrary number of rows to check, to keep it sane
+      if checked > 20:
+        break
+      checked += 1
+
+      if len(row) != columns:
+        continue  # skip rows that have irregular number of columns
+
+      for col in columnTypes.keys():
+
+        for thisType in [int, long, float, complex]:
+          try:
+            thisType(row[col])
+            break
+          except (ValueError, OverflowError):
+            pass
+        else:
+          # fallback to length of string
+          thisType = len(row[col])
+
+        # treat longs as ints
+        if thisType == long:
+          thisType = int
+
+        if thisType != columnTypes[col]:
+          if columnTypes[col] is None:  # add new column type
+            columnTypes[col] = thisType
+          else:
+            # type is inconsistent, remove column from
+            # consideration
+            del columnTypes[col]
+
+    # finally, compare results against first row and "vote"
+    # on whether it's a header
+    hasHeader = 0
+    for col, colType in columnTypes.items():
+      if type(colType) == type(0):  # it's a length
+        if len(header[col]) != colType:
+          hasHeader += 1
+        else:
+          hasHeader -= 1
+      else:  # attempt typecast
+        try:
+          colType(header[col])
+        except (ValueError, TypeError):
+          hasHeader += 1
+        else:
+          hasHeader -= 1
+
+    return hasHeader > 0
+
   @classmethod
   def valid_format(cls, format_):
     valid = super(CSVFormat, cls).valid_format(format_)
@@ -366,7 +434,11 @@ class CSVFormat(FileFormat):
   def _guess_from_file_stream(cls, file_stream):
     for sample_data, sample_lines in cls._get_sample(file_stream):
       try:
-        dialect, has_header = cls._guess_dialect(sample_data)
+        lines = itertools.islice(StringIO.StringIO(sample_data), IMPORT_PEEK_NLINES)
+        sample_data_lines = ''
+        for line in lines:
+          sample_data_lines += line
+        dialect, has_header = cls._guess_dialect(sample_data_lines) # Only use first few lines for guessing. Greatly improves performance of CSV library.
         delimiter = dialect.delimiter
         line_terminator = dialect.lineterminator
         quote_char = dialect.quotechar