From 2191f948c3ce30ceeebc05438541c9fdd98d7200 Mon Sep 17 00:00:00 2001
From: Philip Meier <github.pmeier@posteo.de>
Date: Tue, 22 Apr 2025 23:46:10 +0200
Subject: [PATCH] fix: add missing pydantic model config compat (#2316)

Fixes #2315.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **Refactor**
- Enhanced query processing to maintain smooth functionality across
different dependency versions, ensuring improved stability and
performance.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 python/python/lancedb/query.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py
index 6db7e78a..235e0b16 100644
--- a/python/python/lancedb/query.py
+++ b/python/python/lancedb/query.py
@@ -28,6 +28,8 @@ import pyarrow.compute as pc
 import pyarrow.fs as pa_fs
 import pydantic
 
+from lancedb.pydantic import PYDANTIC_VERSION
+
 from . import __version__
 from .arrow import AsyncRecordBatchReader
 from .dependencies import pandas as pd
@@ -498,10 +500,14 @@ class Query(pydantic.BaseModel):
             )
         return query
 
-    class Config:
-        # This tells pydantic to allow custom types (needed for the `vector` query since
-        # pa.Array wouln't be allowed otherwise)
-        arbitrary_types_allowed = True
+    # This tells pydantic to allow custom types (needed for the `vector` query since
+    # pa.Array wouln't be allowed otherwise)
+    if PYDANTIC_VERSION.major < 2:  # Pydantic 1.x compat
+
+        class Config:
+            arbitrary_types_allowed = True
+    else:
+        model_config = {"arbitrary_types_allowed": True}
 
 
 class LanceQueryBuilder(ABC):