浏览代码

HUE-3442 [core] Improve memory efficiency of XLS download

Also make CSV/XLS download size configurable (by max num of cells)
Jenny Kim 9 年之前
父节点
当前提交
5a813a40ad

+ 6 - 4
apps/beeswax/src/beeswax/conf.py

@@ -96,11 +96,13 @@ LIST_PARTITIONS_LIMIT = Config(
   type=int,
   help=_t('Limit the number of partitions that can be listed. A positive value will be set as the LIMIT.'))
 
-DOWNLOAD_ROW_LIMIT = Config(
-  key='download_row_limit',
-  default=1000000,
+DOWNLOAD_CELL_LIMIT = Config(
+  key='download_cell_limit',
+  default=10000000,
   type=int,
-  help=_t('A limit to the number of rows that can be downloaded from a query. A value of -1 means there will be no limit. A maximum of 30,000 is applied to XLS downloads.'))
+  help=_t('A limit to the number of cells (rows * columns) that can be downloaded from a query '
+          '(e.g. - 10K rows * 1K columns = 10M cells.) '
+          'A value of -1 means there will be no limit.'))
 
 APPLY_NATURAL_SORT_MAX = Config(
   key="apply_natural_sort_max",

+ 20 - 11
apps/beeswax/src/beeswax/data_export.py

@@ -41,7 +41,9 @@ def download(handle, format, db):
     LOG.error('Unknown download format "%s"' % (format,))
     return
 
-  content_generator = HS2DataAdapter(handle, db, conf.DOWNLOAD_ROW_LIMIT.get())
+  max_cells = conf.DOWNLOAD_CELL_LIMIT.get()
+
+  content_generator = HS2DataAdapter(handle, db, max_cells=max_cells, start_over=True)
   generator = export_csvxls.create_generator(content_generator, format)
   return export_csvxls.make_response(generator, format, 'query_result')
 
@@ -57,17 +59,16 @@ def upload(path, handle, user, db, fs):
   else:
     fs.do_as_user(user.username, fs.create, path)
 
-  content_generator = HS2DataAdapter(handle, db, -1, start_over=True)
+  content_generator = HS2DataAdapter(handle, db, max_cells=-1, start_over=True)
   for header, data in content_generator:
     dataset = export_csvxls.dataset(None, data)
     fs.do_as_user(user.username, fs.append, path, dataset.csv)
 
 
-def HS2DataAdapter(handle, db, max_rows=0, start_over=True):
+def HS2DataAdapter(handle, db, max_cells=-1, start_over=True):
   """
   HS2DataAdapter(query_model, db) -> headers, 2D array of data.
   """
-
   results = db.fetch(handle, start_over=start_over, rows=FETCH_SIZE)
 
   while not results.ready:
@@ -75,25 +76,33 @@ def HS2DataAdapter(handle, db, max_rows=0, start_over=True):
     results = db.fetch(handle, start_over=start_over, rows=FETCH_SIZE)
 
   headers = results.cols()
+  num_cols = len(headers)
+
+  # For result sets with high num of columns, fetch in smaller batches to avoid serialization cost
+  if num_cols > 100:
+    LOG.warn('The query results contain %d columns and may take an extremely long time to download, will reduce fetch size to 100.' % num_cols)
+    fetch_size = 100
+  else:
+    fetch_size = FETCH_SIZE
 
-  num_rows_seen = 0
-  limit_rows = max_rows > -1
+  row_ctr = 1
+  limit_cells = max_cells > -1
 
   while results is not None:
     data = []
     for row in results.rows():
-      num_rows_seen += 1
-      if limit_rows and num_rows_seen > max_rows:
-        LOG.warn('The query results exceeded the maximum row limit of %d. Data has been truncated.' % max_rows)
+      row_ctr += 1
+      if limit_cells and (row_ctr * num_cols) > max_cells:
+        LOG.warn('The query results exceeded the maximum cell limit of %d. Data has been truncated to first %d rows.' % (max_cells, row_ctr))
         break
       data.append(row)
 
     yield headers, data
 
-    if limit_rows and num_rows_seen > max_rows:
+    if limit_cells and (row_ctr * num_cols) > max_cells:
       break
 
     if results.has_more:
-      results = db.fetch(handle, start_over=False, rows=FETCH_SIZE)
+      results = db.fetch(handle, start_over=False, rows=fetch_size)
     else:
       results = None

+ 14 - 0
apps/beeswax/src/beeswax/tests.py

@@ -809,6 +809,7 @@ for x in sys.stdin:
     resp = download(handle, 'xls', self.db)
 
     sheet_data = _read_xls_sheet_data(resp)
+    num_cols = len(sheet_data[0])
     # It should have 257 lines (256 + header)
     assert_equal(len(sheet_data), 257, sheet_data)
 
@@ -822,6 +823,19 @@ for x in sys.stdin:
 
     assert_equal(sheet_data, csv_data)
 
+    # Test max cell limit truncation
+    finish = conf.DOWNLOAD_CELL_LIMIT.set_for_testing(num_cols*5)
+    try:
+      hql = 'SELECT * FROM `%(db)s`.`test`' % {'db': self.db_name}
+      query = hql_query(hql)
+      handle = self.db.execute_and_wait(query)
+      resp = download(handle, 'xls', self.db)
+      sheet_data = _read_xls_sheet_data(resp)
+      # It should have 5 lines
+      assert_equal(len(sheet_data), 5, sheet_data)
+    finally:
+      finish()
+
 
   def test_data_upload(self):
     hql = 'SELECT * FROM `%(db)s`.`test`' % {'db': self.db_name}

+ 3 - 3
desktop/conf.dist/hue.ini

@@ -803,10 +803,10 @@
   # The maximum number of partitions that will be included in the SELECT * LIMIT sample query for partitioned tables.
   ## query_partitions_limit=10
 
-  # A limit to the number of rows that can be downloaded from a query.
+  # A limit to the number of cells (rows * columns) that can be downloaded from a query
+  # (e.g. - 10K rows * 1K columns = 10M cells.)
   # A value of -1 means there will be no limit.
-  # A maximum of 30,000 is applied to XLS downloads.
-  ## download_row_limit=1000000
+  ## download_cell_limit=10000000
 
   # Hue will try to close the Hive query when the user leaves the editor page.
   # This will free all the query resources in HiveServer2, but also make its results inaccessible.

+ 3 - 3
desktop/conf/pseudo-distributed.ini.tmpl

@@ -805,10 +805,10 @@
   # The maximum number of partitions that will be included in the SELECT * LIMIT sample query for partitioned tables.
   ## query_partitions_limit=10
 
-  # A limit to the number of rows that can be downloaded from a query.
+  # A limit to the number of cells (rows * columns) that can be downloaded from a query
+  # (e.g. - 10K rows * 1K columns = 10M cells.)
   # A value of -1 means there will be no limit.
-  # A maximum of 30,000 is applied to XLS downloads.
-  ## download_row_limit=1000000
+  ## download_cell_limit=10000000
 
   # Hue will try to close the Hive query when the user leaves the editor page.
   # This will free all the query resources in HiveServer2, but also make its results inaccessible.

+ 20 - 40
desktop/core/src/desktop/lib/export_csvxls.py

@@ -29,16 +29,16 @@ from desktop.lib import i18n
 
 
 LOG = logging.getLogger(__name__)
-MAX_XLS_ROWS = 65000
-MAX_XLS_COLS = 5000
 
 
 def nullify(cell):
   return cell if cell is not None else "NULL"
 
-def format(row, encoding=None):
+
+def encode_row(row, encoding=None):
   return [smart_str(nullify(cell), encoding or i18n.get_site_encoding(), strings_only=True, errors='replace') for cell in row]
 
+
 def dataset(headers, data, encoding=None):
   """
   dataset(headers, data) -> Dataset object
@@ -48,10 +48,10 @@ def dataset(headers, data, encoding=None):
   dataset = tablib.Dataset()
 
   if headers:
-    dataset.headers = format(headers, encoding)
+    dataset.headers = encode_row(headers, encoding)
 
   for row in data:
-    dataset.append(format(row, encoding))
+    dataset.append(encode_row(row, encoding))
 
   return dataset
 
@@ -61,22 +61,10 @@ class XlsWrapper():
     self.xls = xls
 
 
-def xls_dataset(headers, data, encoding=None):
+def xls_dataset(workbook):
   output = StringIO.StringIO()
-
-  workbook = openpyxl.Workbook(write_only=True)
-  worksheet = workbook.create_sheet()
-
-  if headers:
-    worksheet.append(format(headers, encoding))
-
-  for row in data:
-    worksheet.append(format(row, encoding))
-
   workbook.save(output)
-
   output.seek(0)
-
   return XlsWrapper(output.read())
 
 
@@ -86,34 +74,26 @@ def create_generator(content_generator, format, encoding=None):
     for headers, data in content_generator:
       yield dataset(show_headers and headers or None, data, encoding).csv
       show_headers = False
-
   elif format == 'xls':
-    headers = None
-    data = []
-    for _headers, _data in content_generator:
-      # Forced limit on size from tablib
-      if len(data) > MAX_XLS_ROWS:
-        LOG.warn('The query results exceeded the maximum row limit of %d. Data has been truncated.' % MAX_XLS_ROWS)
-        break
+    workbook = openpyxl.Workbook(write_only=True)
+    worksheet = workbook.create_sheet()
+    row_ctr = 0
 
-      if _headers and len(_headers) > MAX_XLS_COLS:
-        LOG.warn('The query results exceeded the maximum column limit of %d. Data has been truncated.' % MAX_XLS_COLS)
-        _headers = _headers[:MAX_XLS_COLS]
-
-      headers = _headers
+    for _headers, _data in content_generator:
+      # Write headers to workbook once
+      if _headers and row_ctr == 0:
+        worksheet.append(encode_row(_headers, encoding))
+        row_ctr += 1
 
+      # Write row data to workbook
       for row in _data:
-        if len(row) >= MAX_XLS_COLS:
-          row = row[:MAX_XLS_COLS]
+        worksheet.append(encode_row(row, encoding))
+        row_ctr += 1
 
-        data.append(row)
-
-    if len(data) > MAX_XLS_ROWS:
-      LOG.warn('The query results exceeded the maximum row limit of %d. Data has been truncated.' % MAX_XLS_ROWS)
-      data = data[:MAX_XLS_ROWS]
-
-    yield xls_dataset(headers, data, encoding).xls
+    yield xls_dataset(workbook).xls
     gc.collect()
+  else:
+    raise Exception("Unknown format: %s" % format)
 
 
 def make_response(generator, format, name, encoding=None):

+ 3 - 36
desktop/core/src/desktop/lib/test_export_csvxls.py

@@ -15,18 +15,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import openpyxl
 import StringIO
 
 from nose.tools import assert_equal
 from openpyxl import load_workbook
 
-from desktop.lib.export_csvxls import MAX_XLS_ROWS, MAX_XLS_COLS, create_generator, make_response
+from desktop.lib.export_csvxls import create_generator, make_response
 
 
 def content_generator(header, data):
   yield header, data
 
+
 def test_export_csv():
   headers = ["x", "y"]
   data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None] ]
@@ -39,6 +39,7 @@ def test_export_csv():
   assert_equal('x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\n', content)
   assert_equal("attachment; filename=foo.csv", response["content-disposition"])
 
+
 def test_export_xls():
   headers = ["x", "y"]
   data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None] ]
@@ -55,40 +56,6 @@ def test_export_xls():
   assert_equal(expected_data, sheet_data)
   assert_equal("attachment; filename=foo.xlsx", response["content-disposition"])
 
-def test_export_xls_truncate_rows():
-  headers = ["a"]
-  data = [["1"]] * (MAX_XLS_ROWS + 1)
-  sheet = [headers] + data
-
-  # Check XLS
-  generator = create_generator(content_generator(headers, data), "xls")
-  response = make_response(generator, "xls", "foo")
-  assert_equal("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", response["content-type"])
-
-  sheet_data = _read_xls_sheet_data(response)
-  assert_equal(len(sheet[:MAX_XLS_ROWS + 1]), len(sheet_data))
-  assert_equal("attachment; filename=foo.xlsx", response["content-disposition"])
-
-def test_export_xls_truncate_cols():
-  headers = [u"a"] * (MAX_XLS_COLS + 1)
-  data = [[u"1"] * (MAX_XLS_COLS + 1)]
-  sheet = [headers] + data
-
-  # Check XLS
-  generator = create_generator(content_generator(headers, data), "xls")
-  response = make_response(generator, "xls", "foo")
-  assert_equal("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", response["content-type"])
-  expected_data = [[cell is not None and cell or "NULL" for cell in row][:MAX_XLS_COLS] for row in sheet]
-
-  sheet_data = _read_xls_sheet_data(response)
-  sheet_data[0] = sheet_data[0][:MAX_XLS_COLS]
-  sheet_data[1] = sheet_data[1][:MAX_XLS_COLS]
-  assert_equal(len(expected_data), len(sheet_data))
-  assert_equal(len(expected_data[0]), len(sheet_data[0]))
-
-  assert_equal(expected_data, sheet_data)
-  assert_equal("attachment; filename=foo.xlsx", response["content-disposition"])
-
 
 def _read_xls_sheet_data(response):
   content = ''.join(response.streaming_content)