Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add numpy 2 support #434

Merged
merged 11 commits into from
Oct 17, 2024
+83 −11
Merged

Add numpy 2 support #434

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
fd2e731
Add numpy 2 support
jparismorgan Jul 3, 2024
a0f8af4
Merge branch 'main' of https://github.com/TileDB-Inc/TileDB-Vector-Se…
jparismorgan Jul 11, 2024
8245f04
update to tiledb-cloud>=0.12.15
jparismorgan Jul 11, 2024
eb58404
fix casting errors and add numpy 2 ruff lint
jparismorgan Jul 12, 2024
f18af1e
fix bug in cast to int
jparismorgan Jul 12, 2024
c8192f4
add workaround for np.in1d bug
jparismorgan Jul 15, 2024
20b6d4b
format
jparismorgan Jul 15, 2024
45853f2
Merge branch 'main' of https://github.com/TileDB-Inc/TileDB-Vector-Se…
jparismorgan Oct 14, 2024
712a276
Update CI to also run python tests w
jparismorgan Oct 15, 2024
5697b1d
do not error if numpy2 CI job fails, fix flaky test
jparismorgan Oct 16, 2024
3888e2a
cleanup
jparismorgan Oct 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ repos:
- id: prettier

- repo: https://github.com/charliermarsh/ruff-pre-commit
rev: "v0.0.265"
rev: "v0.4.4"
hooks:
- id: ruff
args: [--fix, --exit-non-zero-on-fix]
Expand Down
10 changes: 5 additions & 5 deletions apis/python/src/tiledb/vector_search/ingestion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -401,13 +401,13 @@ def read_source_metadata(
) -> Tuple[int, int, np.dtype]:
if source_type == "TILEDB_ARRAY":
schema = tiledb.ArraySchema.load(source_uri)
size = schema.domain.dim(1).domain[1] + 1
dimensions = schema.domain.dim(0).domain[1] + 1
size = np.int64(schema.domain.dim(1).domain[1]) + 1
dimensions = np.int64(schema.domain.dim(0).domain[1]) + 1
Comment on lines +414 to +415
Copy link
Contributor Author

@jparismorgan jparismorgan Jul 12, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We make this change b/c of changes to numpy data type promotion.

With numpy 1 and the original code we were getting:

[ingestion@read_source_metadata] schema.domain.dim(1).domain[1]: 2147483647 <class 'numpy.int32'>
[ingestion@read_source_metadata] size: 2147483648 <class 'numpy.int64'>
[ingestion@read_source_metadata] dimensions: 3 <class 'numpy.int64'>

With numpy 2 and the original code we instead were getting:

[ingestion@read_source_metadata] schema.domain.dim(1).domain[1]: 2147483647 <class 'numpy.int32'>
[ingestion@read_source_metadata] size: -2147483648 <class 'numpy.int32'>
[ingestion@read_source_metadata] dimensions: 3 <class 'numpy.int32'>

This is because in size = schema.domain.dim(1).domain[1] + 1 the + 1 used to cause a cast to int64 (which is required because 2147483647 is int32 max).

As mentioned in https://numpy.org/devdocs/numpy_2_0_migration_guide.html#changes-to-numpy-data-type-promotion, numpy no longer casts automatically:

Screenshot 2024-07-12 at 3 55 46 PM

So here we explicitly cast to int64 so that we return the same value as we did before.

return size, dimensions, schema.attr(0).dtype
if source_type == "TILEDB_SPARSE_ARRAY":
schema = tiledb.ArraySchema.load(source_uri)
size = schema.domain.dim(0).domain[1] + 1
dimensions = schema.domain.dim(1).domain[1] + 1
size = np.int64(schema.domain.dim(0).domain[1]) + 1
dimensions = np.int64(schema.domain.dim(1).domain[1]) + 1
return size, dimensions, schema.attr(0).dtype
if source_type == "TILEDB_PARTITIONED_ARRAY":
with tiledb.open(source_uri, "r", config=config) as source_array:
Expand Down Expand Up @@ -2016,7 +2016,7 @@ def consolidate_partition_udf(
prev_index = partial_indexes[0]
i = 0
for partial_index in partial_indexes[1:]:
s = slice(int(prev_index), int(partial_index - 1))
s = slice(int(prev_index), int(partial_index) - 1)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change is also because numpy 2 does not do data type promotion.

With numpy 1 and the original code we were getting:

[ingestion@consolidate_partition_udf] partial_indexes [  0   0   1   1   5   6  14  14  14  18  23  28  42  44  44  44  50  55
  58  60  61  61  61  66  66  81  81  84  84  86  91  94 101 108 110 114
 118 125 126 127 127 127 135 135 141 142 143 143 149 154 157 161 178 191
 200 201 209 214 214 230 233 236 240 242 243 248 257 257 275 276 278 282
 283 290 290 291 298 316 324 324 332 335 335 343 343 347 350 353 356 373
 374 379 382 391 391 391 398 399 405 412 421] <class 'numpy.ndarray'> uint64
[ingestion@consolidate_partition_udf] prev_index 0 <class 'numpy.uint64'>
[ingestion@consolidate_partition_udf] partial_index 0 <class 'numpy.uint64'>
[ingestion@consolidate_partition_udf] s slice(0, -1, None) <class 'slice'>

With numpy 2 and the original code we were getting:

[ingestion@consolidate_partition_udf] partial_indexes [  0   0   1   1   5   6  14  14  14  18  23  28  42  44  44  44  50  55
  58  60  61  61  61  66  66  81  81  84  84  86  91  94 101 108 110 114
 118 125 126 127 127 127 135 135 141 142 143 143 149 154 157 161 178 191
 200 201 209 214 214 230 233 236 240 242 243 248 257 257 275 276 278 282
 283 290 290 291 298 316 324 324 332 335 335 343 343 347 350 353 356 373
 374 379 382 391 391 391 398 399 405 412 421] <class 'numpy.ndarray'> uint64
[ingestion@consolidate_partition_udf] prev_index 0 <class 'numpy.uint64'>
[ingestion@consolidate_partition_udf] partial_index 0 <class 'numpy.uint64'>
[ingestion@consolidate_partition_udf] s slice(0, 18446744073709551615, None) <class 'slice'>

Notice that we get slice(0, 18446744073709551615, None) instead of slice(0, -1, None). To fix this we can cast before subtracting, which we do here.

if (
s.start <= s.stop
and s.start != np.iinfo(np.dtype("uint64")).max
Expand Down
15 changes: 12 additions & 3 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,14 @@ classifiers = [
"Programming Language :: Python :: 3.11",
]

# These are the runtime depdendencies.
dependencies = [
"tiledb-cloud>=0.11",
"tiledb-cloud>=0.12.15",
"tiledb>=0.30.2",
"typing-extensions", # for tiledb-cloud indirect, x-ref https://github.com/TileDB-Inc/TileDB-Cloud-Py/pull/428
# scikit-learn>=1.4.2 may be needed?
"scikit-learn",
"numpy<2.0.0",
"numpy>=1.25.0",
]

[project.optional-dependencies]
Expand All @@ -34,8 +36,10 @@ benchmarks = ["boto3", "paramiko"]
homepage = "https://tiledb.com"
repository = "https://github.com/TileDB-Inc/tiledb-vector-search"

# These are the build-time depdendencies.
[build-system]
requires = ["scikit-build-core[pyproject]", "pybind11", "setuptools-scm"]
# pybind11>=2.12 may be needed?
requires = ["scikit-build-core[pyproject]", "pybind11", "setuptools-scm", "numpy>=2.0.0"]
build-backend = "scikit_build_core.build"

[tool.scikit-build]
Expand Down Expand Up @@ -65,10 +69,15 @@ TILEDB_PATH = {env="TILEDB_PATH"}
[tool.setuptools_scm]
version_file = "apis/python/src/tiledb/vector_search/version.py"

[tool.ruff.lint]
select = ["NPY201"]

[tool.ruff]
extend-select = ["I"]
ignore = ["F403", "F405", "E501", "E741"]
exclude = [".ipynb"]
# Numpy 2 rule: https://numpy.org/devdocs/numpy_2_0_migration_guide.html#ruff-plugin
# select = ["NPY201"]

[tool.ruff.isort]
known-first-party = ["tiledb"]
Expand Down
Loading