XLSX auto-sized columns, frozen header and auto-filters (#17325)

src/metabase/query_processor/streaming/xlsx.clj

@@ -13,8 +13,9 @@
             [metabase.util.i18n :refer [tru]])
   (:import java.io.OutputStream
            [java.time LocalDate LocalDateTime LocalTime OffsetDateTime OffsetTime ZonedDateTime]
-           [org.apache.poi.ss.usermodel Cell DataFormat DateUtil Sheet Workbook]
-           org.apache.poi.xssf.streaming.SXSSFWorkbook))
+           [org.apache.poi.ss.usermodel Cell DataFormat DateUtil Workbook]
+           org.apache.poi.ss.util.CellRangeAddress
+           [org.apache.poi.xssf.streaming SXSSFRow SXSSFSheet SXSSFWorkbook]))
 
 ;;; +----------------------------------------------------------------------------------------------------------------+
 ;;; |                                         Format string generation                                               |
@@ -366,7 +367,7 @@
   "Adds a row of values to the spreadsheet. Values with the `scaled` viz setting are scaled prior to being added.
 
   This is based on the equivalent function in Docjure, but adapted to support Metabase viz settings."
-  [^Sheet sheet values cols col-settings]
+  [^SXSSFSheet sheet values cols col-settings]
   (let [row-num (if (= 0 (.getPhysicalNumberOfRows sheet))
                   0
                   (inc (.getLastRowNum sheet)))
@@ -378,7 +379,7 @@
             scaled-val (if (and value (::mb.viz/scale settings))
                          (* value (::mb.viz/scale settings))
                          value)]
-        (set-cell! (.createCell row index) scaled-val id-or-name)))
+        (set-cell! (.createCell ^SXSSFRow row ^Integer index) scaled-val id-or-name)))
     row))
 
 (defn- column-titles
@@ -400,15 +401,44 @@
         (str column-title " (" (currency-identifier merged-settings) ")")
         column-title))))
 
+(def ^:dynamic *auto-sizing-threshold*
+  "The maximum number of rows we should use for auto-sizing. If this number is too large, exports
+  of large datasets will be prohibitively slow."
+  100)
+
+(def ^:private extra-column-width
+  "The extra width applied to columns after they have been auto-sized, in units of 1/256 of a character width.
+  This ensures the cells in the header row have enough room for the filter dropdown icon."
+  (* 4 256))
+
+(defn- autosize-columns!
+  "Adjusts each column to fit its largest value, plus a constant amount of extra padding."
+  [sheet]
+  (doseq [i (.getTrackedColumnsForAutoSizing ^SXSSFSheet sheet)]
+    (.autoSizeColumn ^SXSSFSheet sheet i)
+    (.setColumnWidth ^SXSSFSheet sheet i (+ (.getColumnWidth ^SXSSFSheet sheet i) extra-column-width))
+    (.untrackColumnForAutoSizing ^SXSSFSheet sheet i)))
+
+(defn- setup-header-row!
+  "Turns on auto-filter for the header row, which adds a button to each header cell that allows columns to be
+  filtered and sorted. Also freezes the header row so that it floats above the data."
+  [sheet col-count]
+  (when (> col-count 0)
+    (.setAutoFilter ^SXSSFSheet sheet (new CellRangeAddress 0 0 0 (dec col-count)))
+    (.createFreezePane ^SXSSFSheet sheet 0 1)))
+
 (defmethod i/streaming-results-writer :xlsx
   [_ ^OutputStream os]
   (let [workbook            (SXSSFWorkbook.)
         sheet               (spreadsheet/add-sheet! workbook (tru "Query result"))]
     (reify i/StreamingResultsWriter
       (begin! [_ {{:keys [ordered-cols]} :data} {col-settings ::mb.viz/column-settings}]
+        (doseq [i (range (count ordered-cols))]
+          (.trackColumnForAutoSizing ^SXSSFSheet sheet i))
+        (setup-header-row! sheet (count ordered-cols))
         (spreadsheet/add-row! sheet (column-titles ordered-cols col-settings)))
 
-      (write-row! [_ row _ ordered-cols {:keys [output-order] :as viz-settings}]
+      (write-row! [_ row row-num ordered-cols {:keys [output-order] :as viz-settings}]
         (let [ordered-row  (if output-order
                              (let [row-v (into [] row)]
                                (for [i output-order] (row-v i)))
@@ -416,9 +446,14 @@
               col-settings (::mb.viz/column-settings viz-settings)
               cell-styles  (cell-style-delays workbook ordered-cols col-settings)]
           (binding [*cell-styles* cell-styles]
-            (add-row! sheet ordered-row ordered-cols col-settings))))
-
-      (finish! [_ _]
+            (add-row! sheet ordered-row ordered-cols col-settings))
+          (when (= (inc row-num) *auto-sizing-threshold*)
+            (autosize-columns! sheet))))
+
+      (finish! [_ {:keys [row_count]}]
+        (when (or (nil? row_count) (< row_count *auto-sizing-threshold*))
+          ;; Auto-size columns if we never hit the row threshold, or a final row count was not provided
+          (autosize-columns! sheet))
         (spreadsheet/save-workbook-into-stream! os workbook)
         (.dispose workbook)
         (.close os)))))

test/metabase/query_processor/streaming/xlsx_test.clj

@@ -224,9 +224,14 @@
 ;;; |                                               XLSX export tests                                                |
 ;;; +----------------------------------------------------------------------------------------------------------------+
 
+(defn- parse-cell-content
+  [sheet]
+  (for [row (spreadsheet/into-seq sheet)]
+    (map spreadsheet/read-cell row)))
+
 (defn- xlsx-export
   ([ordered-cols viz-settings rows]
-   (xlsx-export ordered-cols viz-settings rows spreadsheet/read-cell))
+   (xlsx-export ordered-cols viz-settings rows parse-cell-content))
 
   ([ordered-cols viz-settings rows parse-fn]
    (with-open [bos (ByteArrayOutputStream.)
@@ -236,19 +241,17 @@
        (doall (map-indexed
                (fn [i row] (i/write-row! results-writer row i ordered-cols viz-settings))
                rows))
-       (i/finish! results-writer {}))
+       (i/finish! results-writer {:row_count (count rows)}))
      (let [bytea (.toByteArray bos)]
        (with-open [is (BufferedInputStream. (ByteArrayInputStream. bytea))]
-        (let [workbook (spreadsheet/load-workbook-from-stream is)
-              sheet    (spreadsheet/select-sheet "Query result" workbook)]
-          (for [row (spreadsheet/into-seq sheet)]
-            (map parse-fn row))))))))
+         (let [workbook (spreadsheet/load-workbook-from-stream is)
+               sheet    (spreadsheet/select-sheet "Query result" workbook)]
+           (parse-fn sheet)))))))
 
 (defn- parse-format-strings
-  [row]
-  (-> row
-      .getCellStyle
-      .getDataFormatString))
+  [sheet]
+  (for [row (spreadsheet/into-seq sheet)]
+    (map #(-> % .getCellStyle .getDataFormatString) row)))
 
 (deftest export-format-test
   (testing "Different format strings are used for ints and numbers that round to ints (with 2 decimal places)"
@@ -414,3 +417,31 @@
            (second (xlsx-export [{:name "val1"} {:name "val2"}]
                                 {}
                                 [[(SampleNastyClass. "Hello XLSX World!") (AnotherNastyClass. "No Encoder")]]))))))
+
+(defn- parse-column-width
+  [sheet]
+  (for [row (spreadsheet/into-seq sheet)]
+    (for [i (range (.getLastCellNum row))]
+      (.getColumnWidth sheet i))))
+
+(deftest auto-sizing-test
+  (testing "Columns in export are autosized to fit their content"
+    (let [[col1-width col2-width] (second (xlsx-export [{:id 0, :name "Col1"} {:id 1, :name "Col2"}]
+                                                       {}
+                                                       [["a" "abcdefghijklmnopqrstuvwxyz"]]
+                                                       parse-column-width))]
+      ;; Provide a marign for error since width measurements end up being slightly different on CI
+      (is (<= 2300 col1-width 2400))
+      (is (<= 7950 col2-width 8200))))
+  (testing "Auto-sizing works when the number of rows is at or above the auto-sizing threshold"
+    (binding [xlsx/*auto-sizing-threshold* 2]
+      (let [[col-width] (second (xlsx-export [{:id 0, :name "Col1"}]
+                                             {}
+                                             [["abcdef"] ["abcedf"]]
+                                             parse-column-width))]
+        (is (<= 2800 col-width 2900)))
+      (let [[col-width] (second (xlsx-export [{:id 0, :name "Col1"}]
+                                             {}
+                                             [["abcdef"] ["abcedf"] ["abcdef"]]
+                                             parse-column-width))]
+        (is (<= 2800 col-width 2900))))))