fix(db): preserve union types in select result extraction

.changeset/fix-ref-union-collapse.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/db': patch
+---
+
+fix: preserve discriminated union types when selecting object fields via `.select()`

packages/db/src/query/builder/types.ts

@@ -322,9 +322,16 @@ export type ResultTypeFromSelect<TSelectObject> =
                 : // includes subquery (bare QueryBuilder) — produces a child Collection
                   TSelectObject[K] extends QueryBuilder<infer TChildContext>
                   ? Collection<GetResult<TChildContext>>
-                  : // Ref (full object ref or spread with RefBrand) - recursively process properties
+                  : // Ref (full object ref or spread with RefBrand)
                     TSelectObject[K] extends Ref<infer _T>
-                    ? ExtractRef<TSelectObject[K]>
+                    ? // When the branded type is a union, extract it directly to
+                      // preserve the discriminated union. ExtractRef would collapse
+                      // it via keyof/Prettify which only yield common keys.
+                      true extends IsUnion<ExtractRefBrand<TSelectObject[K]>>
+                      ? IsNullableRef<TSelectObject[K]> extends true
+                        ? ExtractRefBrand<TSelectObject[K]> | undefined
+                        : ExtractRefBrand<TSelectObject[K]>
+                      : ExtractRef<TSelectObject[K]>
                     : // RefLeaf (simple property ref like user.name)
                       TSelectObject[K] extends RefLeaf<infer T>
                       ? IsNullableRef<TSelectObject[K]> extends true
@@ -372,6 +379,9 @@ export type SelectResult<TSelect> =
     ? ResultTypeFromSelect<TSelect>
     : ResultTypeFromSelectValue<TSelect>
 
+// Extract the raw type stored in a RefLeaf's brand
+type ExtractRefBrand<T> = T extends RefLeaf<infer U> ? U : never
+
 // Extract Ref or subobject with a spread or a Ref
 type ExtractRef<T> = Prettify<ResultTypeFromSelect<WithoutRefBrand<T>>>
 
@@ -1068,6 +1078,14 @@ type IsPlainObject<T> = T extends unknown
 
 type IsAny<T> = 0 extends 1 & T ? true : false
 
+// Detects whether T is a union type (e.g., A | B | C returns true, A returns false)
+type IsUnion<T, U = T> = T extends unknown
+  ? [U] extends [T]
+    ? false
+    : true
+  : false
+
+
 /**
  * JsBuiltIns - List of JavaScript built-ins
  */

packages/db/tests/query/select.test-d.ts

@@ -1,6 +1,6 @@
 import { describe, expectTypeOf, test } from 'vitest'
 import { createCollection } from '../../src/collection/index.js'
-import { createLiveQueryCollection } from '../../src/query/index.js'
+import { createLiveQueryCollection, eq } from '../../src/query/index.js'
 import { mockSyncCollectionOptions } from '../utils.js'
 import { upper } from '../../src/query/builder/functions.js'
 import type { OutputWithVirtual } from '../utils.js'
@@ -109,6 +109,67 @@ describe(`select types`, () => {
     expectTypeOf(results).toMatchTypeOf<OutputWithVirtualKeyed<Expected>>()
   })
 
+  test(`select preserves union types and where works on common keys`, () => {
+    type ItemDocument =
+      | { type: 'pdf'; url: string; pages: number }
+      | { type: 'image'; url: string; width: number; height: number }
+      | { type: 'legacy'; path: string }
+
+    type Item = { id: number; name: string; document: ItemDocument }
+
+    const items = createCollection(
+      mockSyncCollectionOptions<Item>({
+        id: `union-field-items`,
+        getKey: (i) => i.id,
+        initialData: [],
+      }),
+    )
+
+    // Filtering by a common key of the union should compile,
+    // and the result should preserve the full discriminated union
+    const col = createLiveQueryCollection((q) =>
+      q
+        .from({ i: items })
+        .where(({ i }) => eq(i.document.type, `pdf`))
+        .select(({ i }) => ({
+          id: i.id,
+          document: i.document,
+        })),
+    )
+
+    const result = col.toArray[0]!
+    expectTypeOf(result.document).toEqualTypeOf<ItemDocument>()
+  })
+
+  test(`select preserves union when collection type is a union`, () => {
+    type DocV1 = { version: 1; title: string }
+    type DocV2 = { version: 2; title: string; subtitle: string }
+    type Doc = DocV1 | DocV2
+
+    const docs = createCollection(
+      mockSyncCollectionOptions<Doc>({
+        id: `union-collection`,
+        getKey: (d) => d.title,
+        initialData: [],
+      }),
+    )
+
+    // Without select — union preserved
+    const col1 = createLiveQueryCollection((q) => q.from({ d: docs }))
+    const r1 = col1.toArray[0]!
+    expectTypeOf(r1).toMatchTypeOf<Doc>()
+
+    // With select on individual fields — per-field unions (not top-level union)
+    const col2 = createLiveQueryCollection((q) =>
+      q.from({ d: docs }).select(({ d }) => ({
+        version: d.version,
+        title: d.title,
+      })),
+    )
+    const r2 = col2.toArray[0]!
+    expectTypeOf(r2).toMatchTypeOf<{ version: 1 | 2; title: string }>()
+  })
+
   test(`nested spread preserves object structure types`, () => {
     const users = createUsers()
     const col = createLiveQueryCollection((q) => {