diff --git a/lib/esse/active_record/callbacks.rb b/lib/esse/active_record/callbacks.rb
index e5f3873..70b1c9a 100644
--- a/lib/esse/active_record/callbacks.rb
+++ b/lib/esse/active_record/callbacks.rb
@@ -5,8 +5,9 @@ module ActiveRecord
     class Callback
       attr_reader :repo, :options, :block_result
 
-      def initialize(repo:, block_result: nil, **kwargs)
+      def initialize(repo:, block_result: nil, with: nil, **kwargs)
         @repo = repo
+        @with = with
         @options = kwargs
         @block_result = block_result
       end
diff --git a/lib/esse/active_record/callbacks/indexing_on_update.rb b/lib/esse/active_record/callbacks/indexing_on_update.rb
index 0901b38..5593d87 100644
--- a/lib/esse/active_record/callbacks/indexing_on_update.rb
+++ b/lib/esse/active_record/callbacks/indexing_on_update.rb
@@ -3,13 +3,6 @@
 module Esse::ActiveRecord
   module Callbacks
     class IndexingOnUpdate < Callback
-      attr_reader :update_with
-
-      def initialize(with: :index, **kwargs, &block)
-        @update_with = with
-        super(**kwargs, &block)
-      end
-
       def call(model)
         record = block_result || model
 
@@ -43,7 +36,7 @@ def call(model)
       def update_document(document)
         return if document.ignore_on_index?
 
-        if update_with == :update
+        if @with == :update || (@with.nil? && repo.lazy_document_attributes.any?)
           begin
             repo.index.update(document, **options)
           rescue Esse::Transport::NotFoundError
diff --git a/lib/esse/active_record/collection.rb b/lib/esse/active_record/collection.rb
index a7c7852..9c027b7 100644
--- a/lib/esse/active_record/collection.rb
+++ b/lib/esse/active_record/collection.rb
@@ -90,7 +90,7 @@ def each_batch_ids
       end
 
       def count
-        dataset.except(:includes, :preload, :eager_load, :order, :limit, :offset).count
+        dataset.except(:includes, :preload, :eager_load, :group, :order, :limit, :offset).count
       end
       alias_method :size, :count
 
diff --git a/spec/esse/active_record/callbacks/indexing_on_update_spec.rb b/spec/esse/active_record/callbacks/indexing_on_update_spec.rb
index 44619d3..8ad7bd2 100644
--- a/spec/esse/active_record/callbacks/indexing_on_update_spec.rb
+++ b/spec/esse/active_record/callbacks/indexing_on_update_spec.rb
@@ -6,9 +6,9 @@
   describe '.initialize' do
     let(:repo) { instance_double(Esse::Repository) }
 
-    it 'sets update_with' do
+    it 'sets @with' do
       callback = described_class.new(repo: repo, with: :update)
-      expect(callback.update_with).to eq(:update)
+      expect(callback.instance_variable_get(:@with)).to eq(:update)
     end
 
     it 'sets options' do
@@ -96,5 +96,30 @@
         state.update!(name: 'IL')
       end
     end
+
+    context 'when not calling with in a repository with lazy attributes' do
+      before do
+        StatesIndex::State.lazy_document_attribute :total_counties do |docs|
+          ::County.where(state_id: docs.map(&:id)).group(:state_id).count
+        end
+      end
+
+      after do
+        StatesIndex::State.instance_variable_set(:@lazy_document_attributes, {}.freeze)
+      end
+
+      it { expect(StatesIndex::State.lazy_document_attributes).not_to be_empty }
+
+      it 'updates the record using :update action to avoid losing lazy attributes' do
+        expect(StatesIndex).to receive(:update).and_call_original
+        expect(StatesIndex).to esse_receive_request(:update).with(
+          id: state.id,
+          index: StatesIndex.index_name,
+          body: {doc: { name: 'IL' } },
+        ).and_return(ok_response)
+
+        state.update!(name: 'IL')
+      end
+    end
   end
 end