Improve chart performance and gapfilling (#2306)

app/models/account/chartable.rb

@@ -1,139 +1,26 @@
 module Account::Chartable
   extend ActiveSupport::Concern
 
-  class_methods do
-    def balance_series(currency:, period: Period.last_30_days, favorable_direction: "up", view: :balance, interval: nil)
-      raise ArgumentError, "Invalid view type" unless [ :balance, :cash_balance, :holdings_balance ].include?(view.to_sym)
-
-      series_interval = interval || period.interval
-
-      balances = Balance.find_by_sql([
-        balance_series_query,
-        {
-          start_date: period.start_date,
-          end_date: period.end_date,
-          interval: series_interval,
-          target_currency: currency
-        }
-      ])
-
-      balances = gapfill_balances(balances)
-      balances = invert_balances(balances) if favorable_direction == "down"
-
-      values = [ nil, *balances ].each_cons(2).map do |prev, curr|
-        Series::Value.new(
-          date: curr.date,
-          date_formatted: I18n.l(curr.date, format: :long),
-          trend: Trend.new(
-            current: Money.new(balance_value_for(curr, view), currency),
-            previous: prev.nil? ? nil : Money.new(balance_value_for(prev, view), currency),
-            favorable_direction: favorable_direction
-          )
-        )
-      end
-
-      Series.new(
-        start_date: period.start_date,
-        end_date: period.end_date,
-        interval: series_interval,
-        trend: Trend.new(
-          current: Money.new(balance_value_for(balances.last, view) || 0, currency),
-          previous: Money.new(balance_value_for(balances.first, view) || 0, currency),
-          favorable_direction: favorable_direction
-        ),
-        values: values
-      )
-    end
-
-    private
-      def balance_series_query
-        <<~SQL
-          WITH dates as (
-            SELECT generate_series(DATE :start_date, DATE :end_date, :interval::interval)::date as date
-            UNION DISTINCT
-            SELECT :end_date::date AS date -- Ensures we always end on user's "today" date, regardless of interval
-          )
-          SELECT
-            d.date,
-            SUM(CASE WHEN accounts.classification = 'asset' THEN ab.balance ELSE -ab.balance END * COALESCE(er.rate, 1)) as balance,
-            SUM(CASE WHEN accounts.classification = 'asset' THEN ab.cash_balance ELSE -ab.cash_balance END * COALESCE(er.rate, 1)) as cash_balance,
-            SUM(CASE WHEN accounts.classification = 'asset' THEN ab.balance - ab.cash_balance ELSE 0 END * COALESCE(er.rate, 1)) as holdings_balance,
-            COUNT(CASE WHEN accounts.currency <> :target_currency AND er.rate IS NULL THEN 1 END) as missing_rates
-          FROM dates d
-          LEFT JOIN accounts ON accounts.id IN (#{all.select(:id).to_sql})
-          LEFT JOIN balances ab ON (
-            ab.date = d.date AND
-            ab.currency = accounts.currency AND
-            ab.account_id = accounts.id
-          )
-          LEFT JOIN exchange_rates er ON (
-            er.date = ab.date AND
-            er.from_currency = accounts.currency AND
-            er.to_currency = :target_currency
-          )
-          GROUP BY d.date
-          ORDER BY d.date
-        SQL
-      end
-
-      def balance_value_for(balance_record, view)
-        return 0 if balance_record.nil?
-
-        case view.to_sym
-        when :balance then balance_record.balance
-        when :cash_balance then balance_record.cash_balance
-        when :holdings_balance then balance_record.holdings_balance
-        else
-          raise ArgumentError, "Invalid view type: #{view}"
-        end
-      end
-
-      def invert_balances(balances)
-        balances.map do |balance|
-          balance.balance = -balance.balance
-          balance.cash_balance = -balance.cash_balance
-          balance.holdings_balance = -balance.holdings_balance
-          balance
-        end
-      end
-
-      def gapfill_balances(balances)
-        gapfilled = []
-        prev = nil
-
-        balances.each do |curr|
-          if prev.nil?
-            # Initialize first record with zeros if nil
-            curr.balance ||= 0
-            curr.cash_balance ||= 0
-            curr.holdings_balance ||= 0
-          else
-            # Copy previous values for nil fields
-            curr.balance ||= prev.balance
-            curr.cash_balance ||= prev.cash_balance
-            curr.holdings_balance ||= prev.holdings_balance
-          end
-
-          gapfilled << curr
-          prev = curr
-        end
-
-        gapfilled
-      end
-  end
-
   def favorable_direction
     classification == "asset" ? "up" : "down"
   end
 
   def balance_series(period: Period.last_30_days, view: :balance, interval: nil)
-    self.class.where(id: self.id).balance_series(
-      currency: currency,
+    raise ArgumentError, "Invalid view type" unless [ :balance, :cash_balance, :holdings_balance ].include?(view.to_sym)
+
+    @balance_series ||= {}
+
+    memo_key = [ period.start_date, period.end_date, interval ].compact.join("_")
+
+    builder = (@balance_series[memo_key] ||= Balance::ChartSeriesBuilder.new(
+      account_ids: [ id ],
+      currency: self.currency,
       period: period,
-      view: view,
-      interval: interval,
-      favorable_direction: favorable_direction
-    )
+      favorable_direction: favorable_direction,
+      interval: interval
+    ))
+
+    builder.send("#{view}_series")
   end
 
   def sparkline_series

app/models/assistant/function/get_balance_sheet.rb

@@ -50,14 +50,17 @@ class Assistant::Function::GetBalanceSheet < Assistant::Function
       if period.start_date == Date.current
         []
       else
-        balance_series = scope.balance_series(
+        account_ids = scope.pluck(:id)
+
+        builder = Balance::ChartSeriesBuilder.new(
+          account_ids: account_ids,
           currency: family.currency,
           period: period,
-          interval: "1 month",
-          favorable_direction: "up"
+          favorable_direction: "up",
+          interval: "1 month"
         )
 
-        to_ai_time_series(balance_series)
+        to_ai_time_series(builder.balance_series)
       end
     end
 

app/models/balance/chart_series_builder.rb

@@ -0,0 +1,141 @@
+class Balance::ChartSeriesBuilder
+  def initialize(account_ids:, currency:, period: Period.last_30_days, interval: "1 day", favorable_direction: "up")
+    @account_ids = account_ids
+    @currency = currency
+    @period = period
+    @interval = interval
+    @favorable_direction = favorable_direction
+  end
+
+  def balance_series
+    build_series_for(:balance)
+  end
+
+  def cash_balance_series
+    build_series_for(:cash_balance)
+  end
+
+  def holdings_balance_series
+    build_series_for(:holdings_balance)
+  end
+
+  private
+    attr_reader :account_ids, :currency, :period, :favorable_direction
+
+    def interval
+      @interval || period.interval
+    end
+
+    def build_series_for(column)
+      values = query_data.map do |datum|
+        Series::Value.new(
+          date: datum.date,
+          date_formatted: I18n.l(datum.date, format: :long),
+          value: Money.new(datum.send(column), currency),
+          trend: Trend.new(
+            current: Money.new(datum.send(column), currency),
+            previous: Money.new(datum.send("previous_#{column}"), currency),
+            favorable_direction: favorable_direction
+          )
+        )
+      end
+
+      Series.new(
+        start_date: period.start_date,
+        end_date: period.end_date,
+        interval: interval,
+        values: values,
+        favorable_direction: favorable_direction
+      )
+    end
+
+    def query_data
+      @query_data ||= Balance.find_by_sql([
+        query,
+        {
+          account_ids: account_ids,
+          target_currency: currency,
+          start_date: period.start_date,
+          end_date: period.end_date,
+          interval: interval,
+          sign_multiplier: sign_multiplier
+        }
+      ])
+    end
+
+    # Since the query aggregates the *net* of assets - liabilities, this means that if we're looking at
+    # a single liability account, we'll get a negative set of values.  This is not what the user expects
+    # to see.  When favorable direction is "down" (i.e. liability, decrease is "good"), we need to invert
+    # the values by multiplying by -1.
+    def sign_multiplier
+      favorable_direction == "down" ? -1 : 1
+    end
+
+    def query
+      <<~SQL
+        WITH dates AS (
+          SELECT generate_series(DATE :start_date, DATE :end_date, :interval::interval)::date AS date
+          UNION DISTINCT
+          SELECT :end_date::date  -- Pass in date to ensure timezone-aware "today" date
+        ), aggregated_balances AS (
+          SELECT
+            d.date,
+            -- Total balance (assets positive, liabilities negative)
+            SUM(
+              CASE WHEN accounts.classification = 'asset'
+                THEN COALESCE(last_bal.balance, 0)
+                ELSE -COALESCE(last_bal.balance, 0)
+              END * COALESCE(er.rate, 1) * :sign_multiplier::integer
+            ) AS balance,
+            -- Cash-only balance
+            SUM(
+              CASE WHEN accounts.classification = 'asset'
+                THEN COALESCE(last_bal.cash_balance, 0)
+                ELSE -COALESCE(last_bal.cash_balance, 0)
+              END * COALESCE(er.rate, 1) * :sign_multiplier::integer
+            ) AS cash_balance,
+            -- Holdings value (balance ‑ cash)
+            SUM(
+              CASE WHEN accounts.classification = 'asset'
+                THEN COALESCE(last_bal.balance, 0) - COALESCE(last_bal.cash_balance, 0)
+                ELSE 0
+              END * COALESCE(er.rate, 1) * :sign_multiplier::integer
+            ) AS holdings_balance
+          FROM dates d
+          JOIN accounts ON accounts.id = ANY(array[:account_ids]::uuid[])
+
+          -- Last observation carried forward (LOCF), use the most recent balance on or before the chart date
+          LEFT JOIN LATERAL (
+            SELECT b.balance, b.cash_balance
+            FROM balances b
+            WHERE b.account_id = accounts.id
+              AND b.date <= d.date
+            ORDER BY b.date DESC
+            LIMIT 1
+          ) last_bal ON TRUE
+
+          -- Last observation carried forward (LOCF), use the most recent exchange rate on or before the chart date
+          LEFT JOIN LATERAL (
+            SELECT er.rate
+            FROM exchange_rates er
+            WHERE er.from_currency = accounts.currency
+              AND er.to_currency = :target_currency
+              AND er.date <= d.date
+            ORDER BY er.date DESC
+            LIMIT 1
+          ) er ON TRUE
+          GROUP BY d.date
+        )
+        SELECT
+          date,
+          balance,
+          cash_balance,
+          holdings_balance,
+          COALESCE(LAG(balance) OVER (ORDER BY date), 0) AS previous_balance,
+          COALESCE(LAG(cash_balance) OVER (ORDER BY date), 0) AS previous_cash_balance,
+          COALESCE(LAG(holdings_balance) OVER (ORDER BY date), 0) AS previous_holdings_balance
+        FROM aggregated_balances
+        ORDER BY date
+      SQL
+    end
+end

app/models/balance_sheet.rb

@@ -88,7 +88,20 @@ class BalanceSheet
   end
 
   def net_worth_series(period: Period.last_30_days)
-    active_accounts.balance_series(currency: currency, period: period, favorable_direction: "up")
+    memo_key = [ period.start_date, period.end_date ].compact.join("_")
+
+    @net_worth_series ||= {}
+
+    account_ids = active_accounts.pluck(:id)
+
+    builder = (@net_worth_series[memo_key] ||= Balance::ChartSeriesBuilder.new(
+      account_ids: account_ids,
+      currency: currency,
+      period: period,
+      favorable_direction: "up"
+    ))
+
+    builder.balance_series
   end
 
   def currency

app/models/series.rb

@@ -1,9 +1,21 @@
 class Series
-  attr_reader :start_date, :end_date, :interval, :trend, :values
+  # Behave like an Array whose elements are the `Value` structs stored in `values`
+  include Enumerable
+
+  # Forward any undefined method calls (e.g. `first`, `[]`, `map`) to `values`
+  delegate_missing_to :values
+
+  # Enumerable needs `#each`
+  def each(&block)
+    values.each(&block)
+  end
+
+  attr_reader :start_date, :end_date, :interval, :trend, :values, :favorable_direction
 
   Value = Struct.new(
     :date,
     :date_formatted,
+    :value,
     :trend,
     keyword_init: true
   )
@@ -29,6 +41,7 @@ class Series
           Value.new(
             date: curr_value[:date],
             date_formatted: I18n.l(curr_value[:date], format: :long),
+            value: curr_value[:value],
             trend: Trend.new(
               current: curr_value[:value],
               previous: prev_value&.[](:value)
@@ -39,19 +52,29 @@ class Series
     end
   end
 
-  def initialize(start_date:, end_date:, interval:, trend:, values:)
+  def initialize(start_date:, end_date:, interval:, values:, favorable_direction: "up")
     @start_date = start_date
     @end_date = end_date
     @interval = interval
-    @trend = trend
     @values = values
+    @favorable_direction = favorable_direction
   end
 
-  def current
-    values.last.trend.current
+  def trend
+    @trend ||= Trend.new(
+      current: values.last&.value,
+      previous: values.first&.value,
+      favorable_direction: favorable_direction
+    )
   end
 
-  def any?
-    values.any?
+  def as_json
+    {
+      start_date: start_date,
+      end_date: end_date,
+      interval: interval,
+      trend: trend,
+      values: values.map { |v| { date: v.date, date_formatted: v.date_formatted, value: v.value, trend: v.trend } }
+    }
   end
 end