Skip to content
Snippets Groups Projects
Unverified Commit 89bdaa2f authored by adam-james's avatar adam-james Committed by GitHub
Browse files

XLSX Pivot Table Downloads have 'Native' Pivot Table (#43791) 

* XLSX Pivot Table Downloads have 'Native' Pivot Table

Addresses: #2473

The xlsx downloads work fine for regular tables. However, on Pivot Table downloads, the resulting file isn't pivoted
and is formatted somewhat strangely (it contains additional rows related to totals, has an extra 'pivot-grouping'
column).

Now, the downloaded xlsx file contains 2 sheets:
- the data sheet containing the question's rows unpivoted. If you were constructing a pivot table yourself, you would
probably start with this shape of data.
- the pivot sheet containing a 'native' Pivot table whose reference points to the data on the data sheet, and whose
columns, rows, and values align with the cols, rows, and measures set up in Metabase

* Silly typo!

* Exported Pivot uses correct aggregation (sum, count, stddev, avg, min, max)

* data for pivot exports now also uses the formatting

* Add a test asserting that a pivot table exists in the xlsx file

* add-row! is a method dispatching on sheet class

* remove unnecessary hint

* address review feedback:

 - consolidate the aggregation function key fns into one function in the .xlsx namespace
   this was moved out of postprocessing so that the xlsx specific fns don't litter the post processing namespace at
   this time

 - clean up an unnecessary destructiring inside the native-pivot function
 - don't shadow 'name' inside the body of a fn

* Add tests that confirm zero-col and zero-row situations work

* Consolidate the atoms used to store pivot rows/options
parent 301ab32c
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
src/metabase/query_processor/pivot/postprocess.clj
+ 6
− 3
View file @ 89bdaa2f
......@@ -31,6 +31,8 @@
[:column-titles [:sequential [:string]]]
[:pivot-rows [:sequential [:int {:min 0}]]]
[:pivot-cols [:sequential [:int {:min 0}]]]
[:pivot-grouping-key {:optional true}
[:int {:min 0}]]
[:pivot-measures {:optional true}
[:sequential [:int {:min 0}]]]])
......@@ -169,7 +171,7 @@
:else
export-style-row)))
(defn- pivot-grouping-key
(defn pivot-grouping-key
"Get the index into the raw pivot rows for the 'pivot-grouping' column."
[column-titles]
;; a vector is kinda sorta a map of indices->values, so
......@@ -177,7 +179,7 @@
(get (set/map-invert (vec column-titles)) "pivot-grouping"))
(mu/defn ^:private pivot-measures
"Get the indices into the raw pivot rows corresponding to the pivot table's measure."
"Get the indices into the raw pivot rows corresponding to the pivot table's measure(s)."
[{:keys [pivot-rows pivot-cols column-titles]} :- ::pivot-spec]
(-> (set/difference
;; every possible idx is just the range over the count of cols
......@@ -188,7 +190,8 @@
sort
vec))
(mu/defn ^:private add-pivot-measures :- ::pivot-spec
(mu/defn add-pivot-measures :- ::pivot-spec
"Given a pivot-spec map without the `:pivot-measures` key, determine what key(s) the measures will be and assoc that value into `:pivot-measures`."
[pivot-spec :- ::pivot-spec]
(assoc pivot-spec :pivot-measures (pivot-measures pivot-spec)))
......
src/metabase/query_processor/streaming/xlsx.clj
+ 182
− 32
View file @ 89bdaa2f
......@@ -4,9 +4,11 @@
[clojure.string :as str]
[dk.ative.docjure.spreadsheet :as spreadsheet]
[java-time.api :as t]
[medley.core :as m]
[metabase.formatter :as formatter]
[metabase.lib.schema.temporal-bucketing
:as lib.schema.temporal-bucketing]
[metabase.query-processor.pivot.postprocess :as qp.pivot.postprocess]
[metabase.query-processor.streaming.common :as common]
[metabase.query-processor.streaming.interface :as qp.si]
[metabase.shared.models.visualization-settings :as mb.viz]
......@@ -17,8 +19,10 @@
(:import
(java.io OutputStream)
(java.time LocalDate LocalDateTime LocalTime OffsetDateTime OffsetTime ZonedDateTime)
(org.apache.poi.ss.usermodel Cell DataFormat DateUtil Workbook)
(org.apache.poi.ss.util CellRangeAddress)
(org.apache.poi.ss SpreadsheetVersion)
(org.apache.poi.ss.usermodel Cell DataFormat DateUtil Workbook DataConsolidateFunction)
(org.apache.poi.xssf.usermodel XSSFWorkbook XSSFSheet XSSFRow XSSFPivotTable)
(org.apache.poi.ss.util CellReference CellRangeAddress AreaReference)
(org.apache.poi.xssf.streaming SXSSFRow SXSSFSheet SXSSFWorkbook)))
(set! *warn-on-reflection* true)
......@@ -391,11 +395,38 @@
(catch Exception _ value
value))))
(defn- add-row!
(defmulti ^:private add-row!
"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."
{:arglists '([sheet values cols col-settings cell-styles typed-cell-styles])}
(fn [sheet _values _cols _col-settings _cell-styles _typed-cell-styles]
(class sheet)))
(defmethod add-row! org.apache.poi.xssf.streaming.SXSSFSheet
[^SXSSFSheet sheet values cols col-settings cell-styles typed-cell-styles]
(let [row-num (if (= 0 (.getPhysicalNumberOfRows sheet))
0
(inc (.getLastRowNum sheet)))
row (.createRow sheet row-num)]
(doseq [[value col styles index] (map vector values cols cell-styles (range (count values)))]
(let [id-or-name (or (:id col) (:name col))
settings (or (get col-settings {::mb.viz/field-id id-or-name})
(get col-settings {::mb.viz/column-name id-or-name}))
scaled-val (if (and value (::mb.viz/scale settings))
(* value (::mb.viz/scale settings))
value)
;; Temporal values are converted into strings in the format-rows QP middleware, which is enabled during
;; dashboard subscription/pulse generation. If so, we should parse them here so that formatting is applied.
parsed-value (or
(maybe-parse-temporal-value value col)
(maybe-parse-coordinate-value value col)
scaled-val)]
(set-cell! (.createCell ^SXSSFRow row ^Integer index) parsed-value styles typed-cell-styles)))
row))
(defmethod add-row! org.apache.poi.xssf.usermodel.XSSFSheet
[^XSSFSheet sheet values cols col-settings cell-styles typed-cell-styles]
(let [row-num (if (= 0 (.getPhysicalNumberOfRows sheet))
0
(inc (.getLastRowNum sheet)))
......@@ -413,7 +444,7 @@
(maybe-parse-temporal-value value col)
(maybe-parse-coordinate-value value col)
scaled-val)]
(set-cell! (.createCell ^SXSSFRow row ^Integer index) parsed-value styles typed-cell-styles)))
(set-cell! (.createCell ^XSSFRow row ^Integer index) parsed-value styles typed-cell-styles)))
row))
(def ^:dynamic *auto-sizing-threshold*
......@@ -447,38 +478,157 @@
(.setAutoFilter ^SXSSFSheet sheet (new CellRangeAddress 0 0 0 (dec col-count)))
(.createFreezePane ^SXSSFSheet sheet 0 1)))
(defn- cell-range
[rows]
(let [x (dec (count (first rows)))
y (dec (count rows))]
(CellRangeAddress.
0 ;; first row
y ;; last row
0 ;; first col
x ;; last col
)))
(defn- cell-range->area-ref
[cell-range]
(AreaReference. (.formatAsString ^CellRangeAddress cell-range) SpreadsheetVersion/EXCEL2007))
;; Possible Functions: https://poi.apache.org/apidocs/dev/org/apache/poi/ss/usermodel/DataConsolidateFunction.html
;; I'm only including the keys that seem to work for our Pivot Tables as of 2024-06-06
(defn- col->aggregation-fn
[{agg-name :name source :source}]
(when (= :aggregation source)
(let [agg-name (u/lower-case-en agg-name)]
(cond
(str/starts-with? agg-name "sum") DataConsolidateFunction/SUM
(str/starts-with? agg-name "avg") DataConsolidateFunction/AVERAGE
(str/starts-with? agg-name "min") DataConsolidateFunction/MIN
(str/starts-with? agg-name "max") DataConsolidateFunction/MAX
(str/starts-with? agg-name "count") DataConsolidateFunction/COUNT
(str/starts-with? agg-name "stddev") DataConsolidateFunction/STD_DEV))))
(defn pivot-opts->pivot-spec
"Utility that adds :pivot-grouping-key to the pivot-opts map internal to the xlsx streaming response writer."
[pivot-opts cols]
(let [titles (mapv :display_name cols)
agg-fns (mapv col->aggregation-fn cols)]
(-> pivot-opts
(assoc :column-titles titles)
qp.pivot.postprocess/add-pivot-measures
(assoc :aggregation-functions agg-fns)
(assoc :pivot-grouping-key (qp.pivot.postprocess/pivot-grouping-key titles)))))
(defn- native-pivot
[rows
{:keys [pivot-grouping-key] :as pivot-spec}
{:keys [ordered-cols col-settings viz-settings]}]
(let [idx-shift (fn [indices]
(map (fn [idx]
(if (> idx pivot-grouping-key)
(dec idx)
idx)) indices))
ordered-cols (vec (m/remove-nth pivot-grouping-key ordered-cols))
pivot-rows (idx-shift (:pivot-rows pivot-spec))
pivot-cols (idx-shift (:pivot-cols pivot-spec))
pivot-measures (idx-shift (:pivot-measures pivot-spec))
aggregation-functions (vec (m/remove-nth pivot-grouping-key (:aggregation-functions pivot-spec)))
wb (spreadsheet/create-workbook
"pivot" [[]]
"data" [])
data-format (. ^XSSFWorkbook wb createDataFormat)
cell-styles (compute-column-cell-styles wb data-format viz-settings ordered-cols)
typed-cell-styles (compute-typed-cell-styles wb data-format)
data-sheet (spreadsheet/select-sheet "data" wb)
pivot-sheet (spreadsheet/select-sheet "pivot" wb)
area-ref (cell-range->area-ref (cell-range rows))
_ (doseq [row rows]
(add-row! data-sheet row ordered-cols col-settings cell-styles typed-cell-styles))
^XSSFPivotTable pivot-table (.createPivotTable ^XSSFSheet pivot-sheet
^AreaReference area-ref
(CellReference. "A1")
^XSSFSheet data-sheet)]
(doseq [idx pivot-rows]
(.addRowLabel pivot-table idx))
(doseq [idx pivot-cols]
(.addColLabel pivot-table idx))
(doseq [idx pivot-measures]
(.addColumnLabel pivot-table (get aggregation-functions idx DataConsolidateFunction/COUNT) idx))
wb))
(defmethod qp.si/streaming-results-writer :xlsx
[_ ^OutputStream os]
(let [workbook (SXSSFWorkbook.)
sheet (spreadsheet/add-sheet! workbook (tru "Query result"))
data-format (. workbook createDataFormat)
cell-styles (volatile! nil)
typed-cell-styles (volatile! nil)]
(let [workbook (SXSSFWorkbook.)
sheet (spreadsheet/add-sheet! workbook (tru "Query result"))
data-format (. workbook createDataFormat)
cell-styles (volatile! nil)
typed-cell-styles (volatile! nil)
pivot-data! (atom {:rows []})]
(reify qp.si/StreamingResultsWriter
(begin! [_ {{:keys [ordered-cols]} :data} {col-settings ::mb.viz/column-settings :as viz-settings}]
(vreset! cell-styles (compute-column-cell-styles workbook data-format viz-settings ordered-cols))
(vreset! typed-cell-styles (compute-typed-cell-styles workbook data-format))
(doseq [i (range (count ordered-cols))]
(.trackColumnForAutoSizing ^SXSSFSheet sheet i))
(setup-header-row! sheet (count ordered-cols))
(spreadsheet/add-row! sheet (common/column-titles ordered-cols col-settings true)))
(begin! [_ {{:keys [ordered-cols format-rows? pivot-export-options]} :data}
{col-settings ::mb.viz/column-settings :as viz-settings}]
(let [opts (when pivot-export-options
(pivot-opts->pivot-spec (merge {:pivot-cols []
:pivot-rows []}
pivot-export-options) ordered-cols))
;; col-names are created later when exporting a pivot table, so only create them if there are no pivot options
col-names (when-not opts (common/column-titles ordered-cols (::mb.viz/column-settings viz-settings) format-rows?))]
(vreset! cell-styles (compute-column-cell-styles workbook data-format viz-settings ordered-cols))
(vreset! typed-cell-styles (compute-typed-cell-styles workbook data-format))
;; when pivot options exist, we want to save them to access later when processing the complete set of results for export.
(when opts
(swap! pivot-data! assoc
:cell-style-data {:ordered-cols ordered-cols
:col-settings col-settings
:viz-settings viz-settings}
:pivot-options opts))
(when col-names
(doseq [i (range (count ordered-cols))]
(.trackColumnForAutoSizing ^SXSSFSheet sheet i))
(setup-header-row! sheet (count ordered-cols))
(spreadsheet/add-row! sheet (common/column-titles ordered-cols col-settings true)))))
(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)))
row)
col-settings (::mb.viz/column-settings viz-settings)]
(add-row! sheet ordered-row ordered-cols col-settings @cell-styles @typed-cell-styles)
(when (= (inc row-num) *auto-sizing-threshold*)
(autosize-columns! sheet))))
(let [ordered-row (if output-order
(let [row-v (into [] row)]
(for [i output-order] (row-v i)))
row)
col-settings (::mb.viz/column-settings viz-settings)
{:keys [pivot-options]} @pivot-data!]
(if pivot-options
(let [{:keys [pivot-grouping-key]} pivot-options
group (get row pivot-grouping-key)]
(when (= 0 group)
;; TODO: right now, the way I'm building up the native pivot,
;; I end up using the docjure set-cell! (since I create a whole sheet with all the rows at once)
;; I'll want to change that so I can use the set-cell! method we have in this ns, but for now just string everything.
(let [modified-row (->> (vec (m/remove-nth pivot-grouping-key row))
(mapv (fn [value]
(if (number? value)
value
(str value)))))]
(swap! pivot-data! update :rows conj modified-row))))
(do
(add-row! sheet ordered-row ordered-cols col-settings @cell-styles @typed-cell-styles)
(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))
(try
(spreadsheet/save-workbook-into-stream! os workbook)
(finally
(.dispose workbook)
(.close os)))))))
(let [{:keys [pivot-options rows cell-style-data]} @pivot-data!]
(if pivot-options
(let [header (vec (m/remove-nth (:pivot-grouping-key pivot-options) (:column-titles pivot-options)))
wb (native-pivot (concat [header] rows) pivot-options cell-style-data)]
(try
(spreadsheet/save-workbook-into-stream! os wb)
(finally
(.dispose workbook)
(.close os))))
(do
(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))
(try
(spreadsheet/save-workbook-into-stream! os workbook)
(finally
(.dispose workbook)
(.close os))))))))))
test/metabase/api/downloads_exports_test.clj
+ 101
− 2
View file @ 89bdaa2f
......@@ -11,9 +11,15 @@
- Alert attachments"
(:require
[clojure.data.csv :as csv]
[clojure.java.io :as io]
[clojure.set :as set]
[clojure.test :refer :all]
[metabase.test :as mt]))
[dk.ative.docjure.spreadsheet :as spreadsheet]
[metabase.test :as mt])
(:import
(org.apache.poi.xssf.usermodel XSSFSheet)))
(set! *warn-on-reflection* true)
(def ^:private pivot-rows-query
"SELECT *
......@@ -137,7 +143,7 @@
(is (= "Grand Totals"
(first (last result)))))))))))
(deftest ^:parallel multi-measure-pivot-tables-headers-test
(deftest multi-measure-pivot-tables-headers-test
(testing "Pivot tables with multiple measures correctly include the measure titles in the final header row."
(mt/dataset test-data
(mt/with-temp [:model/Card {pivot-card-id :id}
......@@ -171,6 +177,7 @@
"Average of Rating"]]
(take 2 result))))))))
(deftest ^:parallel zero-column-pivot-tables-test
(testing "Pivot tables with zero columns download correctly."
(mt/dataset test-data
......@@ -262,3 +269,95 @@
["Totals for 2019-01-01T00:00:00Z" "" "1060.98" "1060.98"]
["Grand Totals" "" "11149.28" "11149.28"]]
result)))))))
(deftest pivot-table-native-pivot-in-xlsx-test
(testing "Pivot table xlsx downloads produce a 'native pivot' in the workbook."
(mt/dataset test-data
(mt/with-temp [:model/Card {pivot-card-id :id}
{:display :pivot
:visualization_settings {:pivot_table.column_split
{:rows [[:field (mt/id :products :created_at) {:base-type :type/DateTime, :temporal-unit :month}]],
:columns [[:field (mt/id :products :category) {:base-type :type/Text}]],
:values [[:aggregation 0]
[:aggregation 1]]}}
:dataset_query {:database (mt/id)
:type :query
:query
{:source-table (mt/id :products)
:aggregation [[:sum [:field (mt/id :products :price) {:base-type :type/Float}]]
[:avg [:field (mt/id :products :rating) {:base-type :type/Float}]]]
:breakout [[:field (mt/id :products :category) {:base-type :type/Text}]
[:field (mt/id :products :created_at) {:base-type :type/DateTime :temporal-unit :month}]]}}}]
(let [result (mt/user-http-request :crowberto :post 200 (format "card/%d/query/xlsx?format_rows=false" pivot-card-id))
pivot (with-open [in (io/input-stream result)]
(->> (spreadsheet/load-workbook in)
(spreadsheet/select-sheet "pivot")
((fn [s] (.getPivotTables ^XSSFSheet s)))))]
(is (not (nil? pivot))))))))
(deftest ^:parallel zero-column-native-pivot-tables-test
(testing "Pivot tables with zero columns download correctly as xlsx."
(mt/dataset test-data
(mt/with-temp [:model/Card {pivot-card-id :id}
{:display :pivot
:visualization_settings {:pivot_table.column_split
{:rows [[:field (mt/id :products :created_at) {:base-type :type/DateTime :temporal-unit :month}]
[:field (mt/id :products :category) {:base-type :type/Text}]]
:columns []
:values [[:aggregation 0]]}}
:dataset_query {:database (mt/id)
:type :query
:query
{:source-table (mt/id :products)
:aggregation [[:sum [:field (mt/id :products :price) {:base-type :type/Float}]]]
:breakout [[:field (mt/id :products :category) {:base-type :type/Text}]
[:field (mt/id :products :created_at) {:base-type :type/DateTime :temporal-unit :month}]]}}}]
(let [result (mt/user-http-request :crowberto :post 200 (format "card/%d/query/xlsx?format_rows=false" pivot-card-id))
[pivot data] (with-open [in (io/input-stream result)]
(let [wb (spreadsheet/load-workbook in)
pivot (.getPivotTables ^XSSFSheet (spreadsheet/select-sheet "pivot" wb))
data (->> (spreadsheet/select-sheet "data" wb)
spreadsheet/row-seq
(mapv (fn [row] (->> (spreadsheet/cell-seq row)
(mapv spreadsheet/read-cell)))))]
[pivot data]))]
(is (not (nil? pivot)))
(is (= [["Category" "Created At" "Sum of Price"]
["Doohickey" #inst "2016-05-01T00:00:00.000-00:00" 144.12]
["Doohickey" #inst "2016-06-01T00:00:00.000-00:00" 82.92]
["Doohickey" #inst "2016-07-01T00:00:00.000-00:00" 78.22]
["Doohickey" #inst "2016-08-01T00:00:00.000-00:00" 71.09]
["Doohickey" #inst "2016-09-01T00:00:00.000-00:00" 45.65]]
(take 6 data))))))))
(deftest ^:parallel zero-row-native-pivot-tables-test
(testing "Pivot tables with zero rows download correctly as xlsx."
(mt/dataset test-data
(mt/with-temp [:model/Card {pivot-card-id :id}
{:display :pivot
:visualization_settings {:pivot_table.column_split
{:rows []
:columns [[:field (mt/id :products :category) {:base-type :type/Text}]]
:values [[:aggregation 0]]}}
:dataset_query {:database (mt/id)
:type :query
:query
{:source-table (mt/id :products)
:aggregation [[:sum [:field (mt/id :products :price) {:base-type :type/Float}]]]
:breakout [[:field (mt/id :products :category) {:base-type :type/Text}]]}}}]
(let [result (mt/user-http-request :crowberto :post 200 (format "card/%d/query/xlsx?format_rows=false" pivot-card-id))
[pivot data] (with-open [in (io/input-stream result)]
(let [wb (spreadsheet/load-workbook in)
pivot (.getPivotTables ^XSSFSheet (spreadsheet/select-sheet "pivot" wb))
data (->> (spreadsheet/select-sheet "data" wb)
spreadsheet/row-seq
(mapv (fn [row] (->> (spreadsheet/cell-seq row)
(mapv spreadsheet/read-cell)))))]
[pivot data]))]
(is (not (nil? pivot)))
(is (= [["Category" "Sum of Price"]
["Doohickey" 2185.89]
["Gadget" 3019.2]
["Gizmo" 2834.88]
["Widget" 3109.31]]
(take 6 data))))))))
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment