fix(python): cover FP16 GPU index queries (#3819)

Keep FP16 vectors through IVF-PQ index creation and search in the
synchronous and asynchronous quickstarts. Add native and conditional
CUDA regression coverage using a float64 query against an explicitly
typed FP16 column. This is test and documentation-example coverage for
dtype preservation already present in Lance; it does not change
production code.

Related #2120.

Validation on `bbad20940632b21a9ee85bdcb46b4cc19dbd65e3` with Python
3.12.12, Lance 15.0.0-beta.1, PyArrow 24.0.0, and Torch 2.14.1:

- The native FP16 parameter, Arrow FP16 ingestion test, and both
quickstarts passed: 4 passed, 1 CUDA-only skip.
- The two affected test files passed: 386 passed, 1 CUDA-only skip (`-m
"not slow and not s3_test"`).
- Ruff 0.9.9 lint and formatting checks passed for the Python package.

Real CUDA execution has not been performed. Torch is installed, but
`torch.cuda.is_available()` is false; `test_create_f16_table[cuda]`
skips with `CUDA not available`. The native CPU result does not
establish that the original GPU failure is resolved. Issue #2120 remains
open for verification on an actual CUDA device.

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>

python/python/tests/docs/test_basic.py

@@ -91,19 +91,13 @@ def test_quickstart(tmp_path):
         }
     )
     # --8<-- [end:alter_columns_vector]
-    # Change it back since we can get a panic with fp16
-    tbl.alter_columns(
-        {
-            "path": "vector",
-            "data_type": pa.list_(pa.float32(), list_size=2),
-        }
-    )
     # --8<-- [start:drop_columns]
     tbl.drop_columns(["dbl_price"])
     # --8<-- [end:drop_columns]
     # --8<-- [start:create_index]
     tbl.create_index(num_sub_vectors=1)
     # --8<-- [end:create_index]
+    tbl.search([100, 100]).limit(2).to_pandas()
     # --8<-- [start:delete_rows]
     tbl.delete("item = 'fizz'")
     # --8<-- [end:delete_rows]
@@ -185,13 +179,6 @@ async def test_quickstart_async(tmp_path):
         }
     )
     # --8<-- [end:alter_columns_async_vector]
-    # Change it back since we can get a panic with fp16
-    await tbl.alter_columns(
-        {
-            "path": "vector",
-            "data_type": pa.list_(pa.float32(), list_size=2),
-        }
-    )
     # --8<-- [start:drop_columns_async]
     await tbl.drop_columns(["dbl_price"])
     # --8<-- [end:drop_columns_async]
@@ -200,6 +187,7 @@ async def test_quickstart_async(tmp_path):
     # --8<-- [start:create_index_async]
     await tbl.create_index("vector")
     # --8<-- [end:create_index_async]
+    await tbl.vector_search([100, 100]).limit(2).to_pandas()
     # --8<-- [start:delete_rows_async]
     await tbl.delete("item = 'fizz'")
     # --8<-- [end:delete_rows_async]

python/python/tests/test_table.py

@@ -3965,26 +3965,42 @@ def test_create_f16_table_from_arrow_data(mem_db: DBConnection):
     assert "s-2" in expected["text"].to_pylist()
 
 
-def test_create_f16_table(mem_db: DBConnection):
+@pytest.mark.parametrize("accelerator", [None, "cuda"])
+def test_create_f16_table(tmp_path, accelerator):
+    if accelerator == "cuda":
+        torch = pytest.importorskip("torch")
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA not available")
+
     class MyTable(LanceModel):
         text: str
         vector: Vector(32, value_type=pa.float16())
 
     rng = np.random.default_rng(42)
+    vectors = rng.standard_normal((512, 32)).astype(np.float16)
     df = pa.table(
         {
             "text": [f"s-{i}" for i in range(512)],
-            "vector": [rng.standard_normal(32).astype(np.float16) for _ in range(512)],
+            "vector": list(vectors),
         }
     )
-    table = mem_db.create_table(
+    db = lancedb.connect(tmp_path)
+    table = db.create_table(
         "f16_tbl",
         schema=MyTable,
     )
     table.add(df)
-    table.create_index(num_partitions=2, num_sub_vectors=2)
+    table.create_index(
+        "vector",
+        config=IvfPq(
+            num_partitions=2,
+            num_sub_vectors=2,
+            accelerator=accelerator,
+        ),
+    )
 
-    query = df["vector"][2].as_py()
+    # Match the issue's float64 query against an explicitly typed float16 column.
+    query = vectors[2].astype(np.float64)
     expected = table.search(query).limit(2).to_arrow()
 
     assert "s-2" in expected["text"].to_pylist()