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 @@ def get_primary_key(base_name) #:nodoc:
         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 @@
 
 module ActiveRecord
   class AttributeSet # :nodoc:
-    delegate :keys, to: :initialized_attributes
-
     def initialize(attributes)
       @attributes = attributes
     end
@@ -25,6 +23,10 @@ def key?(name)
       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 @@ def freeze
     end
 
     def initialize_dup(_)
-      @attributes = attributes.transform_values(&:dup)
+      @attributes = attributes.dup
       super
     end
 
@@ -58,12 +60,6 @@ def reset(key)
       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

 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
+
+  class LazyAttributeHash
+    delegate :select, :transform_values, to: :materialize
+    delegate :[], :[]=, :freeze, to: :delegate_hash
+
+    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 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)
+    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 add_uninitialized_attributes(attributes)
-        types.each_key do |name|
-          next if attributes.key? name
-          type = types[name]
-          attributes[name] =
-            Attribute.uninitialized(name, type)
-        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 @@ def to_ary # :nodoc:
     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 @@ def table_exists?
       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 @@ class AttributeSetTest < ActiveRecord::TestCase
       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?