fix vector search tests

Author: tswast

bigframes/session/_io/bigquery/read_gbq_query.py

@@ -67,8 +67,11 @@ def create_dataframe_from_row_iterator(
     """
     pa_table = rows.to_arrow()
     bq_schema = list(rows.schema)
+    is_default_index = not index_col or isinstance(
+        index_col, bigframes.enums.DefaultIndexKind
+    )
 
-    if not index_col or isinstance(index_col, bigframes.enums.DefaultIndexKind):
+    if is_default_index:
         # We get a sequential index for free, so use that if no index is specified.
         # TODO(tswast): Use array_value.promote_offsets() instead once that node is
         # supported by the local engine.
@@ -81,6 +84,7 @@ def create_dataframe_from_row_iterator(
         index_columns = (index_col,)
         index_labels = (index_col,)
     else:
+        index_col = cast(Iterable[str], index_col)
         index_columns = tuple(index_col)
         index_labels = cast(Tuple[Optional[str], ...], tuple(index_col))
 
@@ -111,4 +115,7 @@ def create_dataframe_from_row_iterator(
     if columns:
         df = df[list(columns)]
 
+    if not is_default_index:
+        df = df.sort_index()
+
     return df

bigframes/session/loader.py

@@ -956,6 +956,7 @@ def read_gbq_query(
                 True if use_cache is None else use_cache
             )
 
+        _check_duplicates("columns", columns)
         index_cols = _to_index_cols(index_col)
         _check_index_col_param(index_cols, columns)
 

tests/system/small/bigquery/test_vector_search.py

@@ -157,80 +157,43 @@ def test_vector_search_basic_params_with_df():
     )
 
 
-def test_vector_search_different_params_with_query():
-    search_query = bpd.Series([[1.0, 2.0], [3.0, 5.2]])
-    vector_search_result = bbq.vector_search(
-        base_table="bigframes-dev.bigframes_tests_sys.base_table",
-        column_to_search="my_embedding",
-        query=search_query,
-        distance_type="cosine",
-        top_k=2,
-    ).to_pandas()  # type:ignore
-    expected = pd.DataFrame(
+def test_vector_search_different_params_with_query(session):
+    base_df = bpd.DataFrame(
         {
-            "0": [
-                np.array([1.0, 2.0]),
-                np.array([1.0, 2.0]),
-                np.array([3.0, 5.2]),
-                np.array([3.0, 5.2]),
-            ],
-            "id": [2, 1, 1, 2],
+            "id": [1, 2, 3, 4],
             "my_embedding": [
-                np.array([2.0, 4.0]),
-                np.array([1.0, 2.0]),
-                np.array([1.0, 2.0]),
-                np.array([2.0, 4.0]),
+                np.array([0.0, 1.0]),
+                np.array([1.0, 0.0]),
+                np.array([0.0, -1.0]),
+                np.array([-1.0, 0.0]),
             ],
-            "distance": [0.0, 0.0, 0.001777, 0.001777],
         },
-        index=pd.Index([0, 0, 1, 1], dtype="Int64"),
-    )
-    pd.testing.assert_frame_equal(
-        vector_search_result, expected, check_dtype=False, rtol=0.1
-    )
-
-
-def test_vector_search_df_with_query_column_to_search():
-    search_query = bpd.DataFrame(
-        {
-            "query_id": ["dog", "cat"],
-            "embedding": [[1.0, 2.0], [3.0, 5.2]],
-            "another_embedding": [[1.0, 2.5], [3.3, 5.2]],
-        }
-    )
-    vector_search_result = bbq.vector_search(
-        base_table="bigframes-dev.bigframes_tests_sys.base_table",
-        column_to_search="my_embedding",
-        query=search_query,
-        query_column_to_search="another_embedding",
-        top_k=2,
-    ).to_pandas()  # type:ignore
-    expected = pd.DataFrame(
-        {
-            "query_id": ["dog", "dog", "cat", "cat"],
-            "embedding": [
-                np.array([1.0, 2.0]),
-                np.array([1.0, 2.0]),
-                np.array([3.0, 5.2]),
-                np.array([3.0, 5.2]),
-            ],
-            "another_embedding": [
-                np.array([1.0, 2.5]),
-                np.array([1.0, 2.5]),
-                np.array([3.3, 5.2]),
-                np.array([3.3, 5.2]),
-            ],
-            "id": [1, 4, 2, 5],
-            "my_embedding": [
-                np.array([1.0, 2.0]),
-                np.array([1.0, 3.2]),
-                np.array([2.0, 4.0]),
-                np.array([5.0, 5.4]),
-            ],
-            "distance": [0.5, 0.7, 1.769181, 1.711724],
-        },
-        index=pd.Index([0, 0, 1, 1], dtype="Int64"),
-    )
-    pd.testing.assert_frame_equal(
-        vector_search_result, expected, check_dtype=False, rtol=0.1
+        session=session,
     )
+    base_table = base_df.to_gbq()
+    try:
+        search_query = bpd.Series([[0.75, 0.25], [-0.25, -0.75]], session=session)
+        vector_search_result = bbq.vector_search(
+            base_table=base_table,
+            column_to_search="my_embedding",
+            query=search_query,
+            distance_type="cosine",
+            top_k=2,
+        ).to_pandas()  # type:ignore
+        expected = pd.DataFrame(
+            {
+                "0": {np.int64(0): [0.75, 0.25], np.int64(1): [-0.25, -0.75]},
+                "id": {np.int64(0): 1, np.int64(1): 4},
+                "my_embedding": {np.int64(0): [0.0, 1.0], np.int64(1): [-1.0, 0.0]},
+                "distance": {
+                    np.int64(0): 0.683772233983162,
+                    np.int64(1): 0.683772233983162,
+                },
+            },
+            index=pd.Index([0, 0, 1, 1], dtype="Int64"),
+        )
+        pd.testing.assert_frame_equal(
+            vector_search_result, expected, check_dtype=False, rtol=0.1
+        )
+    finally:
+        session.bqclient.delete_table(base_table, not_found_ok=True)

tests/system/small/test_unordered.py

@@ -103,7 +103,7 @@ def test_unordered_mode_read_gbq(unordered_session):
         }
     )
     # Don't need ignore_order as there is only 1 row
-    assert_pandas_df_equal(df.to_pandas(), expected)
+    assert_pandas_df_equal(df.to_pandas(), expected, check_index_type=False)
 
 
 @pytest.mark.parametrize(