Reduce the amount of work performed when instantiating AR models

We don't know which attributes will or won't be used, and we don't want
to create massive bottlenecks at instantiation. Rather than doing *any*
iteration over types and values, we can lazily instantiate the object.

The lazy attribute hash should not fully implement hash, or subclass
hash at any point in the future. It is not meant to be a replacement,
but instead implement its own interface which happens to overlap.

activerecord/lib/active_record/attribute_methods/primary_key.rb

@@ -120,6 +120,7 @@ module ActiveRecord
         def primary_key=(value)
           @primary_key        = value && value.to_s
           @quoted_primary_key = nil
+          @attributes_builder = nil
         end
       end
     end

activerecord/lib/active_record/attribute_set.rb

@@ -2,8 +2,6 @@ require 'active_record/attribute_set/builder'
 
 module ActiveRecord
   class AttributeSet # :nodoc:
-    delegate :keys, to: :initialized_attributes
-
     def initialize(attributes)
       @attributes = attributes
     end
@@ -25,6 +23,10 @@ module ActiveRecord
       attributes.key?(name) && self[name].initialized?
     end
 
+    def keys
+      attributes.initialized_keys
+    end
+
     def fetch_value(name, &block)
       self[name].value(&block)
     end
@@ -43,7 +45,7 @@ module ActiveRecord
     end
 
     def initialize_dup(_)
-      @attributes = attributes.transform_values(&:dup)
+      @attributes = attributes.dup
       super
     end
 
@@ -58,12 +60,6 @@ module ActiveRecord
       end
     end
 
-    def ensure_initialized(key)
-      unless self[key].initialized?
-        write_from_database(key, nil)
-      end
-    end
-
     protected
 
     attr_reader :attributes

activerecord/lib/active_record/attribute_set/builder.rb

@@ -1,35 +1,83 @@
 module ActiveRecord
   class AttributeSet # :nodoc:
     class Builder # :nodoc:
-      attr_reader :types
+      attr_reader :types, :always_initialized
 
-      def initialize(types)
+      def initialize(types, always_initialized = nil)
         @types = types
+        @always_initialized = always_initialized
       end
 
       def build_from_database(values = {}, additional_types = {})
-        attributes = build_attributes_from_values(values, additional_types)
-        add_uninitialized_attributes(attributes)
+        if always_initialized && !values.key?(always_initialized)
+          values[always_initialized] = nil
+        end
+
+        attributes = LazyAttributeHash.new(types, values, additional_types)
         AttributeSet.new(attributes)
       end
 
       private
+    end
+  end
 
-      def build_attributes_from_values(values, additional_types)
-        values.each_with_object({}) do |(name, value), hash|
-          type = additional_types.fetch(name, types[name])
-          hash[name] = Attribute.from_database(name, value, type)
-        end
-      end
+  class LazyAttributeHash
+    delegate :select, :transform_values, to: :materialize
+    delegate :[], :[]=, :freeze, to: :delegate_hash
 
-      def add_uninitialized_attributes(attributes)
-        types.each_key do |name|
-          next if attributes.key? name
-          type = types[name]
-          attributes[name] =
-            Attribute.uninitialized(name, type)
+    def initialize(types, values, additional_types)
+      @types = types
+      @values = values
+      @additional_types = additional_types
+      @materialized = false
+      @delegate_hash = {}
+      assign_default_proc
+    end
+
+    def key?(key)
+      delegate_hash.key?(key) || values.key?(key) || types.key?(key)
+    end
+
+    def initialized_keys
+      delegate_hash.keys | values.keys
+    end
+
+    def initialize_dup(_)
+      @delegate_hash = delegate_hash.transform_values(&:dup)
+      assign_default_proc
+      super
+    end
+
+    def initialize_clone(_)
+      @delegate_hash = delegate_hash.clone
+      super
+    end
+
+    protected
+
+    attr_reader :types, :values, :additional_types, :delegate_hash
+
+    private
+
+    def assign_default_proc
+      delegate_hash.default_proc = proc do |hash, name|
+        type = additional_types.fetch(name, types[name])
+
+        if values.key?(name)
+          hash[name] = Attribute.from_database(name, values[name], type)
+        elsif type
+          hash[name] = Attribute.uninitialized(name, type)
         end
       end
     end
+
+    def materialize
+      unless @materialized
+        values.each_key { |key| delegate_hash[key] }
+        types.each_key { |key| delegate_hash[key] }
+        @materialized = true
+      end
+      delegate_hash
+    end
   end
 end

activerecord/lib/active_record/core.rb

@@ -536,8 +536,6 @@ module ActiveRecord
     end
 
     def init_internals
-      @attributes.ensure_initialized(self.class.primary_key)
-
       @aggregation_cache        = {}
       @association_cache        = {}
       @readonly                 = false

activerecord/lib/active_record/model_schema.rb

@@ -231,7 +231,7 @@ module ActiveRecord
       end
 
       def attributes_builder # :nodoc:
-        @attributes_builder ||= AttributeSet::Builder.new(column_types)
+        @attributes_builder ||= AttributeSet::Builder.new(column_types, primary_key)
       end
 
       def column_types # :nodoc:

activerecord/test/cases/attribute_set_test.rb

@@ -123,6 +123,15 @@ module ActiveRecord
       assert_nil attributes.fetch_value(:bar)
     end
 
+    test "the primary_key is always initialized" do
+      builder = AttributeSet::Builder.new({ foo: Type::Integer.new }, :foo)
+      attributes = builder.build_from_database
+
+      assert attributes.key?(:foo)
+      assert_equal [:foo], attributes.keys
+      assert attributes[:foo].initialized?
+    end
+
     class MyType
       def type_cast_from_user(value)
         return if value.nil?