rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-01-12 23:02:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

bug: add a test for fp16 (#837)

Browse Source 
Add test to ingest fp16 to a database
This commit is contained in:
Lei Xu
2024-01-20 16:23:28 -08:00
committed by Weston Pace
parent 0c580abd70
commit 97d033dfd6
4 changed files with 36 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/lancedb/remote/client.py
View File

@@ -23,8 +23,8 @@ from pydantic import BaseModel
from lancedb.common import Credential
from lancedb.remote import VectorQuery, VectorQueryResult
from lancedb.remote.errors import LanceDBClientError
from lancedb.remote.connection_timeout import LanceDBClientHTTPAdapterFactory
from lancedb.remote.errors import LanceDBClientError
ARROW_STREAM_CONTENT_TYPE = "application/vnd.apache.arrow.stream"

6
python/lancedb/remote/connection_timeout.py
View File

@@ -20,13 +20,13 @@
# https://github.com/urllib3/urllib3/pull/3275
import datetime
import os
import logging
import os
from requests.adapters import HTTPAdapter
from urllib3.poolmanager import PoolManager
from urllib3.connectionpool import HTTPSConnectionPool
from urllib3.connection import HTTPSConnection
from urllib3.connectionpool import HTTPSConnectionPool
from urllib3.poolmanager import PoolManager
def get_client_connection_timeout() -> int:

12
python/lancedb/table.py
View File

@@ -1327,13 +1327,15 @@ def _sanitize_vector_column(
elif not pa.types.is_fixed_size_list(vec_arr.type):
raise TypeError(f"Unsupported vector column type: {vec_arr.type}")
vec_arr = ensure_fixed_size_list_of_f32(vec_arr)
vec_arr = ensure_fixed_size_list(vec_arr)
data = data.set_column(
data.column_names.index(vector_column_name), vector_column_name, vec_arr
)
has_nans = pc.any(pc.is_nan(vec_arr.values)).as_py()
if has_nans:
# Use numpy to check for NaNs, because as pyarrow 14.0.2 does not have `is_nan`
# kernel over f16 types.
values_np = vec_arr.values.to_numpy(zero_copy_only=False)
if np.isnan(values_np).any():
data = _sanitize_nans(
data, fill_value, on_bad_vectors, vec_arr, vector_column_name
)
@@ -1341,9 +1343,9 @@ def _sanitize_vector_column(
return data
def ensure_fixed_size_list_of_f32(vec_arr):
def ensure_fixed_size_list(vec_arr) -> pa.FixedSizeListArray:
values = vec_arr.values
if not pa.types.is_float32(values.type):
if not (pa.types.is_float16(values.type) or pa.types.is_float32(values.type)):
values = values.cast(pa.float32())
if pa.types.is_fixed_size_list(vec_arr.type):
list_size = vec_arr.type.list_size