diff --git a/lance_ray/index.py b/lance_ray/index.py index 2fd07591..1cb386fe 100755 --- a/lance_ray/index.py +++ b/lance_ray/index.py @@ -602,12 +602,13 @@ def create_scalar_index( is_supported = ( pa.types.is_integer(value_type) or pa.types.is_floating(value_type) + or pa.types.is_date32(value_type) or pa.types.is_string(value_type) or pa.types.is_large_string(value_type) ) if not is_supported: raise TypeError( - f"Column {column} must be numeric or string type for " + f"Column {column} must be numeric, date32, or string type for " f"{index_type} index, got {value_type}" ) case "LABEL_LIST": diff --git a/tests/test_vector_index_options.py b/tests/test_vector_index_options.py index 34312f3e..a04a5155 100644 --- a/tests/test_vector_index_options.py +++ b/tests/test_vector_index_options.py @@ -72,7 +72,7 @@ def id(self) -> int: class _FakeLanceSchema: def field(self, column: str) -> "_FakeLanceField": - if column not in {"value", "text", "labels"}: + if column not in {"value", "text", "labels", "event_date"}: raise KeyError(column) return _FakeLanceField() @@ -87,6 +87,8 @@ def field(self, column: str) -> "_FakeField": return _FakeField(column, pa.string()) if column == "labels": return _FakeField(column, pa.list_(pa.string())) + if column == "event_date": + return _FakeField(column, pa.date32()) else: raise KeyError(column) @@ -97,6 +99,7 @@ def __iter__(self) -> Any: _FakeField("value", pa.int64()), _FakeField("text", pa.string()), _FakeField("labels", pa.list_(pa.string())), + _FakeField("event_date", pa.date32()), ] ) @@ -467,6 +470,7 @@ def test_create_index_rejects_invalid_num_segments( ("index_type", "column"), [ ("BTREE", "value"), + ("BTREE", "event_date"), ("BITMAP", "value"), ("INVERTED", "text"), ("FTS", "text"), @@ -474,6 +478,7 @@ def test_create_index_rejects_invalid_num_segments( ("BLOOMFILTER", "value"), ("RTREE", "value"), ("LABEL_LIST", "labels"), + ("ZONEMAP", "event_date"), ], ) def test_create_scalar_index_uses_segment_path(