From e005f7bb76a79e167d7c31f07444b6dba8c3443f Mon Sep 17 00:00:00 2001
From: adam-james <21064735+adam-james-v@users.noreply.github.com>
Date: Mon, 4 Nov 2024 12:23:08 -0800
Subject: [PATCH] Pivot Measures Order Used in Pivot Exports (#49367)

* Pivot Measures Order Used in Pivot Exports

Fixes #48442

A pivot table can have any number of measures, and the user can order these by dragging in the UI. Before this PR,
that order was ignored and measures would alway be in index order, which is confusing for any user who needs the
measures to be displayed in a particular order, especially if they've re-ordered them in the pivot viz settings UI.

A test has been added to check that measure order is used.

A few minor changes to the pivot qp and post-processor
 - measure indices are looked up in the pivot qp and added in viz-settings order
 - the pivot measures are only added if the qp has not already added them.
 - pivot-opts Malli spec has been made in the namespace, adjusted to allow `nil` as a valid pivot-opts output, and
 used in relevant functions

* address review points.

* add a rows order test
---
 src/metabase/query_processor/pivot.clj        | 97 ++++++++++---------
 .../query_processor/pivot/postprocess.clj     | 17 +++-
 .../query_processor/streaming/csv.clj         |  9 +-
 .../query_processor/streaming/json.clj        |  2 -
 test/metabase/api/downloads_exports_test.clj  | 92 ++++++++++++++++--
 test/metabase/query_processor/pivot_test.clj  |  6 +-
 6 files changed, 154 insertions(+), 69 deletions(-)

diff --git a/src/metabase/query_processor/pivot.clj b/src/metabase/query_processor/pivot.clj
index ecbca94d081..86c736a1968 100644
--- a/src/metabase/query_processor/pivot.clj
+++ b/src/metabase/query_processor/pivot.clj
@@ -45,8 +45,15 @@
 (mr/def ::num-breakouts ::lib.schema.common/int-greater-than-or-equal-to-zero)
 (mr/def ::index         ::lib.schema.common/int-greater-than-or-equal-to-zero)
 
-(mr/def ::pivot-rows [:sequential ::index])
-(mr/def ::pivot-cols [:sequential ::index])
+(mr/def ::pivot-rows     [:sequential ::index])
+(mr/def ::pivot-cols     [:sequential ::index])
+(mr/def ::pivot-measures [:sequential ::index])
+
+(mr/def ::pivot-opts [:maybe
+                      [:map
+                       [:pivot-rows     {:optional true} [:maybe ::pivot-rows]]
+                       [:pivot-cols     {:optional true} [:maybe ::pivot-cols]]
+                       [:pivot-measures {:optional true} [:maybe ::pivot-measures]]]])
 
 (mu/defn- group-bitmask :- ::bitmask
   "Come up with a display name given a combination of breakout `indexes` e.g.
@@ -173,9 +180,7 @@
 (mu/defn- generate-queries :- [:sequential ::lib.schema/query]
   "Generate the additional queries to perform a generic pivot table"
   [query                                               :- ::lib.schema/query
-   {:keys [pivot-rows pivot-cols], :as _pivot-options} :- [:map
-                                                           [:pivot-rows {:optional true} [:maybe ::pivot-rows]]
-                                                           [:pivot-cols {:optional true} [:maybe ::pivot-cols]]]]
+   {:keys [pivot-rows pivot-cols], :as _pivot-options} :- ::pivot-opts]
   (try
     (let [all-breakouts (lib/breakouts query)
           all-queries   (for [breakout-indexes (u/prog1 (breakout-combinations (count all-breakouts) pivot-rows pivot-cols)
@@ -342,48 +347,45 @@
               qp.pipeline/*reduce*  (or reduce qp.pipeline/*reduce*)]
       (qp/process-query first-query rff))))
 
-(mu/defn- pivot-options :- [:map
-                            [:pivot-rows [:maybe [:sequential [:int {:min 0}]]]]
-                            [:pivot-cols [:maybe [:sequential [:int {:min 0}]]]]]
+(mu/defn- pivot-options :- ::pivot-opts
   "Given a pivot table query and a card ID, looks at the `pivot_table.column_split` key in the card's visualization
   settings and generates pivot-rows and pivot-cols to use for generating subqueries."
   [query        :- [:map
                     [:database ::lib.schema.id/database]]
    viz-settings :- [:maybe :map]]
-  (let [column-split         (:pivot_table.column_split viz-settings)
-        column-split-rows    (seq (:rows column-split))
-        column-split-columns (seq (:columns column-split))
-        index-in-breakouts   (when (or column-split-rows
-                                       column-split-columns)
-                               (let [metadata-provider (or (:lib/metadata query)
-                                                           (lib.metadata.jvm/application-database-metadata-provider (:database query)))
-                                     mlv2-query        (lib/query metadata-provider query)
-                                     breakouts         (into []
-                                                             (map-indexed (fn [i col]
-                                                                            (cond-> col
-                                                                              true                         (assoc ::i i)
-                                                                              ;; if the col has a card-id, we swap the :lib/source to say source/card
-                                                                              ;; this allows `lib/find-matching-column` to properly match a column that has a join-alias
-                                                                              ;; but whose source is a model
-                                                                              (contains? col :lib/card-id) (assoc :lib/source :source/card))))
-                                                             (lib/breakouts-metadata mlv2-query))]
-                                 (fn [legacy-ref]
-                                   (try
-                                     (::i (lib.equality/find-column-for-legacy-ref
-                                           mlv2-query
-                                           -1
-                                           legacy-ref
-                                           breakouts))
-                                     (catch Throwable e
-                                       (log/errorf e "Error finding matching column for ref %s" (pr-str legacy-ref))
-                                       nil)))))
-
-        pivot-rows (when column-split-rows
-                     (into [] (keep index-in-breakouts) column-split-rows))
-        pivot-cols (when column-split-columns
-                     (into [] (keep index-in-breakouts) column-split-columns))]
-    {:pivot-rows pivot-rows
-     :pivot-cols pivot-cols}))
+  (let [{:keys [rows columns values]} (:pivot_table.column_split viz-settings)
+        metadata-provider             (or (:lib/metadata query)
+                                          (lib.metadata.jvm/application-database-metadata-provider (:database query)))
+        mlv2-query                    (lib/query metadata-provider query)
+        breakouts                     (into []
+                                            (map-indexed (fn [i col]
+                                                           (cond-> col
+                                                             true                         (assoc ::idx i)
+                                                             ;; if the col has a card-id, we swap the :lib/source to say
+                                                             ;; source/card this allows `lib/find-matching-column` to properly
+                                                             ;; match a column that has a join-alias but whose source is a
+                                                             ;; model
+                                                             (contains? col :lib/card-id) (assoc :lib/source :source/card))))
+                                            (concat (lib/breakouts-metadata mlv2-query)
+                                                    (lib/aggregations-metadata mlv2-query)))
+        index-in-breakouts            (fn index-in-breakouts [legacy-ref]
+                                        (try
+                                          (::idx (lib.equality/find-column-for-legacy-ref
+                                                  mlv2-query
+                                                  -1
+                                                  legacy-ref
+                                                  breakouts))
+                                          (catch Throwable e
+                                            (log/errorf e "Error finding matching column for ref %s" (pr-str legacy-ref))
+                                            nil)))
+        process-refs                  (fn process-refs [refs]
+                                        (when (seq refs)
+                                          (into [] (keep index-in-breakouts) refs)))
+        pivot-opts                    {:pivot-rows     (process-refs rows)
+                                       :pivot-cols     (process-refs columns)
+                                       :pivot-measures (process-refs values)}]
+    (when (some some? (vals pivot-opts))
+      pivot-opts)))
 
 (mu/defn- column-mapping-for-subquery :- ::pivot-column-mapping
   [num-canonical-cols            :- ::lib.schema.common/int-greater-than-or-equal-to-zero
@@ -529,10 +531,11 @@
      (qp.setup/with-qp-setup [query query]
        (let [rff               (or rff qp.reducible/default-rff)
              query             (lib/query (qp.store/metadata-provider) query)
-             pivot-options     (or
-                                (not-empty (select-keys query [:pivot-rows :pivot-cols]))
-                                (pivot-options query (get-in query [:info :visualization-settings])))
-             query             (assoc-in query [:middleware :pivot-options] pivot-options)
-             all-queries       (generate-queries query pivot-options)
+             pivot-opts        (or
+                                (pivot-options query (get query :viz-settings))
+                                (pivot-options query (get-in query [:info :visualization-settings]))
+                                (not-empty (select-keys query [:pivot-rows :pivot-cols :pivot-measures])))
+             query             (assoc-in query [:middleware :pivot-options] pivot-opts)
+             all-queries       (generate-queries query pivot-opts)
              column-mapping-fn (make-column-mapping-fn query)]
          (process-multiple-queries all-queries rff column-mapping-fn))))))
diff --git a/src/metabase/query_processor/pivot/postprocess.clj b/src/metabase/query_processor/pivot/postprocess.clj
index 7b277a27001..a70414848ed 100644
--- a/src/metabase/query_processor/pivot/postprocess.clj
+++ b/src/metabase/query_processor/pivot/postprocess.clj
@@ -62,10 +62,19 @@
 
 (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]
-  (-> pivot-spec
-      (assoc :pivot-measures (pivot-measures pivot-spec))
-      (assoc :pivot-grouping (pivot-grouping-key (:column-titles pivot-spec)))))
+  [{measure-indices :pivot-measures :as pivot-spec} :- ::pivot-spec]
+  (let [pivot-grouping-key (pivot-grouping-key (:column-titles pivot-spec))]
+    (cond-> pivot-spec
+      ;; if pivot-measures don't already exist (from the pivot qp), we add them ourselves, assuming lowest ID -> highest ID sort order
+      (not (seq measure-indices)) (assoc :pivot-measures (pivot-measures pivot-spec))
+      ;; otherwise, we modify indices to skip over whatever the pivot-grouping idx is, so we pull the correct values per row
+      (seq measure-indices)       (update :pivot-measures (fn [indices]
+                                                            (mapv (fn [idx]
+                                                                    (if (>= idx pivot-grouping-key)
+                                                                      (inc idx)
+                                                                      idx))
+                                                                  indices)))
+      true                       (assoc :pivot-grouping pivot-grouping-key))))
 
 (mu/defn add-totals-settings :- ::pivot-spec
   "Given a pivot-spec map and `viz-settings`, add the `:row-totals?` and `:col-totals?` keys."
diff --git a/src/metabase/query_processor/streaming/csv.clj b/src/metabase/query_processor/streaming/csv.clj
index 99b2d0173e5..ab1a6529d3d 100644
--- a/src/metabase/query_processor/streaming/csv.clj
+++ b/src/metabase/query_processor/streaming/csv.clj
@@ -78,15 +78,14 @@
       (begin! [_ {{:keys [ordered-cols results_timezone format-rows? pivot-export-options pivot?]
                    :or   {format-rows? true
                           pivot?       false}} :data} viz-settings]
-        (let [opts               (when (and pivot? pivot-export-options)
+        (let [col-names          (vec (common/column-titles ordered-cols (::mb.viz/column-settings viz-settings) format-rows?))
+              opts               (when (and pivot? pivot-export-options)
                                    (-> (merge {:pivot-rows []
                                                :pivot-cols []}
                                               pivot-export-options)
-                                       (assoc :column-titles (mapv :display_name ordered-cols))
+                                       (assoc :column-titles col-names)
                                        (qp.pivot.postprocess/add-totals-settings viz-settings)
                                        qp.pivot.postprocess/add-pivot-measures))
-              ;; 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?))
               pivot-grouping-key (qp.pivot.postprocess/pivot-grouping-key col-names)]
 
           ;; initialize the pivot-data
@@ -101,7 +100,7 @@
                    (mapv #(formatter/create-formatter results_timezone % viz-settings format-rows?) ordered-cols))
 
           ;; write the column names for non-pivot tables
-          (when col-names
+          (when (not opts)
             (let [header (m/remove-nth (or pivot-grouping-key (inc (count col-names))) col-names)]
               (write-csv writer [header])
               (.flush writer)))))
diff --git a/src/metabase/query_processor/streaming/json.clj b/src/metabase/query_processor/streaming/json.clj
index 2f851442d98..05cf3de2c42 100644
--- a/src/metabase/query_processor/streaming/json.clj
+++ b/src/metabase/query_processor/streaming/json.clj
@@ -41,8 +41,6 @@
     (reify qp.si/StreamingResultsWriter
       (begin! [_ {{:keys [ordered-cols results_timezone format-rows?]
                    :or   {format-rows? true}} :data} viz-settings]
-        ;; TODO -- wouldn't it make more sense if the JSON downloads used `:name` preferentially? Seeing how JSON is
-        ;; probably going to be parsed programmatically
         (let [cols           (common/column-titles ordered-cols (::mb.viz/column-settings viz-settings) format-rows?)
               pivot-grouping (qp.pivot.postprocess/pivot-grouping-key cols)]
           (when pivot-grouping (vreset! pivot-grouping-idx pivot-grouping))
diff --git a/test/metabase/api/downloads_exports_test.clj b/test/metabase/api/downloads_exports_test.clj
index fb7299965a8..283f34bac7d 100644
--- a/test/metabase/api/downloads_exports_test.clj
+++ b/test/metabase/api/downloads_exports_test.clj
@@ -1061,32 +1061,108 @@
                                                        :type     :native
                                                        :native   {:query "SELECT 1234.567 as A"}}}]
         (testing "CSV downloads respect the formatted/unformatted setting"
-          (let [formatted-json-results   (all-downloads card {:export-format :csv :format-rows true})
-                unformatted-json-results (all-downloads card {:export-format :csv :format-rows false})]
+          (let [formatted-results   (all-downloads card {:export-format :csv :format-rows true})
+                unformatted-results (all-downloads card {:export-format :csv :format-rows false})]
             (is (= {:unsaved-card-download    [["A"] ["1,234.57"]]
                     :card-download            [["A"] ["1,234.57"]]
                     :public-question-download [["A"] ["1,234.57"]]
                     :dashcard-download        [["A"] ["1,234.57"]]
                     :public-dashcard-download [["A"] ["1,234.57"]]}
-                   formatted-json-results))
+                   formatted-results))
             (is (= {:unsaved-card-download    [["A"] ["1234.567"]]
                     :card-download            [["A"] ["1234.567"]]
                     :public-question-download [["A"] ["1234.567"]]
                     :dashcard-download        [["A"] ["1234.567"]]
                     :public-dashcard-download [["A"] ["1234.567"]]}
-                   unformatted-json-results))))
+                   unformatted-results))))
         (testing "JSON downloads respect the formatted/unformatted setting"
-          (let [formatted-json-results   (all-downloads card {:export-format :json :format-rows true})
-                unformatted-json-results (all-downloads card {:export-format :json :format-rows false})]
+          (let [formatted-results   (all-downloads card {:export-format :json :format-rows true})
+                unformatted-results (all-downloads card {:export-format :json :format-rows false})]
             (is (= {:unsaved-card-download    [["A"] ["1,234.57"]]
                     :card-download            [["A"] ["1,234.57"]]
                     :public-question-download [["A"] ["1,234.57"]]
                     :dashcard-download        [["A"] ["1,234.57"]]
                     :public-dashcard-download [["A"] ["1,234.57"]]}
-                   formatted-json-results))
+                   formatted-results))
             (is (= {:unsaved-card-download    [["A"] [1234.567]]
                     :card-download            [["A"] [1234.567]]
                     :public-question-download [["A"] [1234.567]]
                     :dashcard-download        [["A"] [1234.567]]
                     :public-dashcard-download [["A"] [1234.567]]}
-                   unformatted-json-results))))))))
+                   unformatted-results))))))))
+
+(deftest pivot-measures-order-test
+  (testing "A pivot download will use the user-configured measures order (#48442)."
+    (mt/dataset test-data
+      (mt/with-temp [:model/Card card {:display                :pivot
+                                       :dataset_query          {:database (mt/id)
+                                                                :type     :query
+                                                                :query
+                                                                {:source-table (mt/id :products)
+                                                                 :aggregation  [[:count]
+                                                                                [: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}]]}}
+                                       :visualization_settings {:pivot_table.column_split
+                                                                {:rows    [[:field (mt/id :products :category) {:base-type :type/Text}]]
+                                                                 :columns []
+                                                                 :values  [[:aggregation 1]
+                                                                           [:aggregation 0]
+                                                                           [:aggregation 2]]}
+                                                                :column_settings
+                                                                {"[\"name\",\"count\"]" {:column_title "Count Renamed"}
+                                                                 "[\"name\",\"sum\"]"   {:column_title "Sum Renamed"}
+                                                                 "[\"name\",\"avg\"]"   {:column_title "Average Renamed"}}}}]
+        (let [expected-header   ["Category" "Sum of Price" "Count" "Average of Rating"]
+              formatted-results (all-downloads card {:export-format :csv :format-rows false :pivot true})]
+          (is (= {:unsaved-card-download    expected-header
+                  :card-download            expected-header
+                  :public-question-download expected-header
+                  :dashcard-download        expected-header
+                  :public-dashcard-download expected-header}
+                 (update-vals formatted-results first))))
+        (testing "The column title changes are used when format-rows is true"
+          (let [expected-header   ["Category" "Sum Renamed" "Count Renamed" "Average Renamed"]
+                formatted-results (all-downloads card {:export-format :csv :format-rows true :pivot true})]
+            (is (= {:unsaved-card-download    expected-header
+                    :card-download            expected-header
+                    :public-question-download expected-header
+                    :dashcard-download        expected-header
+                    :public-dashcard-download expected-header}
+                   (update-vals formatted-results first)))))))))
+
+(deftest pivot-rows-order-test
+  (testing "A pivot download will use the user-configured rows order."
+    (mt/dataset test-data
+      (mt/with-temp [:model/Card card {:display                :pivot
+                                       :dataset_query          {:database (mt/id)
+                                                                :type     :query
+                                                                :query
+                                                                {:source-table (mt/id :products)
+                                                                 :aggregation  [[:count]]
+                                                                 :breakout     [[:field (mt/id :products :category) {:base-type :type/Text}]
+                                                                                [:field (mt/id :products :created_at) {:base-type :type/DateTime :temporal-unit :year}]]}}
+                                       :visualization_settings {:pivot_table.column_split
+                                                                {:rows    [[:field (mt/id :products :created_at) {:base-type :type/DateTime :temporal-unit :year}]
+                                                                           [:field (mt/id :products :category) {:base-type :type/Text}]]
+                                                                 :columns []
+                                                                 :values  [[:aggregation 0]]}
+                                                                :column_settings
+                                                                {"[\"name\",\"count\"]" {:column_title "Count Renamed"}}}}]
+        (let [expected-header   ["Created At" "Category" "Count"]
+              formatted-results (all-downloads card {:export-format :csv :format-rows false :pivot true})]
+          (is (= {:unsaved-card-download    expected-header
+                  :card-download            expected-header
+                  :public-question-download expected-header
+                  :dashcard-download        expected-header
+                  :public-dashcard-download expected-header}
+                 (update-vals formatted-results first))))
+        (testing "The column title changes are used when format-rows is true"
+          (let [expected-header   ["Created At" "Category" "Count Renamed"]
+                formatted-results (all-downloads card {:export-format :csv :format-rows true :pivot true})]
+            (is (= {:unsaved-card-download    expected-header
+                    :card-download            expected-header
+                    :public-question-download expected-header
+                    :dashcard-download        expected-header
+                    :public-dashcard-download expected-header}
+                   (update-vals formatted-results first)))))))))
diff --git a/test/metabase/query_processor/pivot_test.clj b/test/metabase/query_processor/pivot_test.clj
index ca255bae82c..e2bf1cb0024 100644
--- a/test/metabase/query_processor/pivot_test.clj
+++ b/test/metabase/query_processor/pivot_test.clj
@@ -191,7 +191,7 @@
   (testing "`pivot-options` correctly generates pivot-rows and pivot-cols from a card's viz settings"
     (let [query         (api.pivots/pivot-query false)
           viz-settings  (:visualization_settings (api.pivots/pivot-card))
-          pivot-options {:pivot-rows [1 0], :pivot-cols [2]}]
+          pivot-options {:pivot-rows [1 0], :pivot-cols [2] :pivot-measures nil}]
       (is (= pivot-options
              (#'qp.pivot/pivot-options query viz-settings)))
       (are [num-breakouts expected] (= expected
@@ -217,7 +217,7 @@
                              :columns [[:field "RATING" {:base-type :type/Integer}]]}})
                          [:pivot_table.column_split])
           pivot-options (#'qp.pivot/pivot-options query viz-settings)]
-      (is (= {:pivot-rows [], :pivot-cols []}
+      (is (= {:pivot-rows [], :pivot-cols [] :pivot-measures nil}
              pivot-options))
       (is (= [[0 1] [1] [0] []]
              (#'qp.pivot/breakout-combinations 2 (:pivot-rows pivot-options) (:pivot-cols pivot-options)))))))
@@ -245,7 +245,7 @@
                                  {:rows    [$category]
                                   :columns [$created_at]}})
                 pivot-options (#'qp.pivot/pivot-options query viz-settings)]
-            (is (= {:pivot-rows [0], :pivot-cols [1]}
+            (is (= {:pivot-rows [0], :pivot-cols [1] :pivot-measures nil}
                    pivot-options))
             (is (= [[0 1] [1] [0] []]
                    (#'qp.pivot/breakout-combinations 2 (:pivot-rows pivot-options) (:pivot-cols pivot-options))))
-- 
GitLab