diff --git a/.github/workflows/test-python-bindings.yml b/.github/workflows/test-python-bindings.yml index cd60f6ea..f747756c 100644 --- a/.github/workflows/test-python-bindings.yml +++ b/.github/workflows/test-python-bindings.yml @@ -369,7 +369,7 @@ jobs: # to_dataframe() tests skipped themselves in this job while it # reported green. The guard step below now fails the build if any # test skips for a missing import, so the next drift is loud. - uv pip install --system test-install/*.whl pytest pytest-cov requests numpy "psycopg[binary]" redis neo4j pyarrow pandas "adbc-driver-postgresql>=1.12" + uv pip install --system test-install/*.whl pytest pytest-cov requests numpy "psycopg[binary]" redis neo4j pyarrow pandas polars "adbc-driver-postgresql>=1.12" - name: Run pytest on host id: pytest diff --git a/.github/workflows/test-python-examples.yml b/.github/workflows/test-python-examples.yml index 8868aa08..7b5ed5b7 100644 --- a/.github/workflows/test-python-examples.yml +++ b/.github/workflows/test-python-examples.yml @@ -274,11 +274,14 @@ jobs: if [[ "$OSTYPE" == "darwin"* ]]; then # macOS: Install CPU-only PyTorch (supports MPS acceleration) uv pip install --system torch --index-url https://download.pytorch.org/whl/cpu - uv pip install --system "numpy<2.0" requests sentence-transformers + # numpy is not capped: the old macOS `numpy<2.0` was for Intel Macs, where PyTorch stops at 2.2 (built against + # NumPy 1.x). That leg (darwin/amd64) is disabled; re-enabling it needs the cap back for that platform only. + # On Apple Silicon the cap breaks pyarrow, which requires NumPy 2 (example 27). + uv pip install --system numpy requests sentence-transformers pandas polars pyarrow else # Linux & Windows: Install CPU-only PyTorch to save space (avoid CUDA) uv pip install --system torch --index-url https://download.pytorch.org/whl/cpu - uv pip install --system numpy requests sentence-transformers + uv pip install --system numpy requests sentence-transformers pandas polars pyarrow fi # Restored before the dataset step, which is what triggers the model diff --git a/examples/27_dataframe_arrow_ingest.py b/examples/27_dataframe_arrow_ingest.py new file mode 100644 index 00000000..3287699c --- /dev/null +++ b/examples/27_dataframe_arrow_ingest.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +"""Example 27: DataFrames and Arrow tables in, Polars out. + +A pandas or Polars DataFrame, or an Arrow table, loads straight into a type with +Database.insert_dataframe / Database.insert_arrow, the way DuckDB scans a DataFrame +in place: no loop over rows and no CSV in between. A result set comes back out as a +Polars DataFrame with ResultSet.to_polars(). + +Workflow covered: +- a pandas DataFrame into a DOCUMENT type (timestamps, nulls, a list column) +- a Polars DataFrame into a VERTEX type, then edges between those vertices +- an Arrow table into a type that declares some of its properties, so the int64 column + converts to the declared INTEGER and DATETIME_MICROS properties +- a column that would lose data in a declared property (a float into an INTEGER) + refused before a row is written +- the same data out again through to_polars() + +Requirements: +- pandas, polars and pyarrow + (pip install 'arcadedb-embedded[pandas,polars]') +""" + +from __future__ import annotations + +import argparse +import shutil +from datetime import datetime +from pathlib import Path + +import pandas as pd +import polars as pl +import pyarrow as pa + +import arcadedb_embedded as arcadedb + + +def documents_from_pandas(db, rows: int) -> None: + print(f"\n=== pandas DataFrame into a document type: {rows:,} rows ===") + db.command("sql", "CREATE DOCUMENT TYPE Trip") + df = pd.DataFrame( + { + "id": range(rows), + # NaN is pandas' missing value: it loads as a null + "fare": [None if i % 10 == 0 else 2.5 + i % 40 for i in range(rows)], + "picked_up": pd.Timestamp("2026-01-01") + + pd.to_timedelta(range(rows), "min"), + "stops": [[i % 7, i % 11] for i in range(rows)], + } + ) + written = db.insert_dataframe("Trip", df, batch_size=10_000) + stats = db.query( + "sql", "SELECT count(*) AS n, count(fare) AS with_fare FROM Trip" + ).to_list()[0] + print( + f"written={written:,} stored={stats['n']:,} with a fare={stats['with_fare']:,}" + ) + + +def vertices_from_polars(db) -> None: + print("\n=== Polars DataFrame into a vertex type, then edges ===") + db.command("sql", "CREATE VERTEX TYPE Station") + db.command("sql", "CREATE EDGE TYPE Line") + stations = pl.DataFrame( + {"name": ["Central", "North", "Harbour"], "opened": [1998, 2004, 2011]} + ) + print(f"stations written={db.insert_dataframe('Station', stations)}") + rids = { + row["name"]: str(row["rid"]) + for row in db.query("sql", "SELECT name, @rid AS rid FROM Station").to_list() + } + # An edge needs the RIDs of both ends, so edges go through graph_batch().new_edges() + with db.graph_batch() as batch: + batch.new_edges( + [rids["Central"], rids["North"]], "Line", [rids["North"], rids["Harbour"]] + ) + print("edges:", db.query("sql", "SELECT count(*) AS n FROM Line").to_list()[0]["n"]) + + +def arrow_into_declared_types(db) -> None: + print("\n=== Arrow table into declared properties ===") + db.command("sql", "CREATE DOCUMENT TYPE Reading") + db.command("sql", "CREATE PROPERTY Reading.sensor INTEGER") + db.command("sql", "CREATE PROPERTY Reading.at DATETIME_MICROS") + table = pa.table( + { + "sensor": pa.array([1, 2, 3], pa.int64()), # converts to INTEGER + "at": pa.array( + [datetime(2026, 5, 1, 12, 0, 0, 123456)] * 3, pa.timestamp("us") + ), + "value": pa.array([0.5, None, 2.5], pa.float32()), # undeclared: DOUBLE + } + ) + print(f"written={db.insert_arrow('Reading', table)}") + try: + db.insert_arrow("Reading", pa.table({"sensor": [1.5]})) + except TypeError as exc: + print(f"refused before anything was written: {exc}") + + +def polars_out(db) -> None: + print("\n=== to_polars() ===") + frame = db.query( + "sql", "SELECT sensor, at, value FROM Reading ORDER BY sensor" + ).to_polars() + print(frame) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--rows", type=int, default=100_000) + parser.add_argument("--db-path", default="./my_test_databases/dataframe_ingest_db") + args = parser.parse_args() + + path = Path(args.db_path) + shutil.rmtree(path, ignore_errors=True) + path.parent.mkdir(parents=True, exist_ok=True) + with arcadedb.create_database(str(path)) as db: + documents_from_pandas(db, args.rows) + vertices_from_polars(db) + arrow_into_declared_types(db) + polars_out(db) + + +if __name__ == "__main__": + main() diff --git a/examples/README.md b/examples/README.md index e1637eaf..e180531b 100644 --- a/examples/README.md +++ b/examples/README.md @@ -391,6 +391,23 @@ half-applied after an interruption --- +### 🐼 [27_dataframe_arrow_ingest.py](./27_dataframe_arrow_ingest.py) +**DataFrame And Arrow Ingest | pandas | Polars | Arrow | Declared Types** + +DataFrames and Arrow tables loaded straight into a type, and a result back out as Polars: +- a pandas DataFrame into a document type (timestamps, nulls, a list column), in batches +- a Polars DataFrame into a vertex type, then edges through `graph_batch().new_edges()` +- an Arrow table into declared `INTEGER` and `DATETIME_MICROS` properties, and a float + column refused before a row is written +- `to_polars()` on the result +- synthetic data; flags: `--rows` (default 100,000) and `--db-path` +- requires `pip install 'arcadedb-embedded[pandas,polars]'` + +**Learn:** How a DataFrame becomes documents or vertices without a row loop, and what each +Arrow type is stored as + +--- + ### 🧠 [29_hybrid_memory_retrieval.py](./29_hybrid_memory_retrieval.py) **Agent Memory | Tenant Pre-filter | Vector + Full-Text + Graph | `vector.fuse`** diff --git a/pyproject.toml b/pyproject.toml index 37269b7b..edcd6f9d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,7 @@ test = [ "redis>=5.0", # Redis wire "neo4j>=5.0", # Bolt "pyarrow>=23.0.1", # ResultSet.to_arrow(); <23.0.1 has PYSEC-2026-113 + "polars>=1.0", # ResultSet.to_polars(), insert_dataframe() with a Polars frame # ResultSet.to_dataframe(). Absent until 2026-08-04, which meant its two # tests called importorskip("pandas") and skipped in every CI run this # package has ever had: a public, documented method with no executed @@ -100,6 +101,15 @@ pandas = [ "pandas>=2.2.0", "numpy>=1.26.4", ] +# Optional: ResultSet.to_polars() and insert_dataframe() with a Polars frame. Both go +# through Arrow, so pyarrow comes with it. insert_dataframe() with a pandas frame needs +# `pip install 'arcadedb-embedded[arrow,pandas]'` (the pandas extra alone has no pyarrow, +# because to_dataframe() does not need it). +polars = [ + "polars>=1.0", + "pyarrow>=23.0.1", + "numpy>=1.26.4", +] examples = [ "py7zr>=1.1.3", # Required for Stack Exchange data downloads ] diff --git a/src/arcadedb_embedded/core.py b/src/arcadedb_embedded/core.py index 2a85037a..429c79bf 100644 --- a/src/arcadedb_embedded/core.py +++ b/src/arcadedb_embedded/core.py @@ -996,6 +996,174 @@ def insert_columns( ) return n + def insert_dataframe( + self, + type_name: str, + df, + *, + batch_size: int = 10_000, + parallel: bool = False, + ) -> int: + """Load a pandas or Polars DataFrame into a type, in batches. + + The DataFrame is read as an Arrow table (see :meth:`insert_arrow`, which + takes that table directly and has the type mapping, the declared-type + rules, and the failure contract). Each column becomes a typed array + that crosses the bridge once per batch: integers, floats, and booleans + without nulls as one buffer copy, anything else per element. A document + type gets documents; a vertex type gets vertices. An edge type is + refused: load the vertices, then buffer the edges with + ``graph_batch().new_edges()``. + + What differs by library: a pandas ``NaN`` in a float column is pandas' + missing value and loads as a null, while a Polars ``NaN`` stays a + ``NaN`` (Polars keeps null and NaN apart). The pandas index is not + loaded (``df.reset_index()`` makes it a column), and pandas column + names must be strings. A Polars ``LazyFrame`` must be ``.collect()``-ed + first. + + Args: + type_name: Target type (must exist, with or without declared + properties). + df: A pandas or Polars ``DataFrame``. Any object that speaks the + Arrow PyCapsule interface also works. + batch_size: Rows per batch. Each batch is converted, then written + and committed on its own, so memory stays bounded for a large + frame. + parallel: Document types only: hand each batch to the async + executor's parallel bucket writers (see :meth:`insert_columns`). + + Returns: + Number of records written. + + Raises: + ImportError: If pyarrow is not installed. + ValueError: If the frame has no columns, duplicate or non-string + column names, ``batch_size`` is not a positive integer, or + ``parallel`` is used with a vertex type. + TypeError: If a column has a type with no ArcadeDB mapping (time of + day, duration, map), or one that does not load into the property + type ``type_name`` declares for it (a float column into an + INTEGER property would truncate). Raised before anything is + written. + ArcadeDBError: If the type does not exist or is an edge type, or the + load fails (a duplicate key, a value the declared type refuses). + The batch in flight is rolled back; batches before it stay + committed, and the message says how many rows that is. A + transaction the caller opened is left to the caller (for + documents; a vertex load refuses to run inside one): the + batches join it, the message then says the rows were written + in your transaction, not committed, and a rollback removes them. + + Example: + >>> import pandas as pd + >>> df = pd.DataFrame({"id": [1, 2], "name": ["a", "b"]}) + >>> db.insert_dataframe("User", df) + 2 + """ + from .dataframe import insert_table, to_arrow_table + + self._check_not_closed() + table = to_arrow_table(df, what="insert_dataframe()") + return insert_table( + self, + type_name, + table, + batch_size=batch_size, + parallel=parallel, + what="insert_dataframe()", + ) + + def insert_arrow( + self, + type_name: str, + table, + *, + batch_size: int = 10_000, + parallel: bool = False, + ) -> int: + """Load an Arrow table into a type, in batches. + + The Arrow counterpart of :meth:`insert_columns`: each Arrow column is + converted to the typed array ``insert_columns`` sends across the bridge + (``DocumentBatcher.insertColumns``), so a document type loads at that + path's speed; a vertex type loads through ``GraphBatch``. Documents and + vertices only: an edge type is refused. + + A property the type does not declare takes the stored type below; a + declared property converts to its declared type. + + | Arrow type | Stored as | + |---|---| + | int8 to int64, uint8 to uint32 | LONG (uint64 up to 2**63 - 1) | + | float16, float32, float64 | DOUBLE (float32 widens exactly) | + | bool | BOOLEAN | + | string, large_string, string_view | STRING | + | dictionary (categorical) | its value type | + | binary, large_binary | byte[] | + | timestamp without a time zone, any unit | a UTC date-time, the wall clock as it shows | + | timestamp with a time zone, any unit | a UTC date-time, converted to UTC | + | date32, date64 | DATE | + | decimal128, decimal256 | DECIMAL | + | list, large_list, fixed_size_list | LIST (ARRAY_OF_* when declared so) | + | struct | MAP | + | null (only nulls) | no value | + + A date-time is a UTC wall clock. A declared property keeps its own + precision: DATETIME milliseconds, DATETIME_MICROS microseconds, + DATETIME_NANOS nanoseconds, DATETIME_SECOND seconds; a column the type + does not declare keeps the input's precision. Nanosecond input is read + from the integer epoch value, so it is not cut to microseconds on the + way in. Reading back is coarser: ``Result.get()`` returns a Python + ``datetime`` (microseconds at most) and ``to_arrow()``, ``to_polars()``, + ``to_dataframe()`` and ``to_columns()`` return milliseconds. Time-of-day, duration, and map + columns have no mapping and raise ``TypeError``. A null is stored as a + missing property. + + A column loads into a declared property when the engine converts it + without guessing: integers into BYTE, SHORT, INTEGER, LONG, FLOAT, + DOUBLE, DECIMAL, or STRING (a value past the target's range fails the + load); floats into FLOAT, DOUBLE, DECIMAL, or STRING (never into an + integer property, which would truncate); booleans into BOOLEAN or + STRING; timestamps into any DATETIME precision or STRING; dates into DATE, + any DATETIME, or STRING; strings into STRING or LINK (RIDs). Anything else + raises ``TypeError`` before a row is written. + + Args: + type_name: Target type (must exist). + table: A ``pyarrow.Table``, ``RecordBatch``, or ``RecordBatchReader`` + (read whole), or any object that speaks the Arrow PyCapsule + interface. + batch_size: Rows per batch; each batch is written and committed on + its own. + parallel: Document types only: as in :meth:`insert_columns`. + + Returns: + Number of records written. + + Raises: + ImportError: If pyarrow is not installed. + ValueError, TypeError, ArcadeDBError: As for :meth:`insert_dataframe`. + + Example: + >>> import pyarrow as pa + >>> table = pa.table({"id": [1, 2], "name": ["a", "b"]}) + >>> db.insert_arrow("User", table) + 2 + """ + from .dataframe import insert_table, to_arrow_table + + self._check_not_closed() + arrow_table = to_arrow_table(table, what="insert_arrow()") + return insert_table( + self, + type_name, + arrow_table, + batch_size=batch_size, + parallel=parallel, + what="insert_arrow()", + ) + def close(self): """Close the database.""" if not self._closed and self._java_db is not None: @@ -1166,6 +1334,8 @@ def create_vector_index( pq_clusters: Optional[int] = None, pq_center_globally: Optional[bool] = None, pq_training_limit: Optional[int] = None, + rerank_oversample: Optional[int] = None, + fused_pq: Optional[bool] = None, build_graph_now: bool = True, ) -> "VectorIndex": """ @@ -1199,6 +1369,7 @@ def create_vector_index( In current ArcadeDB engine builds, PRODUCT also requires enough indexed vectors per bucket for PQ training; for tiny corpora, set `pq_clusters` explicitly to a small value or prefer INT8/BINARY/NONE. + PRODUCT is not recommended; see ArcadeDB #10514. encoding: Optional storage encoding for the underlying vector property. Use "INT8" when the document property stores pre-quantized bytes in a `BINARY` property. When using INT8 encoding, set quantization to "NONE" @@ -1230,7 +1401,17 @@ def create_vector_index( pq_center_globally: Whether to globally center vectors before PQ. Requires quantization="PRODUCT". pq_training_limit: Max vectors to use for PQ training. Requires - quantization="PRODUCT". + quantization="PRODUCT". The codebook is trained on a random sample + of this size (the engine caps it at 128000). + rerank_oversample: Candidates per requested neighbour that a search + walking on the PQ codes scores again on the stored vectors + (maps to the index setting "rerankOversample"; the engine default + is 10, 0 returns the PQ scores alone, negative is refused). + Requires quantization="PRODUCT". A query can override it with the + `rerank` option of `vectorNeighbors`. + fused_pq: Whether the graph file stores the PQ codes of every node's + neighbours inline (jvector FusedPQ; maps to the index setting + "fusedPQ"; off by default). Requires quantization="PRODUCT". build_graph_now: If True (default), eagerly builds the vector graph immediately after index creation. If False, graph preparation is deferred and may happen lazily on first search. @@ -1240,7 +1421,10 @@ def create_vector_index( lookups. add_hierarchy: Whether to build hierarchical layers in the HNSW graph (Default is True). If None, uses the engine default. Set explicitly to - True/False to force the behavior. + True/False to force the behavior. The default of True is a deliberate + choice of these bindings: the engine itself defaults to False, so a + SQL ``CREATE INDEX ... LSM_VECTOR`` without ``"addHierarchy"`` builds a + flat graph. Record the value you use when you compare results. Returns: VectorIndex object @@ -1274,6 +1458,8 @@ def create_vector_index( pq_clusters, pq_center_globally, pq_training_limit, + rerank_oversample, + fused_pq, ) ): if not quantization or quantization.upper() != "PRODUCT": @@ -1334,6 +1520,10 @@ def create_vector_index( builder.withPQTrainingLimit(int(pq_training_limit)) metadata_cfg = {} + if rerank_oversample is not None: + metadata_cfg["rerankOversample"] = int(rerank_oversample) + if fused_pq is not None: + metadata_cfg["fusedPQ"] = bool(fused_pq) if store_vectors_in_graph: metadata_cfg["storeVectorsInGraph"] = True if add_hierarchy is not None: diff --git a/src/arcadedb_embedded/dataframe.py b/src/arcadedb_embedded/dataframe.py new file mode 100644 index 00000000..270f397c --- /dev/null +++ b/src/arcadedb_embedded/dataframe.py @@ -0,0 +1,342 @@ +""" +ArcadeDB Python Bindings - DataFrame and Arrow ingest + +Loads a pandas or Polars DataFrame, or an Arrow table, into a type in batches +(``Database.insert_dataframe`` and ``Database.insert_arrow``). Arrow is the one +internal form: a pandas frame becomes an Arrow table, a Polars frame is already +one, and each Arrow column becomes the typed array ``Database.insert_columns`` +sends across the bridge (a document type), or the property rows +``GraphBatch.create_vertices`` takes (a vertex type). No Java was added: both +paths use the existing glue (``DocumentBatcher.insertColumns`` and the +``GraphBatch`` vertex batcher). +""" + +from typing import Any, Dict, List, Optional, Tuple + +import jpype + +from .exceptions import ArcadeDBError + +_PYARROW_HINT = ( + "pyarrow is required for {what}. Install with: pip install pyarrow " + "(or pip install 'arcadedb-embedded[arrow]')" +) + +# Declared property types (the engine's Type names) that an Arrow column kind +# loads into. The engine converts the value (an integer into a LONG or a +# DECIMAL, a list into an ARRAY_OF_FLOATS) and refuses what does not fit (a +# LONG past INTEGER's range). What it also does, silently, is truncate a float +# into an integer property and parse a string into a number, so the kinds that +# would lose or guess are refused here, before anything is written. +_DATETIMES = {"DATETIME", "DATETIME_SECOND", "DATETIME_MICROS", "DATETIME_NANOS"} +_ARRAYS = { + "ARRAY_OF_SHORTS", + "ARRAY_OF_INTEGERS", + "ARRAY_OF_LONGS", + "ARRAY_OF_FLOATS", + "ARRAY_OF_DOUBLES", +} +_ALLOWED = { + "integer": { + "BYTE", + "SHORT", + "INTEGER", + "LONG", + "FLOAT", + "DOUBLE", + "DECIMAL", + "STRING", + }, + "floating": {"FLOAT", "DOUBLE", "DECIMAL", "STRING"}, + "bool": {"BOOLEAN", "STRING"}, + "string": {"STRING", "LINK"}, + "binary": {"BINARY"}, + "timestamp": _DATETIMES | {"STRING"}, + "date": {"DATE"} | _DATETIMES | {"STRING"}, + "decimal": {"DECIMAL", "DOUBLE", "FLOAT", "STRING"}, + "list": {"LIST"} | _ARRAYS, + "struct": {"MAP", "EMBEDDED"}, +} + + +def _import_pyarrow(what: str): + try: + import pyarrow as pa + except ImportError as exc: + raise ImportError(_PYARROW_HINT.format(what=what)) from exc + return pa + + +def _kind(pa, arrow_type) -> Optional[str]: + """The mapping family of an Arrow type, or None when it has no ArcadeDB mapping.""" + types = pa.types + if types.is_dictionary(arrow_type): + return _kind(pa, arrow_type.value_type) + if types.is_integer(arrow_type): + return "integer" + if types.is_floating(arrow_type): + return "floating" + if types.is_boolean(arrow_type): + return "bool" + if ( + types.is_string(arrow_type) + or types.is_large_string(arrow_type) + or (hasattr(types, "is_string_view") and types.is_string_view(arrow_type)) + ): + return "string" + if types.is_binary(arrow_type) or types.is_large_binary(arrow_type): + return "binary" + if types.is_timestamp(arrow_type): + return "timestamp" + if types.is_date(arrow_type): + return "date" + if types.is_decimal(arrow_type): + return "decimal" + if ( + types.is_list(arrow_type) + or types.is_large_list(arrow_type) + or types.is_fixed_size_list(arrow_type) + ): + return "list" + if types.is_struct(arrow_type): + return "struct" + if types.is_null(arrow_type): + return "null" + return None + + +def to_arrow_table(data, *, what: str): + """Normalize a pandas or Polars DataFrame, or anything Arrow can read, to a ``pyarrow.Table``.""" + pa = _import_pyarrow(what) + module = type(data).__module__.split(".")[0] + if isinstance(data, pa.Table): + return data + if isinstance(data, pa.RecordBatch): + return pa.Table.from_batches([data]) + if isinstance(data, pa.RecordBatchReader): + return data.read_all() + if module == "pandas" and hasattr(data, "columns"): + names = list(data.columns) + bad = [n for n in names if not isinstance(n, str)] + if bad: + raise ValueError( + f"column names must be strings, got {bad[0]!r}; rename the columns first" + ) + if len(set(names)) != len(names): + raise ValueError("a DataFrame with duplicate column names cannot be loaded") + # NaN in a float column is pandas' missing value: Arrow stores it as a null. + return pa.Table.from_pandas(data, preserve_index=False) + if module == "polars": + if hasattr(data, "collect") and not hasattr(data, "to_arrow"): + raise TypeError( + "a Polars LazyFrame is not loaded as it is; call .collect() first" + ) + return data.to_arrow() + if hasattr(data, "__arrow_c_stream__") or hasattr(data, "__arrow_c_array__"): + return pa.table(data) # any library that speaks the Arrow PyCapsule interface + raise TypeError( + f"{what} takes a pandas or Polars DataFrame or an Arrow table, " + f"not {type(data).__name__}" + ) + + +def _declared_types( + db, type_name: str, names: List[str] +) -> Tuple[bool, Dict[str, str]]: + """``(is_vertex_type, {column: declared property type name})`` for the columns the type declares. + + Raises ArcadeDBError when the type does not exist or is an edge type. + """ + java_schema = db._java_db.getSchema() + if not java_schema.existsType(type_name): + raise ArcadeDBError( + f"Cannot load into '{type_name}': the type does not exist " + f"(create it first, with or without declared properties)" + ) + java_type = java_schema.getType(type_name) + if isinstance(java_type, jpype.JClass("com.arcadedb.schema.EdgeType")): + raise ArcadeDBError( + f"Cannot load into '{type_name}': it is an edge type. Load the vertices, " + f"then buffer the edges with graph_batch().new_edges()" + ) + is_vertex = isinstance(java_type, jpype.JClass("com.arcadedb.schema.VertexType")) + declared = {} + for name in names: + prop = java_type.getPolymorphicPropertyIfExists(name) + if prop is not None: + declared[name] = str(prop.getType().name()) + return is_vertex, declared + + +def _check_declared( + type_name: str, schema, kinds: Dict[str, str], declared: Dict[str, str] +): + """Refuse a column whose Arrow kind does not load into the property type the type declares.""" + for name, declared_type in declared.items(): + kind = kinds[name] + if kind == "null": + continue # nothing but nulls: any property accepts it + if declared_type not in _ALLOWED[kind]: + raise TypeError( + f"column {name!r} has Arrow type {schema.field(name).type} but " + f"'{type_name}.{name}' is declared {declared_type}; cast the column to a " + f"matching type first (e.g. df[{name!r}].astype(...))" + ) + + +def _temporal_to_java(chunk, pa): + """A timestamp or date chunk as a list of ``LocalDateTime``/``LocalDate``, keeping the unit's precision. + + DATETIME is a UTC wall clock: a timestamp with a zone is already UTC epoch time, and a timestamp without + one is taken as the wall clock it shows (the same rule ``Document.set`` applies to a ``datetime``), so both + are read from the integer epoch value. Going through ``datetime`` would cut a nanosecond value to microseconds. + """ + import numpy as np + + types = pa.types + mask = chunk.is_null().to_numpy(zero_copy_only=False) + if types.is_date(chunk.type): + local_date = jpype.JClass("java.time.LocalDate") + raw = chunk.cast(pa.int64() if types.is_date64(chunk.type) else pa.int32()) + days = raw.fill_null(0).to_numpy(zero_copy_only=False).astype(np.int64) + if types.is_date64(chunk.type): + days = days // 86_400_000 + return [ + None if m else local_date.ofEpochDay(int(d)) for d, m in zip(days, mask) + ] + per_second = {"s": 1, "ms": 1_000, "us": 1_000_000, "ns": 1_000_000_000}[ + chunk.type.unit + ] + values = chunk.cast(pa.int64()).fill_null(0).to_numpy(zero_copy_only=False) + seconds, rest = np.divmod(values.astype(np.int64), per_second) + nanos = rest * (1_000_000_000 // per_second) + local_datetime = jpype.JClass("java.time.LocalDateTime") + utc = jpype.JClass("java.time.ZoneOffset").UTC + return [ + None if m else local_datetime.ofEpochSecond(int(s), int(n), utc) + for s, n, m in zip(seconds, nanos, mask) + ] + + +def _column_values(name: str, chunk, pa): + """One Arrow array as what ``Database.insert_columns`` takes: a numpy array when the column crosses as one + buffer (integers, floats and booleans without nulls), else a list of values (``None`` is a null). + """ + types = pa.types + if types.is_dictionary(chunk.type): # a pandas categorical or a Polars Categorical + chunk = chunk.cast(chunk.type.value_type) + kind = _kind(pa, chunk.type) + if kind in ("integer", "floating", "bool"): + if chunk.null_count == 0: + values = chunk.to_numpy(zero_copy_only=False) + if values.dtype.kind == "f" and values.dtype.itemsize < 8: + values = values.astype("float64") # float16/32 widen exactly + return values + return chunk.to_pylist() + if kind in ("timestamp", "date"): + return _temporal_to_java(chunk, pa) + if kind == "null": + return [None] * len(chunk) + # strings, binary, decimal, list, struct: Python values, which insert_columns converts element by element + # (str, bytes -> byte[], Decimal -> BigDecimal, list -> LIST, dict -> MAP) + return chunk.to_pylist() + + +def insert_table( + db, + type_name: str, + table, + *, + batch_size: int, + parallel: bool, + what: str, +) -> int: + """Load a ``pyarrow.Table`` into ``type_name`` in batches of ``batch_size`` rows; see ``Database.insert_arrow``.""" + pa = _import_pyarrow(what) + db._check_not_closed() + if ( + isinstance(batch_size, bool) + or not isinstance(batch_size, int) + or batch_size < 1 + ): + raise ValueError(f"batch_size must be a positive integer, got {batch_size!r}") + names = list(table.column_names) + if not names: + raise ValueError(f"{what} needs at least one column") + if len(set(names)) != len(names): + raise ValueError("a table with duplicate column names cannot be loaded") + + kinds = {} + for field in table.schema: + kind = _kind(pa, field.type) + if kind is None: + raise TypeError( + f"column {field.name!r} has Arrow type {field.type}, which has no ArcadeDB " + f"mapping; convert it first (supported: integers, floats, bool, string, " + f"binary, timestamp, date, decimal, list, struct, null)" + ) + kinds[field.name] = kind + is_vertex, declared = _declared_types(db, type_name, names) + _check_declared(type_name, table.schema, kinds, declared) + if parallel and is_vertex: + raise ValueError( + "parallel=True writes documents through the async executor; " + f"'{type_name}' is a vertex type, which loads through graph_batch()" + ) + if table.num_rows == 0: + return 0 + + written = 0 + # a document load inside the caller's own transaction joins it (insert_columns commits + # nothing then); a vertex load refuses to run inside one + in_caller_transaction = (not is_vertex) and bool(db._java_db.isTransactionActive()) + graph_batch = db.graph_batch() if is_vertex else None + try: + for batch in table.to_batches(max_chunksize=batch_size): + if batch.num_rows == 0: + continue + columns = { + name: _column_values(name, batch.column(i), pa) + for i, name in enumerate(names) + } + try: + if is_vertex: + _insert_vertices( + graph_batch, type_name, names, columns, batch.num_rows + ) + else: + db.insert_columns( + type_name, columns, commit_every=batch_size, parallel=parallel + ) + except ArcadeDBError as exc: + if written and in_caller_transaction: + # the batches joined the caller's transaction: nothing is committed + # until the caller commits, and a rollback removes them + raise ArcadeDBError( + f"{exc}; the {written} row(s) of the batches before this one " + f"were written in your transaction, not committed" + ) from exc + if written: + raise ArcadeDBError( + f"{exc}; the {written} row(s) of the batches before this one " + f"were committed" + ) from exc + raise + written += batch.num_rows + finally: + if graph_batch is not None: + graph_batch.close() + return written + + +def _insert_vertices( + graph_batch, type_name: str, names: List[str], columns: Dict[str, Any], n: int +): + """One batch of vertices: the columns as one property dict per row, through ``GraphBatch.create_vertices``.""" + lists = [] + for name in names: + values = columns[name] + lists.append(values.tolist() if hasattr(values, "tolist") else values) + graph_batch.create_vertices( + type_name, [dict(zip(names, row)) for row in zip(*lists)] + ) diff --git a/src/arcadedb_embedded/results.py b/src/arcadedb_embedded/results.py index a8f78f8f..dbe948c0 100644 --- a/src/arcadedb_embedded/results.py +++ b/src/arcadedb_embedded/results.py @@ -645,6 +645,48 @@ def to_dataframe(self, convert_types: bool = True): return pd.DataFrame(self.to_list(convert_types=convert_types)) + def to_polars( + self, batch_size: int = 25_000, columns: Optional[Sequence[str]] = None + ): + """ + Bulk-materialize all rows as a ``polars.DataFrame``, through Arrow. + + Reads :meth:`to_arrow` and hands the table to Polars, so nulls keep + their type (a nullable integer stays an integer column) and a DECIMAL + column is a Polars ``Decimal``. Requires polars and pyarrow + (``pip install 'arcadedb-embedded[polars]'``). + + Args: + batch_size: Rows per bridge crossing, as in :meth:`to_arrow`. + columns: Optional column names to pin, as in :meth:`to_columns`. + + Returns: + polars DataFrame + + Raises: + ImportError: If polars or pyarrow is not installed, or the bundled + bridge jar is missing (the cases in which :meth:`to_arrow` + returns None). + + Example: + >>> df = db.query("sql", "SELECT FROM User").to_polars() + """ + try: + import polars as pl + except ImportError as exc: + raise ImportError( + "polars is required for to_polars(). " + "Install with: pip install 'arcadedb-embedded[polars]'" + ) from exc + table = self.to_arrow(batch_size=batch_size, columns=columns) + if table is None: + raise ImportError( + "to_polars() reads through Arrow, which needs pyarrow and numpy " + "and the bundled bridge jar. " + "Install with: pip install 'arcadedb-embedded[polars]'" + ) + return pl.from_arrow(table) + def to_columns( self, batch_size: int = 25_000, columns: Optional[Sequence[str]] = None ): diff --git a/src/arcadedb_embedded/schema.py b/src/arcadedb_embedded/schema.py index fa54d21b..88b7de52 100644 --- a/src/arcadedb_embedded/schema.py +++ b/src/arcadedb_embedded/schema.py @@ -28,6 +28,7 @@ class IndexType(Enum): LSM_TREE = "LSM_TREE" FULL_TEXT = "FULL_TEXT" LSM_VECTOR = "LSM_VECTOR" + LSM_SPARSE_VECTOR = "LSM_SPARSE_VECTOR" GEOSPATIAL = "GEOSPATIAL" HASH = "HASH" @@ -461,7 +462,8 @@ def create_index( type_name: Name of the type property_names: List of property names to index unique: Whether the index should enforce uniqueness - index_type: Type of index (LSM_TREE, FULL_TEXT, LSM_VECTOR, GEOSPATIAL, HASH) + index_type: Type of index (LSM_TREE, FULL_TEXT, LSM_VECTOR, + LSM_SPARSE_VECTOR, GEOSPATIAL, HASH) Returns: Java Index object @@ -536,7 +538,8 @@ def get_or_create_index( type_name: Name of the type property_names: List of property names to index unique: Whether the index should enforce uniqueness - index_type: Type of index (LSM_TREE, FULL_TEXT, LSM_VECTOR, GEOSPATIAL, HASH) + index_type: Type of index (LSM_TREE, FULL_TEXT, LSM_VECTOR, + LSM_SPARSE_VECTOR, GEOSPATIAL, HASH) Returns: Java Index object diff --git a/src/arcadedb_embedded/vector.py b/src/arcadedb_embedded/vector.py index 717553f6..468f5d16 100644 --- a/src/arcadedb_embedded/vector.py +++ b/src/arcadedb_embedded/vector.py @@ -447,6 +447,26 @@ def _find_neighbor_pairs( ) return idx.findNeighborsFromVectorApproximate(java_vector, k) + # An engine without ArcadeData/arcadedb#10545 (ArcadeDB before 26.11.1) has no + # findNeighbors and answers SQL with the exact walk too: keep that walk there. + if str(idx.getMetadata().quantizationType) == "PRODUCT" and hasattr( + idx, "findNeighbors" + ): + # The entry point SQL vectorNeighbors, Cypher and the select share: a + # PRODUCT index walks on its PQ codes and re-ranks the best + # k * rerankOversample candidates on the stored vectors, so this + # helper answers like the statement does (findNeighborsFromVector is + # the exact walk). -1 for the beam and the re-rank: index defaults. + guard = jpype.JClass("com.arcadedb.query.sql.executor.WorkGuard") + return idx.findNeighbors( + java_vector, + k, + -1 if ef_search is None else int(ef_search), + -1, + allowed_rids_set, + guard.forCommandDeadline(None), + ) + if allowed_rids_set is not None and ef_search is not None: return idx.findNeighborsFromVector( java_vector, k, ef_search, allowed_rids_set @@ -518,6 +538,13 @@ def find_nearest( """ Find k nearest neighbors to the query vector. + On a ``PRODUCT`` index this takes the same engine path as SQL + ``vectorNeighbors``: the walk scores on the in-memory PQ codes and the best + ``k * rerankOversample`` candidates are re-ranked on the stored vectors, so + the distances returned are exact and the ranking can differ from an exact + walk's. Every other quantization takes the exact walk. For the per-query + ``rerank`` option, use ``vectorNeighbors`` in SQL. + Args: query_vector: Query vector as Python list, NumPy array, or array-like k: Number of nearest neighbors to return (final k). Default is 10. @@ -568,12 +595,14 @@ def get_metadata(self): Returns: dict: Index metadata including names, dimensions, similarity function, id property, quantization, and selected tuning settings. + ``rerank_oversample`` and ``fused_pq`` are present when the engine has + them (ArcadeDB 26.11.1 or later). """ try: index = self._get_primary_lsm_index() meta = self._get_primary_metadata() - return { + info = { "index_name": str(self._java_index.getName()), "bucket_index_name": str(index.getName()), "type_name": str(index.getTypeName()), @@ -598,6 +627,11 @@ def get_metadata(self): "pq_center_globally": bool(meta.pqCenterGlobally), "pq_training_limit": int(meta.pqTrainingLimit), } + # ArcadeDB before 26.11.1 (ArcadeData/arcadedb#10545) has neither setting + if hasattr(meta, "rerankOversample"): + info["rerank_oversample"] = int(meta.rerankOversample) + info["fused_pq"] = bool(meta.fusedPQ) + return info except Exception as e: raise ArcadeDBError(f"Failed to read vector index metadata: {e}") from e @@ -687,6 +721,11 @@ def find_nearest_approximate( """ Find k nearest neighbors using Product Quantization (PQ) approximate search. + The answer is the PQ scores alone: no re-rank on the stored vectors and no + beam, so recall stops improving at about 0.72 on SIFT-1M with 32 subspaces + and the distances are PQ estimates. ``find_nearest`` re-ranks and returns + exact distances. + Args: query_vector: Query vector as Python list, NumPy array, or array-like k: Number of nearest neighbors to return (final k) diff --git a/tests/test_dataframe_ingest.py b/tests/test_dataframe_ingest.py new file mode 100644 index 00000000..a0b4bd3c --- /dev/null +++ b/tests/test_dataframe_ingest.py @@ -0,0 +1,731 @@ +"""Tests for Database.insert_dataframe, Database.insert_arrow and ResultSet.to_polars (humemai/arcadedb-embedded-python#387). + +A pandas or Polars DataFrame, or an Arrow table, loads into a document or vertex type in batches. These tests +hold what a user relies on: each Arrow type stores as the documented ArcadeDB type and reads back equal, nulls +are missing properties, a declared property converts (or the load is refused before a row is written), a frame +larger than one batch loads whole, and a result set goes back out as a Polars frame. +""" + +import datetime as dt +import decimal +import sys + +import pytest + +np = pytest.importorskip("numpy") +pd = pytest.importorskip("pandas") +pa = pytest.importorskip("pyarrow") +pl = pytest.importorskip("polars") + +from arcadedb_embedded.exceptions import ArcadeDBError # noqa: E402 + + +def _rows(db, sql): + return db.query("sql", sql).to_list() + + +def _count(db, type_name, where=""): + q = f"SELECT count(*) AS n FROM {type_name} {where}" # nosec B608 - test-controlled + return int(_rows(db, q)[0]["n"]) + + +def _one_column(db, arrow_type, values, declared=None): + """Load ``values`` as one column ``val`` of ``arrow_type`` and return what reads back, in row order.""" + db.command("sql", "CREATE DOCUMENT TYPE One") + if declared: + db.command("sql", f"CREATE PROPERTY One.val {declared}") + table = pa.table( + { + "idx": pa.array(range(len(values)), pa.int64()), + "val": pa.array(values, arrow_type), + } + ) + assert db.insert_arrow("One", table) == len(values) + return [r.get("val") for r in _rows(db, "SELECT FROM One ORDER BY idx")] + + +class TestTypeMapping: + """Every Arrow type of the documented table, with no declared property: write, read back, compare.""" + + @pytest.mark.parametrize( + "arrow_type", + [ + pa.int8(), + pa.int16(), + pa.int32(), + pa.int64(), + pa.uint8(), + pa.uint16(), + pa.uint32(), + ], + ) + def test_integers_store_as_long(self, temp_db, arrow_type): + got = _one_column(temp_db, arrow_type, [0, 1, 100]) + assert got == [0, 1, 100] and all(type(v) is int for v in got) + + def test_signed_integer_extremes(self, temp_db): + values = [-(2**63), -1, 2**63 - 1] + assert _one_column(temp_db, pa.int64(), values) == values + + @pytest.mark.parametrize("arrow_type", [pa.float32(), pa.float64(), pa.float16()]) + def test_floats_store_as_double(self, temp_db, arrow_type): + # 0.5, -2.25 and 1024.0 are exact in every width, so the read-back is equal and not approximate + values = [0.5, -2.25, 1024.0] + got = _one_column( + temp_db, arrow_type, np.array(values, dtype=arrow_type.to_pandas_dtype()) + ) + assert got == values and all(type(v) is float for v in got) + + def test_bool(self, temp_db): + assert _one_column(temp_db, pa.bool_(), [True, False, True]) == [ + True, + False, + True, + ] + + @pytest.mark.parametrize("arrow_type", [pa.string(), pa.large_string()]) + def test_string(self, temp_db, arrow_type): + values = ["a", "", "naïve ✓"] + assert _one_column(temp_db, arrow_type, values) == values + + def test_dictionary_column_stores_its_values(self, temp_db): + # a pandas categorical and a Polars Categorical both arrive as an Arrow dictionary + temp_db.command("sql", "CREATE DOCUMENT TYPE Cat") + df = pd.DataFrame({"c": pd.Categorical(["x", "y", "x"])}) + assert temp_db.insert_dataframe("Cat", df) == 3 + assert sorted(r["c"] for r in _rows(temp_db, "SELECT FROM Cat")) == [ + "x", + "x", + "y", + ] + + def test_timestamp_without_zone_keeps_the_wall_clock(self, temp_db): + values = [ + dt.datetime(2020, 1, 2, 3, 4, 5, 123000), + dt.datetime(1969, 12, 31, 23, 59, 59, 500000), + ] + for unit in ("s", "ms", "us", "ns"): + temp_db.command("sql", "DROP TYPE One IF EXISTS UNSAFE") + expected = ( + [v.replace(microsecond=0) for v in values] if unit == "s" else values + ) + assert _one_column(temp_db, pa.timestamp(unit), expected) == expected + + def test_timestamp_with_zone_is_stored_as_utc(self, temp_db): + local = dt.datetime( + 2020, 1, 2, 3, 4, 5, 123000, tzinfo=dt.timezone(dt.timedelta(hours=5)) + ) + got = _one_column(temp_db, pa.timestamp("us", tz="Asia/Karachi"), [local]) + assert got == [ + dt.datetime(2020, 1, 1, 22, 4, 5, 123000) + ] # UTC wall clock, naive on the way out + + def test_date(self, temp_db): + values = [dt.date(2020, 2, 29), dt.date(1960, 1, 1)] + assert _one_column(temp_db, pa.date32(), values) == values + + def test_date64(self, temp_db): + values = [dt.date(2020, 2, 29)] + assert _one_column(temp_db, pa.date64(), values) == values + + def test_decimal_keeps_every_digit(self, temp_db): + values = [ + decimal.Decimal("12345678901234567890.123456"), + decimal.Decimal("-0.000001"), + ] + assert _one_column(temp_db, pa.decimal128(38, 6), values) == values + + def test_list_columns(self, temp_db): + assert _one_column(temp_db, pa.list_(pa.int64()), [[1, 2, 3], [], [4]]) == [ + [1, 2, 3], + [], + [4], + ] + + def test_fixed_size_list_of_floats(self, temp_db): + values = [[0.5, 1.5], [2.5, 3.5]] + assert _one_column(temp_db, pa.list_(pa.float32(), 2), values) == values + + def test_list_of_strings_and_large_list(self, temp_db): + values = [["a", "b"], ["c"]] + assert _one_column(temp_db, pa.large_list(pa.string()), values) == values + + def test_struct_stores_as_a_map(self, temp_db): + values = [{"a": 1, "b": "x"}, {"a": 2, "b": "y"}] + assert ( + _one_column( + temp_db, pa.struct([("a", pa.int64()), ("b", pa.string())]), values + ) + == values + ) + + def test_binary(self, temp_db): + got = _one_column(temp_db, pa.binary(), [b"\x00\x01\xff"]) + assert [b & 0xFF for b in got[0]] == [ + 0, + 1, + 255, + ] # a byte[] reads back as signed ints + + def test_all_null_column_stores_no_value(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Nothing") + table = pa.table({"id": [1, 2], "gone": pa.nulls(2)}) + assert temp_db.insert_arrow("Nothing", table) == 2 + assert all(r.get("gone") is None for r in _rows(temp_db, "SELECT FROM Nothing")) + + def test_uint64_past_the_signed_range_is_refused(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Big") + table = pa.table({"v": pa.array([2**63], pa.uint64())}) + with pytest.raises(ValueError): + temp_db.insert_arrow("Big", table) + assert _count(temp_db, "Big") == 0 + + +class TestPrecision: + def test_declared_precisions(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Prec") + for name, decl in ( + ("sec", "DATETIME_SECOND"), + ("ms", "DATETIME"), + ("us", "DATETIME_MICROS"), + ("ns", "DATETIME_NANOS"), + ): + temp_db.command("sql", f"CREATE PROPERTY Prec.{name} {decl}") + # 2020-01-02 03:04:05.123456789 as nanoseconds since the epoch + ns = 1_577_934_245_123_456_789 + table = pa.table( + {n: pa.array([ns], pa.timestamp("ns")) for n in ("sec", "ms", "us", "ns")} + ) + assert temp_db.insert_arrow("Prec", table) == 1 + row = _rows(temp_db, "SELECT FROM Prec")[0] + assert row["sec"] == dt.datetime(2020, 1, 2, 3, 4, 5) + assert row["ms"] == dt.datetime(2020, 1, 2, 3, 4, 5, 123000) + # a Python datetime holds microseconds, so the nanosecond column reads back to the microsecond + assert row["us"] == dt.datetime(2020, 1, 2, 3, 4, 5, 123456) + assert row["ns"] == dt.datetime(2020, 1, 2, 3, 4, 5, 123456) + + def test_dates_before_the_epoch(self, temp_db): + values = [ + dt.datetime(1900, 1, 1, 0, 0, 1), + dt.datetime(1969, 12, 31, 23, 59, 59, 999000), + ] + assert _one_column(temp_db, pa.timestamp("ms"), values) == values + + +class TestNulls: + def test_every_kind_stores_nulls_as_missing(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Nul") + table = pa.table( + { + "i": pa.array([1, None, 3], pa.int64()), + "f": pa.array([1.5, None, 3.5], pa.float64()), + "b": pa.array([True, None, False], pa.bool_()), + "s": pa.array(["a", None, "c"], pa.string()), + "t": pa.array( + [dt.datetime(2020, 1, 1), None, dt.datetime(2021, 1, 1)], + pa.timestamp("us"), + ), + "d": pa.array( + [dt.date(2020, 1, 1), None, dt.date(2021, 1, 1)], pa.date32() + ), + "n": pa.array( + [decimal.Decimal("1.5"), None, decimal.Decimal("2.5")], + pa.decimal128(10, 1), + ), + "l": pa.array([[1], None, [3]], pa.list_(pa.int64())), + } + ) + table = table.append_column("k", pa.array([0, 1, 2], pa.int64())) + assert temp_db.insert_arrow("Nul", table) == 3 + middle = _rows(temp_db, "SELECT FROM Nul WHERE k = 1")[0] + assert [middle.get(c) for c in "ifbstdnl"] == [None] * 8 + first = _rows(temp_db, "SELECT FROM Nul WHERE k = 0")[0] + assert ( + first["i"] == 1 + and first["f"] == 1.5 + and first["b"] is True + and first["s"] == "a" + ) + assert first["t"] == dt.datetime(2020, 1, 1) and first["d"] == dt.date( + 2020, 1, 1 + ) + assert first["n"] == decimal.Decimal("1.5") and first["l"] == [1] + + def test_pandas_nullable_integer_keeps_its_values(self, temp_db): + # a float64 with NaN would lose the precision past 2**53; Int64 with must not + temp_db.command("sql", "CREATE DOCUMENT TYPE Nullable") + big = 2**60 + 1 + df = pd.DataFrame({"k": [0, 1], "v": pd.array([big, None], dtype="Int64")}) + assert temp_db.insert_dataframe("Nullable", df) == 2 + rows = {r["k"]: r.get("v") for r in _rows(temp_db, "SELECT FROM Nullable")} + assert rows == {0: big, 1: None} + + def test_pandas_nan_is_a_null_and_polars_nan_stays_nan(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE FromPandas") + temp_db.command("sql", "CREATE DOCUMENT TYPE FromPolars") + temp_db.insert_dataframe( + "FromPandas", pd.DataFrame({"k": [0, 1], "x": [1.5, float("nan")]}) + ) + temp_db.insert_dataframe( + "FromPolars", pl.DataFrame({"k": [0, 1], "x": [1.5, float("nan")]}) + ) + pandas_row = _rows(temp_db, "SELECT FROM FromPandas WHERE k = 1")[0] + polars_row = _rows(temp_db, "SELECT FROM FromPolars WHERE k = 1")[0] + assert pandas_row.get("x") is None + assert polars_row["x"] != polars_row["x"] # NaN + + +class TestDeclaredTypes: + def test_values_convert_to_the_declared_type(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Decl") + for prop, decl in ( + ("small", "INTEGER"), + ("short", "SHORT"), + ("as_double", "DOUBLE"), + ("as_decimal", "DECIMAL"), + ("as_text", "STRING"), + ("vec", "ARRAY_OF_FLOATS"), + ("born", "DATE"), + ("seen", "DATETIME_MICROS"), + ): + temp_db.command("sql", f"CREATE PROPERTY Decl.{prop} {decl}") + df = pd.DataFrame( + { + "small": np.array([1, 2], dtype=np.int64), + "short": np.array([3, 4], dtype=np.int8), + "as_double": np.array([5, 6], dtype=np.int32), + "as_decimal": pd.array( + [decimal.Decimal("1.25"), decimal.Decimal("2.5")], + dtype=pd.ArrowDtype(pa.decimal128(5, 2)), + ), + "as_text": np.array([7, 8], dtype=np.int64), + "vec": [[0.5, 1.5], [2.5, 3.5]], + "born": [dt.date(2020, 1, 1), dt.date(2021, 1, 1)], + "seen": pd.to_datetime( + ["2020-01-01 00:00:00.000001", "2020-01-02 00:00:00.000000"] + ), + } + ) + assert temp_db.insert_dataframe("Decl", df) == 2 + row = _rows(temp_db, "SELECT FROM Decl WHERE small = 1")[0] + assert row["small"] == 1 and row["short"] == 3 + assert row["as_double"] == 5.0 and type(row["as_double"]) is float + assert row["as_decimal"] == decimal.Decimal("1.25") + assert row["as_text"] == "7" + assert list(row["vec"]) == [0.5, 1.5] + assert row["born"] == dt.date(2020, 1, 1) + assert row["seen"] == dt.datetime(2020, 1, 1, 0, 0, 0, 1) + + def test_a_value_past_the_declared_range_fails_and_rolls_back(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Range") + temp_db.command("sql", "CREATE PROPERTY Range.v INTEGER") + with pytest.raises(ArcadeDBError): + temp_db.insert_arrow( + "Range", pa.table({"v": pa.array([1, 2**40], pa.int64())}) + ) + assert _count(temp_db, "Range") == 0 + + @pytest.mark.parametrize( + "declared, column", + [ + ("INTEGER", pa.array([1.5], pa.float64())), # would truncate to 1 + ("LONG", pa.array(["12"], pa.string())), # would be parsed + ( + "DATE", + pa.array([dt.datetime(2020, 1, 1, 12)], pa.timestamp("us")), + ), # would lose the time + ("BOOLEAN", pa.array([1], pa.int64())), + ("STRING", pa.array([[1]], pa.list_(pa.int64()))), + ("DOUBLE", pa.array([dt.date(2020, 1, 1)], pa.date32())), + ], + ) + def test_a_kind_that_would_lose_or_guess_is_refused_before_a_row_is_written( + self, temp_db, declared, column + ): + temp_db.command("sql", "CREATE DOCUMENT TYPE Refuse") + temp_db.command("sql", f"CREATE PROPERTY Refuse.v {declared}") + with pytest.raises(TypeError, match=f"Refuse.v. is declared {declared}"): + temp_db.insert_arrow( + "Refuse", pa.table({"ok": pa.array([0] * len(column)), "v": column}) + ) + assert _count(temp_db, "Refuse") == 0 + + def test_declared_property_of_a_parent_type_applies(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Parent") + temp_db.command("sql", "CREATE PROPERTY Parent.v INTEGER") + temp_db.command("sql", "CREATE DOCUMENT TYPE Child EXTENDS Parent") + with pytest.raises(TypeError, match="declared INTEGER"): + temp_db.insert_arrow("Child", pa.table({"v": [1.5]})) + + def test_an_undeclared_column_next_to_declared_ones_is_stored_as_it_is( + self, temp_db + ): + temp_db.command("sql", "CREATE DOCUMENT TYPE Mixed") + temp_db.command("sql", "CREATE PROPERTY Mixed.a INTEGER") + assert ( + temp_db.insert_dataframe("Mixed", pd.DataFrame({"a": [1], "b": [2.5]})) == 1 + ) + row = _rows(temp_db, "SELECT FROM Mixed")[0] + assert row["a"] == 1 and row["b"] == 2.5 + + +class TestInputsAndErrors: + def test_pandas_polars_and_arrow_inputs_store_the_same_records(self, temp_db): + data = { + "id": [1, 2, 3], + "name": ["a", "b", "c"], + "score": [0.5, 1.5, 2.5], + "ok": [True, False, True], + } + sources = { + "FromPandas": pd.DataFrame(data), + "FromPolars": pl.DataFrame(data), + "FromTable": pa.table(data), + "FromBatch": pa.RecordBatch.from_pydict(data), + "FromReader": pa.RecordBatchReader.from_batches( + pa.table(data).schema, pa.table(data).to_batches() + ), + } + stored = {} + for name, source in sources.items(): + temp_db.command("sql", f"CREATE DOCUMENT TYPE {name}") + method = ( + temp_db.insert_dataframe + if name in ("FromPandas", "FromPolars") + else temp_db.insert_arrow + ) + assert method(name, source) == 3 + sql = f"SELECT FROM {name} ORDER BY id" # nosec B608 - test-controlled + stored[name] = [ + {k: v for k, v in r.items() if not k.startswith("@")} + for r in _rows(temp_db, sql) + ] + assert all(rows == stored["FromTable"] for rows in stored.values()) + assert stored["FromTable"][0] == { + "id": 1, + "name": "a", + "score": 0.5, + "ok": True, + } + + def test_insert_dataframe_takes_an_arrow_capsule_object(self, temp_db): + class Wrapper: # a library that speaks the Arrow PyCapsule interface, and nothing else + def __init__(self, table): + self._table = table + + def __arrow_c_stream__(self, requested_schema=None): + return self._table.__arrow_c_stream__(requested_schema) + + temp_db.command("sql", "CREATE DOCUMENT TYPE Capsule") + assert ( + temp_db.insert_dataframe("Capsule", Wrapper(pa.table({"id": [1, 2]}))) == 2 + ) + + def test_pandas_index_is_not_loaded(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Indexed") + df = pd.DataFrame({"v": [1, 2]}, index=["x", "y"]) + temp_db.insert_dataframe("Indexed", df) + assert set(_rows(temp_db, "SELECT FROM Indexed")[0]) - { + "@rid", + "@type", + "@cat", + } == {"v"} + temp_db.insert_dataframe("Indexed", df.rename_axis("name").reset_index()) + assert _count(temp_db, "Indexed", "WHERE name = 'y'") == 1 + + def test_empty_frame_writes_nothing(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Empty") + assert ( + temp_db.insert_dataframe( + "Empty", pd.DataFrame({"a": pd.Series([], dtype="int64")}) + ) + == 0 + ) + assert _count(temp_db, "Empty") == 0 + + def test_missing_type_and_edge_type_are_refused(self, temp_db): + temp_db.command("sql", "CREATE EDGE TYPE Link") + with pytest.raises(ArcadeDBError, match="does not exist"): + temp_db.insert_dataframe("Nope", pd.DataFrame({"a": [1]})) + with pytest.raises(ArcadeDBError, match="edge type"): + temp_db.insert_dataframe("Link", pd.DataFrame({"a": [1]})) + + def test_bad_frames_are_refused_before_anything_is_written(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Bad") + with pytest.raises(ValueError, match="at least one column"): + temp_db.insert_dataframe("Bad", pd.DataFrame()) + with pytest.raises(ValueError, match="strings"): + temp_db.insert_dataframe("Bad", pd.DataFrame({0: [1]})) + with pytest.raises(ValueError, match="duplicate"): + temp_db.insert_dataframe("Bad", pd.DataFrame([[1, 2]], columns=["a", "a"])) + with pytest.raises(TypeError, match="LazyFrame"): + temp_db.insert_dataframe("Bad", pl.DataFrame({"a": [1]}).lazy()) + with pytest.raises(TypeError, match="DataFrame"): + temp_db.insert_dataframe("Bad", {"a": [1]}) + with pytest.raises(TypeError, match="no ArcadeDB mapping"): + temp_db.insert_arrow("Bad", pa.table({"t": pa.array([1], pa.time64("us"))})) + with pytest.raises(TypeError, match="no ArcadeDB mapping"): + temp_db.insert_arrow( + "Bad", pa.table({"d": pa.array([1], pa.duration("s"))}) + ) + with pytest.raises(TypeError, match="no ArcadeDB mapping"): + temp_db.insert_arrow( + "Bad", + pa.table( + {"m": pa.array([[("k", 1)]], pa.map_(pa.string(), pa.int64()))} + ), + ) + for size in (0, -1, 1.5, True): + with pytest.raises(ValueError, match="batch_size"): + temp_db.insert_dataframe( + "Bad", pd.DataFrame({"a": [1]}), batch_size=size + ) + assert _count(temp_db, "Bad") == 0 + + def test_missing_pyarrow_is_an_import_error_with_the_fix( + self, temp_db, monkeypatch + ): + temp_db.command("sql", "CREATE DOCUMENT TYPE NoArrow") + monkeypatch.setitem(sys.modules, "pyarrow", None) + with pytest.raises(ImportError, match=r"pip install pyarrow"): + temp_db.insert_dataframe("NoArrow", pd.DataFrame({"a": [1]})) + + def test_closed_database_is_refused(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Closed") + temp_db.close() + with pytest.raises(ArcadeDBError): + temp_db.insert_dataframe("Closed", pd.DataFrame({"a": [1]})) + + +class TestBatching: + def test_a_frame_of_100k_rows_loads_whole(self, temp_db): + n = 100_000 + temp_db.command("sql", "CREATE DOCUMENT TYPE Big") + temp_db.command("sql", "CREATE PROPERTY Big.id LONG") + ids = np.arange(n, dtype=np.int64) + df = pd.DataFrame( + { + "id": ids, + "price": ids * 0.25, + "label": [f"row_{i}" for i in range(n)], + "flag": ids % 3 == 0, + "ts": pd.Timestamp("2020-01-01") + pd.to_timedelta(ids, unit="s"), + } + ) + assert temp_db.insert_dataframe("Big", df) == n + assert _count(temp_db, "Big") == n + total = _rows( + temp_db, "SELECT sum(id) AS s, count(DISTINCT label) AS d FROM Big" + )[0] + assert total["s"] == int(ids.sum()) and total["d"] == n + last_sql = f"SELECT FROM Big WHERE id = {n - 1}" # nosec B608 - test-controlled + last = _rows(temp_db, last_sql)[0] + assert last["ts"] == dt.datetime(2020, 1, 1) + dt.timedelta(seconds=n - 1) + + @pytest.mark.parametrize("batch_size", [1, 7, 10, 25]) + def test_a_batch_size_that_does_not_divide_the_rows(self, temp_db, batch_size): + temp_db.command("sql", "CREATE DOCUMENT TYPE Odd") + assert ( + temp_db.insert_dataframe( + "Odd", pd.DataFrame({"id": range(23)}), batch_size=batch_size + ) + == 23 + ) + assert _count(temp_db, "Odd") == 23 + + def test_a_table_of_many_chunks_loads_every_chunk(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Chunks") + table = pa.concat_tables( + [pa.table({"id": list(range(i * 5, i * 5 + 5))}) for i in range(4)] + ) + assert temp_db.insert_arrow("Chunks", table, batch_size=8) == 20 + assert _count(temp_db, "Chunks") == 20 + + def test_a_failure_rolls_back_its_batch_and_says_what_was_committed(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Dup") + temp_db.command("sql", "CREATE PROPERTY Dup.id LONG") + temp_db.command("sql", "CREATE INDEX ON Dup (id) UNIQUE") + ids = list(range(10)) + [3] # the 11th row repeats id 3, in the second batch + with pytest.raises(ArcadeDBError, match=r"10 row\(s\) of the batches before"): + temp_db.insert_dataframe("Dup", pd.DataFrame({"id": ids}), batch_size=10) + assert _count(temp_db, "Dup") == 10 + + def test_the_failure_message_tells_committed_from_written_in_your_transaction( + self, temp_db + ): + """Outside a transaction the earlier batches are committed; inside the caller's own they + joined it, so the message must not claim they are committed (a rollback removes them). + """ + temp_db.command("sql", "CREATE DOCUMENT TYPE Dup2") + temp_db.command("sql", "CREATE PROPERTY Dup2.id LONG") + temp_db.command("sql", "CREATE INDEX ON Dup2 (id) UNIQUE") + frame = pd.DataFrame({"id": list(range(10)) + [3]}) # the 11th row repeats id 3 + + with pytest.raises(ArcadeDBError) as outside: + temp_db.insert_dataframe("Dup2", frame, batch_size=10) + assert "10 row(s) of the batches before this one were committed" in str( + outside.value + ) + assert _count(temp_db, "Dup2") == 10 + with temp_db.transaction(): + temp_db.command("sql", "DELETE FROM Dup2") + + temp_db.begin() + try: + with pytest.raises(ArcadeDBError) as inside: + temp_db.insert_dataframe("Dup2", frame, batch_size=10) + finally: + temp_db.rollback() + assert "10 row(s) of the batches before this one were written in your" in str( + inside.value + ) + assert "not committed" in str(inside.value) + assert "were committed" not in str(inside.value) + assert _count(temp_db, "Dup2") == 0 # the caller's rollback removed them + + def test_a_callers_transaction_is_left_to_the_caller(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Mine") + with pytest.raises(RuntimeError, match="stop"): + with temp_db.transaction(): + temp_db.insert_dataframe( + "Mine", pd.DataFrame({"id": [1, 2, 3]}), batch_size=2 + ) + raise RuntimeError("stop") + assert _count(temp_db, "Mine") == 0 + + def test_parallel_documents(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Par") + assert ( + temp_db.insert_dataframe( + "Par", pd.DataFrame({"id": range(1000)}), batch_size=300, parallel=True + ) + == 1000 + ) + assert _count(temp_db, "Par") == 1000 + + +class TestVertices: + def test_a_vertex_type_gets_vertices(self, temp_db): + temp_db.command("sql", "CREATE VERTEX TYPE Person") + temp_db.command("sql", "CREATE PROPERTY Person.id LONG") + df = pd.DataFrame( + { + "id": np.arange(25, dtype=np.int64), + "name": [f"p{i}" for i in range(25)], + "born": pd.to_datetime(["2000-01-01"] * 25), + "score": [None if i % 5 == 0 else i * 0.5 for i in range(25)], + } + ) + assert temp_db.insert_dataframe("Person", df, batch_size=10) == 25 + assert _count(temp_db, "Person") == 25 + row = _rows(temp_db, "SELECT FROM Person WHERE id = 7")[0] + assert ( + row["name"] == "p7" + and row["born"] == dt.datetime(2000, 1, 1) + and row["score"] == 3.5 + ) + assert ( + _rows(temp_db, "SELECT outE().size() AS e FROM Person WHERE id = 7")[0]["e"] + == 0 + ) # a vertex + assert _rows(temp_db, "SELECT FROM Person WHERE id = 5")[0].get("score") is None + + def test_vertices_from_polars_and_edges_between_them(self, temp_db): + temp_db.command("sql", "CREATE VERTEX TYPE City") + temp_db.command("sql", "CREATE EDGE TYPE Road") + assert temp_db.insert_dataframe("City", pl.DataFrame({"name": ["a", "b"]})) == 2 + a, b = ( + r["rid"] + for r in _rows(temp_db, "SELECT @rid AS rid FROM City ORDER BY name") + ) + with temp_db.graph_batch() as batch: + batch.new_edges([str(a)], "Road", [str(b)]) + assert _count(temp_db, "Road") == 1 + + def test_declared_types_apply_to_vertices_too(self, temp_db): + temp_db.command("sql", "CREATE VERTEX TYPE Typed") + temp_db.command("sql", "CREATE PROPERTY Typed.v INTEGER") + with pytest.raises(TypeError, match="declared INTEGER"): + temp_db.insert_dataframe("Typed", pd.DataFrame({"v": [1.5]})) + + def test_parallel_is_for_documents_only(self, temp_db): + temp_db.command("sql", "CREATE VERTEX TYPE Pv") + with pytest.raises(ValueError, match="vertex type"): + temp_db.insert_dataframe("Pv", pd.DataFrame({"a": [1]}), parallel=True) + + +class TestToPolars: + def test_round_trip_through_polars(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Trip") + frame = pl.DataFrame( + { + "id": pl.Series([1, 2, 3], dtype=pl.Int64), + "score": [0.5, None, 2.5], + "name": ["a", None, "c"], + "ok": [True, False, None], + "at": pl.Series( + [ + dt.datetime(2020, 1, 1), + None, + dt.datetime(2020, 1, 3, 4, 5, 6, 789000), + ] + ).cast(pl.Datetime("ms")), + "amount": pl.Series( + [decimal.Decimal("1.5"), None, decimal.Decimal("3.25")], + dtype=pl.Decimal(10, 2), + ), + } + ) + assert temp_db.insert_dataframe("Trip", frame) == 3 + back = temp_db.query( + "sql", "SELECT id, score, name, ok, at, amount FROM Trip ORDER BY id" + ).to_polars() + assert isinstance(back, pl.DataFrame) + assert back.schema["id"] == pl.Int64 and back.schema["name"] == pl.String + assert back.schema["score"] == pl.Float64 and back.schema["ok"] == pl.Boolean + assert back.schema["at"] == pl.Datetime("ms") + assert ( + back.get_column("amount").to_list() == frame.get_column("amount").to_list() + ) + assert back.drop("amount").equals( + frame.drop("amount").cast({"at": pl.Datetime("ms")}) + ) + + def test_matches_to_arrow(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Same") + temp_db.insert_dataframe("Same", pd.DataFrame({"id": [1, 2], "v": [0.5, 1.5]})) + sql = "SELECT id, v FROM Same ORDER BY id" + assert ( + temp_db.query("sql", sql).to_polars().to_dicts() + == pl.from_arrow(temp_db.query("sql", sql).to_arrow()).to_dicts() + ) + + def test_columns_and_batch_size_pass_through(self, temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Pass") + temp_db.insert_dataframe( + "Pass", pd.DataFrame({"id": range(10), "v": range(10)}) + ) + frame = temp_db.query("sql", "SELECT id, v FROM Pass ORDER BY id").to_polars( + batch_size=3, columns=["id"] + ) + assert frame.columns == ["id"] and frame.get_column("id").to_list() == list( + range(10) + ) + + def test_missing_polars_is_an_import_error_with_the_fix(self, temp_db, monkeypatch): + temp_db.command("sql", "CREATE DOCUMENT TYPE NoPolars") + monkeypatch.setitem(sys.modules, "polars", None) + with pytest.raises( + ImportError, + match=r"polars is required for to_polars\(\).*arcadedb-embedded\[polars\]", + ): + temp_db.query("sql", "SELECT FROM NoPolars").to_polars() + + def test_missing_pyarrow_is_an_import_error_with_the_fix( + self, temp_db, monkeypatch + ): + temp_db.command("sql", "CREATE DOCUMENT TYPE NoArrow2") + temp_db.insert_dataframe("NoArrow2", pd.DataFrame({"a": [1]})) + monkeypatch.setitem(sys.modules, "pyarrow", None) + with pytest.raises(ImportError, match="needs pyarrow"): + temp_db.query("sql", "SELECT FROM NoArrow2").to_polars() diff --git a/tests/test_vector_product_quantization.py b/tests/test_vector_product_quantization.py new file mode 100644 index 00000000..d5d6d9d0 --- /dev/null +++ b/tests/test_vector_product_quantization.py @@ -0,0 +1,353 @@ +""" +PRODUCT (PQ) vector quantization after ArcadeData/arcadedb#10545 (issue #10514). + +A PRODUCT index walks its graph on the in-memory PQ codes and re-ranks the best +`k * rerankOversample` candidates on the stored vectors. The vector guide's "Product +quantization in SQL" section shows the statements these tests run: the +`CREATE INDEX ... METADATA` settings (`rerankOversample`, `fusedPQ`, `pqTrainingLimit`) +and the per-query `rerank` option of `vectorNeighbors`. + +CI builds its wheel from upstream's 26.11.1-SNAPSHOT image, which is republished only when +upstream's deploy job succeeds and at the time of this change still held an engine from before +#10545, with the same version string as the engine that has it. So the version cannot tell the +two apart, and `has_pq_rerank` probes the one thing that changed: the index setting +`rerankOversample` is accepted. On an engine without #10545 the same tests assert the old +answer (the settings and the option are refused, `find_nearest` is the exact walk SQL also +takes), with no skips, as `test_vector_sql.py` does for `where` (the repository accepts no +skips, scripts/check_test_skips.py). When 26.11.1 is released and the wheel the suite runs on +is built from it, delete `has_pq_rerank` and the old-answer branches. The new-answer branch runs in CI +only on a build from source (`engine-from-source: true`). +""" + +import numpy as np +import pytest + +import arcadedb_embedded as arcadedb +from arcadedb_embedded import create_database + +DIMS = 32 +COUNT = 2000 +K = 10 + + +@pytest.fixture +def test_db(tmp_path): + db = create_database(str(tmp_path / "test_vector_pq_db")) + yield db + db.drop() + + +@pytest.fixture +def has_pq_rerank(tmp_path): + """True when the engine has #10545: it accepts the `rerankOversample` index setting.""" + with create_database(str(tmp_path / "probe")) as probe: + probe.command("sql", "CREATE DOCUMENT TYPE Probe") + probe.command("sql", "CREATE PROPERTY Probe.embedding ARRAY_OF_FLOATS") + try: + probe.command( + "sql", + "CREATE INDEX ON Probe (embedding) LSM_VECTOR METADATA " + '{"dimensions": 4, "quantization": "PRODUCT", "rerankOversample": 1}', + ) + except Exception as e: + assert "Unsupported metadata key" in str(e), e + return False + return True + + +@pytest.fixture(scope="module") +def corpus(): + rng = np.random.default_rng(7) + vectors = rng.standard_normal((COUNT, DIMS)).astype(np.float32) + queries = rng.standard_normal((20, DIMS)).astype(np.float32) + return vectors, queries + + +def _exact_neighbors(vectors, query, k=K): + distances = ((vectors - query) ** 2).sum(axis=1) + order = np.argsort(distances)[:k] + return [int(i) for i in order], distances + + +def _make_index(db, vectors, metadata=""): + """The statement the guide shows, on a corpus big enough to train 64 clusters.""" + db.command("sql", "CREATE DOCUMENT TYPE PqDoc") + db.command("sql", "CREATE PROPERTY PqDoc.id INTEGER") + db.command("sql", "CREATE PROPERTY PqDoc.embedding ARRAY_OF_FLOATS") + db.command( + "sql", + "CREATE INDEX ON PqDoc (embedding) LSM_VECTOR METADATA " + f'{{"dimensions": {DIMS}, "similarity": "EUCLIDEAN", ' + '"quantization": "PRODUCT", "pqSubspaces": 8, "pqClusters": 64' + f"{metadata}}}", + ) + with db.transaction(): + for i, vec in enumerate(vectors): + db.command( + "sql", + "INSERT INTO PqDoc SET id = ?, embedding = ?", + i, + arcadedb.to_java_float_array(vec), + ) + _build(db) + + +def _build(db): + """Searches before the graph exists are answered exactly from the delta scan.""" + db.schema.get_vector_index("PqDoc", "embedding").build_graph_now() + + +def _search(db, query, options="{ efSearch: 100 }", k=K): + return db.query( + "sql", + "SELECT id, distance FROM " # nosec B608 - test-owned option literals + f"(SELECT expand(vectorNeighbors('PqDoc[embedding]', ?, {k}, {options})))", # nosec B608 + arcadedb.to_java_float_array(query), + ).to_list() + + +def _metadata(db): + index = db.schema.get_index_by_name("PqDoc[embedding]") + if "TypeIndex" in index.getClass().getName(): + index = index.getSubIndexes().get(0) + return index.getMetadata() + + +def _recall(db, vectors, queries, options): + hits = 0 + for query in queries: + wanted, _ = _exact_neighbors(vectors, query) + found = [r["id"] for r in _search(db, query, options)] + hits += len(set(wanted) & set(found)) + return hits / (K * len(queries)) + + +class TestProductQuantizationSql: + def test_index_settings_reach_the_index(self, test_db, corpus, has_pq_rerank): + vectors, _ = corpus + if not has_pq_rerank: + with pytest.raises(Exception, match="Unsupported metadata key"): + _make_index(test_db, vectors, ', "rerankOversample": 30') + return + _make_index( + test_db, + vectors, + ', "rerankOversample": 30, "pqTrainingLimit": 1500, "fusedPQ": false', + ) + + meta = _metadata(test_db) + + assert str(meta.quantizationType) == "PRODUCT" + assert int(meta.rerankOversample) == 30 + assert int(meta.pqTrainingLimit) == 1500 + assert bool(meta.fusedPQ) is False + assert int(meta.pqSubspaces) == 8 + assert int(meta.pqClusters) == 64 + + def test_rerank_oversample_defaults_to_ten_and_fused_pq_to_off( + self, test_db, corpus, has_pq_rerank + ): + _make_index(test_db, corpus[0]) + + meta = _metadata(test_db) + + if not has_pq_rerank: + assert not hasattr(meta, "rerankOversample") + return + + assert int(meta.rerankOversample) == 10 + assert bool(meta.fusedPQ) is False + + def test_negative_rerank_oversample_is_refused(self, test_db, has_pq_rerank): + test_db.command("sql", "CREATE DOCUMENT TYPE PqDoc") + test_db.command("sql", "CREATE PROPERTY PqDoc.embedding ARRAY_OF_FLOATS") + refused = ( + "rerankOversample must be at least 0" + if has_pq_rerank + else "Unsupported metadata key 'rerankOversample'" + ) + + with pytest.raises(Exception, match=refused): + test_db.command( + "sql", + "CREATE INDEX ON PqDoc (embedding) LSM_VECTOR METADATA " + f'{{"dimensions": {DIMS}, "quantization": "PRODUCT", ' + '"rerankOversample": -1}', + ) + + def test_default_search_returns_exact_distances_and_good_recall( + self, test_db, corpus + ): + vectors, queries = corpus + _make_index(test_db, vectors) + + for query in queries[:5]: + _, exact = _exact_neighbors(vectors, query) + rows = _search(test_db, query) + assert len(rows) == K + for row in rows: + assert row["distance"] == pytest.approx( + float(exact[row["id"]]), rel=1e-4 + ) + + assert _recall(test_db, vectors, queries, "{ efSearch: 100 }") >= 0.9 + + def test_rerank_zero_returns_pq_estimates_with_lower_recall( + self, test_db, corpus, has_pq_rerank + ): + vectors, queries = corpus + _make_index(test_db, vectors) + if not has_pq_rerank: + with pytest.raises(Exception, match="Unknown option 'rerank'"): + _search(test_db, queries[0], "{ efSearch: 100, rerank: 0 }") + return + + reranked = _recall(test_db, vectors, queries, "{ efSearch: 100 }") + pq_only = _recall(test_db, vectors, queries, "{ efSearch: 100, rerank: 0 }") + assert pq_only < reranked + + # the distances are PQ estimates, not the exact ones + _, exact = _exact_neighbors(vectors, queries[0]) + rows = _search(test_db, queries[0], "{ efSearch: 100, rerank: 0 }") + assert any( + row["distance"] != pytest.approx(float(exact[row["id"]]), rel=1e-4) + for row in rows + ) + + def test_negative_rerank_option_is_refused(self, test_db, corpus, has_pq_rerank): + _make_index(test_db, corpus[0]) + refused = "rerank must be >= 0" if has_pq_rerank else "Unknown option 'rerank'" + + with pytest.raises(Exception, match=refused): + _search(test_db, corpus[1][0], "{ rerank: -1 }") + + def test_fused_pq_index_returns_k_results(self, test_db, corpus, has_pq_rerank): + vectors, queries = corpus + test_db.command("sql", "CREATE DOCUMENT TYPE PqDoc") + test_db.command("sql", "CREATE PROPERTY PqDoc.id INTEGER") + test_db.command("sql", "CREATE PROPERTY PqDoc.embedding ARRAY_OF_FLOATS") + create_fused = lambda: test_db.command( # noqa: E731 + "sql", + "CREATE INDEX ON PqDoc (embedding) LSM_VECTOR METADATA " + f'{{"dimensions": {DIMS}, "similarity": "EUCLIDEAN", ' + '"quantization": "PRODUCT", "pqSubspaces": 8, "pqClusters": 256, ' + '"fusedPQ": true}', + ) + if not has_pq_rerank: + with pytest.raises(Exception, match="Unsupported metadata key 'fusedPQ'"): + create_fused() + return + create_fused() + with test_db.transaction(): + for i, vec in enumerate(vectors): + test_db.command( + "sql", + "INSERT INTO PqDoc SET id = ?, embedding = ?", + i, + arcadedb.to_java_float_array(vec), + ) + + _build(test_db) + + assert bool(_metadata(test_db).fusedPQ) is True + assert len(_search(test_db, queries[0])) == K + + +class TestProductQuantizationHelpers: + def test_find_nearest_answers_like_sql_on_a_product_index( + self, test_db, corpus, has_pq_rerank + ): + """With rerankOversample 0 both return the PQ scores; the exact walk would not. + + An engine without #10545 has no such setting and answers both exactly. + """ + vectors, queries = corpus + _make_index( + test_db, vectors, ', "rerankOversample": 0' if has_pq_rerank else "" + ) + index = test_db.schema.get_vector_index("PqDoc", "embedding") + query = queries[0] + + sql_rows = _search(test_db, query) + found = index.find_nearest(query, k=K, ef_search=100) + + assert [rec.get("id") for rec, _ in found] == [r["id"] for r in sql_rows] + assert [score for _, score in found] == pytest.approx( + [r["distance"] for r in sql_rows] + ) + if not has_pq_rerank: + return + _, exact = _exact_neighbors(vectors, query) + assert any( + score != pytest.approx(float(exact[rec.get("id")]), rel=1e-4) + for rec, score in found + ) + + def test_find_nearest_returns_exact_distances_by_default(self, test_db, corpus): + vectors, queries = corpus + _make_index(test_db, vectors) + index = test_db.schema.get_vector_index("PqDoc", "embedding") + + found = index.find_nearest(queries[0], k=K, ef_search=100) + + _, exact = _exact_neighbors(vectors, queries[0]) + assert len(found) == K + for rec, score in found: + assert score == pytest.approx(float(exact[rec.get("id")]), rel=1e-4) + + def test_create_vector_index_passes_the_new_settings(self, test_db, has_pq_rerank): + test_db.command("sql", "CREATE VERTEX TYPE PqVertex") + test_db.command("sql", "CREATE PROPERTY PqVertex.embedding ARRAY_OF_FLOATS") + create = lambda: test_db.create_vector_index( # noqa: E731 + "PqVertex", + "embedding", + dimensions=DIMS, + distance_function="euclidean", + quantization="PRODUCT", + pq_subspaces=8, + pq_clusters=64, + rerank_oversample=25, + fused_pq=False, + ) + if not has_pq_rerank: + with pytest.raises( + arcadedb.ArcadeDBError, match="Unsupported metadata key" + ): + create() + return + index = create() + + meta = index.get_metadata() + + assert meta["rerank_oversample"] == 25 + assert meta["fused_pq"] is False + assert meta["pq_subspaces"] == 8 + + def test_get_metadata_reports_the_engine_defaults(self, test_db, has_pq_rerank): + test_db.command("sql", "CREATE VERTEX TYPE PqVertex") + test_db.command("sql", "CREATE PROPERTY PqVertex.embedding ARRAY_OF_FLOATS") + index = test_db.create_vector_index( + "PqVertex", "embedding", dimensions=DIMS, quantization="PRODUCT" + ) + + meta = index.get_metadata() + + if not has_pq_rerank: + assert "rerank_oversample" not in meta and "fused_pq" not in meta + return + assert meta["rerank_oversample"] == 10 + assert meta["fused_pq"] is False + + @pytest.mark.parametrize("setting", ["rerank_oversample", "fused_pq"]) + def test_new_settings_require_product_quantization(self, test_db, setting): + test_db.command("sql", "CREATE VERTEX TYPE PqVertex") + test_db.command("sql", "CREATE PROPERTY PqVertex.embedding ARRAY_OF_FLOATS") + + with pytest.raises(arcadedb.ArcadeDBError, match="require quantization"): + test_db.create_vector_index( + "PqVertex", + "embedding", + dimensions=DIMS, + quantization="INT8", + **{setting: 5 if setting == "rerank_oversample" else True}, + ) diff --git a/tests/test_vector_sql.py b/tests/test_vector_sql.py index 3310837c..03bf0c57 100644 --- a/tests/test_vector_sql.py +++ b/tests/test_vector_sql.py @@ -1109,3 +1109,251 @@ def test_find_nearest_repeated_calls_consistent(self, test_db): ids1 = [rec.get("id") for rec, _ in r1] ids2 = [rec.get("id") for rec, _ in r2] assert ids1 == ids2 and len(ids1) == 3 + + +def _engine_version(): + import jpype + + raw = str(jpype.JClass("com.arcadedb.Constants").getRawVersion()) + parts = [int(p) if p.isdigit() else 0 for p in raw.split("-")[0].split(".")[:3]] + return tuple(parts + [0] * (3 - len(parts))) + + +# `where` in the vector.neighbors options is new in ArcadeDB 26.11.1. The JVM only runs +# once a database fixture has started it, so the version is read inside each test, and +# both engines are asserted on (the repository accepts no skips, scripts/check_test_skips.py). +WHERE_SINCE = (26, 11, 1) + + +class TestVectorSearchOptionsSql: + """The vector.neighbors option map and index metadata the vector guide documents. + + The guide's "Search Options in SQL" and "Graph build settings in SQL" + sections show these exact forms, so each test runs one of them and checks + the answer. + """ + + @staticmethod + def _make_docs(db, metadata=""): + db.command("sql", "CREATE DOCUMENT TYPE OptDoc") + db.command("sql", "CREATE PROPERTY OptDoc.id INTEGER") + db.command("sql", "CREATE PROPERTY OptDoc.lang STRING") + db.command("sql", "CREATE PROPERTY OptDoc.year INTEGER") + db.command("sql", "CREATE PROPERTY OptDoc.embedding ARRAY_OF_FLOATS") + db.command( + "sql", + "CREATE INDEX ON OptDoc (embedding) LSM_VECTOR METADATA " + f'{{"dimensions": 2, "similarity": "EUCLIDEAN"{metadata}}}', + ) + # Squared euclidean distance to the origin is id * id. + with db.transaction(): + for i in range(10): + db.command( + "sql", + "INSERT INTO OptDoc SET id = ?, lang = ?, year = ?, embedding = ?", + i, + "en" if i % 2 == 0 else "fr", + 2015 + i, + arcadedb.to_java_float_array([float(i), 0.0]), + ) + + @staticmethod + def _rids(db): + rows = db.query("sql", "SELECT @rid AS rid, id FROM OptDoc").to_list() + return {r["id"]: str(r["rid"]) for r in rows} + + @staticmethod + def _search(db, options, k=4): + return db.query( + "sql", + "SELECT id, distance FROM " # nosec B608 - test-owned option literals + f"(SELECT expand(vectorNeighbors('OptDoc[embedding]', ?, {k}, {options})))", # nosec B608 + arcadedb.to_java_float_array([0.0, 0.0]), + ).to_list() + + def test_max_distance_drops_far_hits(self, test_db): + self._make_docs(test_db) + + rows = self._search(test_db, "{ maxDistance: 4.5 }") + + assert [r["id"] for r in rows] == [0, 1, 2] + assert [r["distance"] for r in rows] == [0.0, 1.0, 4.0] + + def test_filter_and_ef_search_options(self, test_db): + self._make_docs(test_db) + rids = self._rids(test_db) + allowed = f"['{rids[3]}', '{rids[5]}', '{rids[8]}']" + + rows = self._search(test_db, f"{{ efSearch: 50, filter: {allowed} }}") + + assert [r["id"] for r in rows] == [3, 5, 8] + + def test_group_by_with_group_size_two(self, test_db): + self._make_docs(test_db) + + rows = self._search(test_db, "{ groupBy: 'lang', groupSize: 2 }") + + assert [r["id"] for r in rows] == [0, 1, 2, 3] + + def test_where_map_and_condition(self, test_db): + self._make_docs(test_db) + + if _engine_version() < WHERE_SINCE: + with pytest.raises(Exception, match="Unknown option 'where'"): + self._search(test_db, "{ where: { lang: 'fr' } }") + return + + by_map = self._search(test_db, "{ where: { lang: 'fr' } }") + by_condition = self._search( + test_db, "{ where: \"lang = 'fr' AND year > 2018\" }" + ) + + assert [r["id"] for r in by_map] == [1, 3, 5, 7] + assert [r["id"] for r in by_condition] == [5, 7, 9] + + def test_where_combines_with_filter(self, test_db): + self._make_docs(test_db) + rids = self._rids(test_db) + allowed = f"['{rids[2]}', '{rids[3]}', '{rids[5]}']" + options = f"{{ filter: {allowed}, where: {{ lang: 'fr' }} }}" + + if _engine_version() < WHERE_SINCE: + with pytest.raises(Exception, match="Unknown option 'where'"): + self._search(test_db, options) + return + + rows = self._search(test_db, options) + + assert [r["id"] for r in rows] == [3, 5] + + def test_unknown_option_lists_the_allowed_ones(self, test_db): + self._make_docs(test_db) + + with pytest.raises(Exception, match="Unknown option 'maxDist'"): + self._search(test_db, "{ maxDist: 1.0 }") + + @staticmethod + def _metadata(db, index_name): + index = db.schema.get_index_by_name(index_name) + if "TypeIndex" in index.getClass().getName(): + index = index.getSubIndexes().get(0) + return index.getMetadata() + + def test_graph_build_settings_reach_the_index(self, test_db): + self._make_docs( + test_db, + ', "neighborOverflowFactor": 1.5, "alphaDiversityRelaxation": 1.4, ' + '"addHierarchy": true', + ) + + meta = self._metadata(test_db, "OptDoc[embedding]") + + assert float(meta.neighborOverflowFactor) == pytest.approx(1.5) + assert float(meta.alphaDiversityRelaxation) == pytest.approx(1.4) + assert bool(meta.addHierarchy) is True + + def test_add_hierarchy_defaults_differ_between_sql_and_python(self, test_db): + """SQL follows the engine (flat graph); create_vector_index() builds a hierarchy.""" + self._make_docs(test_db) + meta = self._metadata(test_db, "OptDoc[embedding]") + assert bool(meta.addHierarchy) is False + # The vector guide says the same of quantization: SQL is the engine's NONE. + assert str(meta.quantizationType) == "NONE" + assert float(meta.neighborOverflowFactor) == pytest.approx(1.2) + assert float(meta.alphaDiversityRelaxation) == pytest.approx(1.2) + + test_db.command("sql", "CREATE VERTEX TYPE PyHier") + test_db.command("sql", "CREATE PROPERTY PyHier.embedding ARRAY_OF_FLOATS") + python_index = test_db.create_vector_index("PyHier", "embedding", dimensions=2) + assert python_index.get_metadata()["add_hierarchy"] is True + assert python_index.get_metadata()["quantization"] == "INT8" + + test_db.command("sql", "CREATE VERTEX TYPE PyEngineDefault") + test_db.command( + "sql", "CREATE PROPERTY PyEngineDefault.embedding ARRAY_OF_FLOATS" + ) + engine_default_index = test_db.create_vector_index( + "PyEngineDefault", "embedding", dimensions=2, add_hierarchy=None + ) + assert engine_default_index.get_metadata()["add_hierarchy"] is False + + def test_sql_metadata_gives_the_helper_defaults(self, test_db): + """The vector guide's quick start sets quantization and addHierarchy in METADATA.""" + test_db.command("sql", "CREATE VERTEX TYPE QuickDoc") + test_db.command("sql", "CREATE PROPERTY QuickDoc.embedding ARRAY_OF_FLOATS") + test_db.command( + "sql", + """ + CREATE INDEX ON QuickDoc (embedding) + LSM_VECTOR + METADATA { + "dimensions": 3, + "similarity": "COSINE", + "quantization": "INT8", + "addHierarchy": true + } + """, + ) + meta = self._metadata(test_db, "QuickDoc[embedding]") + assert str(meta.quantizationType) == "INT8" + assert bool(meta.addHierarchy) is True + + +class TestSparseSearchOptionsSql: + """The sparse index forms the vector guide documents.""" + + @staticmethod + def _make_sparse(db): + db.command("sql", "CREATE DOCUMENT TYPE SparseOpt") + db.command("sql", "CREATE PROPERTY SparseOpt.id STRING") + db.command("sql", "CREATE PROPERTY SparseOpt.tokens ARRAY_OF_INTEGERS") + db.command("sql", "CREATE PROPERTY SparseOpt.weights ARRAY_OF_FLOATS") + + @staticmethod + def _insert(db): + with db.transaction(): + for doc_id, weight in (("a", 1.0), ("b", 0.5), ("c", 0.1)): + db.command( + "sql", + "INSERT INTO SparseOpt SET id = ?, tokens = [5], weights = [?]", + doc_id, + weight, + ) + + @staticmethod + def _search(db, options): + return db.query( + "sql", + "SELECT id, score FROM " # nosec B608 - test-owned option literals + "(SELECT expand(`vector.sparseNeighbors`('SparseOpt[tokens,weights]', " # nosec B608 + f"?, ?, 5, {options}))) ORDER BY score DESC", # nosec B608 + jtypes.JArray(jtypes.JInt)([5]), + arcadedb.to_java_float_array([1.0]), + ).to_list() + + def test_min_score_drops_low_scores(self, test_db): + self._make_sparse(test_db) + test_db.command( + "sql", + "CREATE INDEX ON SparseOpt (tokens, weights) LSM_SPARSE_VECTOR " + 'METADATA {"dimensions": 128, "weightQuantization": "FP32"}', + ) + self._insert(test_db) + + rows = self._search(test_db, "{ minScore: 0.4 }") + + assert [r["id"] for r in rows] == ["a", "b"] + assert rows[0]["score"] > rows[1]["score"] >= 0.4 + + def test_index_type_enum_creates_a_sparse_index(self, test_db): + self._make_sparse(test_db) + test_db.schema.create_index( + "SparseOpt", + ["tokens", "weights"], + index_type=arcadedb.IndexType.LSM_SPARSE_VECTOR, + ) + self._insert(test_db) + + rows = self._search(test_db, "{ minScore: 0.4 }") + + assert [r["id"] for r in rows] == ["a", "b"]