import json
import struct
import time
from typing import ClassVar, Dict, List
from unittest import mock
from unittest.mock import MagicMock, patch

import numpy as np
import orjson
import pytest
from pymilvus.client import entity_helper
from pymilvus.client.entity_helper import (
    convert_to_array,
    convert_to_array_of_vector,
    convert_to_json,
    convert_to_str_array,
    entity_to_array_arr,
    entity_to_field_data,
    entity_to_json_arr,
    entity_to_str_arr,
    entity_type_to_dtype,
    extract_array_row_data,
    extract_array_row_data_no_validity,
    extract_array_row_data_with_validity,
    extract_dynamic_field_from_result,
    extract_row_data_from_fields_data,
    extract_row_data_from_fields_data_v2,
    extract_struct_array_from_column_data,
    extract_vector_array_row_data,
    flush_vector_bytes,
    get_array_length,
    get_array_value_at_index,
    get_max_len_of_var_char,
    pack_field_value_to_field_data,
    sparse_proto_to_rows,
    sparse_rows_to_proto,
)
from pymilvus.client.types import DataType
from pymilvus.exceptions import DataNotMatchException, ParamError
from pymilvus.grpc_gen import schema_pb2
from pymilvus.grpc_gen import schema_pb2 as schema_types
from pymilvus.settings import Config
from scipy.sparse import csr_matrix

# Alias for backward compatibility
convert_to_entity = extract_row_data_from_fields_data_v2


class TestEntityHelperSparse:
    """Test entity_helper module functions"""

    @pytest.mark.parametrize(
        "valid_sparse_matrix",
        [
            [{0: 1.0, 5: 2.5, 10: 3.0}],  # list of one dict
            [{0: 1.0, 5: 2.5}, {10: 3.0, 15: 4.0}],  # list of dicts
            [{}, {10: 3.0, 15: 4.0}],  # list of dicts partial empty is allowed
            [[(1, 0.5), (10, 0.3)], [(2, 0.7), (20, 0.1)]],  # list of list
            [[("1", "0.5"), (10, 0.3)]],  # str representation of int
            csr_matrix(
                ([1.0, 2.0, 3.0], [0, 2, 3], [0, 2, 3, 3]), shape=(3, 4)
            ),  # scipy sparse matrix
            [
                csr_matrix([[1.0, 0, 2.0]]),
                csr_matrix([[0, 0, 3.0]]),
            ],  # list of scipy sparse matrices
        ],
    )
    def test_entity_is_sparse_matrix(self, valid_sparse_matrix: list):
        assert entity_helper.entity_is_sparse_matrix(valid_sparse_matrix) is True

    @pytest.mark.parametrize(
        "not_sparse_matrix",
        [
            [{"a": 1.0, "b": 2.0}],  # invalid dict for non-numeric keys
            [],  # empty
            [{0: 1.0}, "not a dict", {5: 2.0}],  # mixed lists
            None,
            123,
            "string",
            [1, 2, 3],
            [[1, 2, 3]],
            [[(1, 0.5, 0.2)]],
            [[(1, "invalid")]],
            [csr_matrix([[1, 0], [0, 1]])],  # list of multi-row is not sparse
        ],
    )
    def test_entity_isnot_sparse_matrix(self, not_sparse_matrix: any):
        assert entity_helper.entity_is_sparse_matrix(not_sparse_matrix) is False

    def test_get_input_num_rows_list(self):
        """Test getting number of rows from list input"""
        # Regular list
        data = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
        assert entity_helper.get_input_num_rows(data) == 3
        assert entity_helper.get_input_num_rows([1, 2, 3]) == 3
        assert entity_helper.get_input_num_rows({"a": 1, "b": 2}) == 2

        data = [[1, 2, 3]]
        assert entity_helper.get_input_num_rows(data) == 1

        data = []
        assert entity_helper.get_input_num_rows(data) == 0

        matrix = csr_matrix([[1, 0], [0, 1], [1, 1]])
        assert entity_helper.get_input_num_rows(matrix) == 3

        sparse_list = [{0: 1.0}, {5: 2.5}, {10: 3.0}]
        assert entity_helper.get_input_num_rows(sparse_list) == 3

        data = np.array([[1, 2, 3], [4, 5, 6]])
        assert entity_helper.get_input_num_rows(data) == 2

    @pytest.mark.parametrize(
        "sparse_list", [[{0: 1.0, 2: 2.0}, {2: 3.0}], csr_matrix([[1, 0, 2], [0, 0, 3]])]
    )
    def test_sparse_rows_to_proto_dict(self, sparse_list: any):
        """Test converting sparse rows to protobuf format"""
        proto = entity_helper.sparse_rows_to_proto(sparse_list)
        assert isinstance(proto, schema_pb2.SparseFloatArray)
        assert len(proto.contents) == 2
        assert proto.dim == 3

    def test_sparse_proto_to_rows(self):
        """Test converting protobuf sparse vectors back to rows"""
        # Create a mock sparse proto
        proto = schema_pb2.SparseFloatArray(dim=100)

        # Add some sparse vectors in binary format
        # Format: pairs of (uint32 index, float32 value)
        vec1_data = b""
        for idx, val in [(0, 1.0), (5, 2.5), (10, 3.0)]:
            vec1_data += struct.pack("I", idx) + struct.pack("f", val)
        proto.contents.append(vec1_data)

        vec2_data = b""
        for idx, val in [(15, 4.0), (20, 5.0)]:
            vec2_data += struct.pack("I", idx) + struct.pack("f", val)
        proto.contents.append(vec2_data)

        # Convert back to rows
        rows = entity_helper.sparse_proto_to_rows(proto, 0, 2)
        assert len(rows) == 2
        assert rows[0] == {0: 1.0, 5: 2.5, 10: 3.0}
        assert rows[1] == {15: 4.0, 20: 5.0}

    def test_sparse_proto_to_rows_with_range(self):
        """Test converting specific range of sparse vectors"""
        proto = schema_pb2.SparseFloatArray(dim=100)

        # Add multiple sparse vectors in binary format
        for i in range(5):
            vec_data = struct.pack("I", i) + struct.pack("f", float(i))
            proto.contents.append(vec_data)

        # Get middle range
        rows = entity_helper.sparse_proto_to_rows(proto, 1, 4)
        assert len(rows) == 3
        assert rows[0] == {1: 1.0}
        assert rows[1] == {2: 2.0}
        assert rows[2] == {3: 3.0}

    def test_convert_to_json_nested(self):
        """Test JSON conversion with nested structures"""
        obj = {
            "level1": {
                "level2": {
                    "array": np.array([[1, 2], [3, 4]]),
                    "list": [np.int32(1), np.float64(2.5)],
                }
            }
        }
        result = entity_helper.convert_to_json(obj)
        parsed = json.loads(result)
        assert parsed["level1"]["level2"]["array"] == [[1, 2], [3, 4]]

    @pytest.mark.parametrize(
        "data",
        [
            {"key": "value", "number": 42},  # dict
            {"outer": {"inner": "value"}},  # nested dict
            [1, 2, 3, "four"],  # list
            [{"a": 1}, {"b": 2}],  # list of dict
            None,
            pytest.param({"array": np.array([1, 2, 3])}, marks=pytest.mark.xfail(reason="fix me")),
            {"int": np.int64(42), "float": np.float32(3.14), "bool": np.bool_(True)},
            [{"val": np.int64(10)}, {"val": np.float32(3.14)}],
        ],
    )
    def test_convert_to_json_dict(self, data: dict):
        """Test JSON conversion for dict input"""
        result = entity_helper.convert_to_json(data)
        assert isinstance(result, bytes)
        assert json.loads(result.decode()) == data

    @pytest.mark.parametrize(
        "json_string,expected",
        [
            ('{"key": "value", "number": 42}', {"key": "value", "number": 42}),
            ('{"nested": {"inner": "value"}}', {"nested": {"inner": "value"}}),
            ('[1, 2, 3, "four"]', [1, 2, 3, "four"]),
            ('{"name": "Alice", "age": 30}', {"name": "Alice", "age": 30}),
            ("null", None),
            ("true", True),
            ("false", False),
            ("123", 123),
            ('"simple string"', "simple string"),
        ],
    )
    def test_convert_to_json_string_valid(self, json_string: str, expected):
        """Test JSON conversion for valid JSON string input"""
        result = entity_helper.convert_to_json(json_string)
        assert isinstance(result, bytes)
        # Verify the result is valid JSON
        parsed = json.loads(result.decode())
        assert parsed == expected

    def test_convert_to_json_from_json_dumps(self):
        """Test JSON conversion from json.dumps() output"""
        original_dict = {"key": "value", "count": 100, "nested": {"inner": "data"}}
        json_string = json.dumps(original_dict)

        result = entity_helper.convert_to_json(json_string)
        assert isinstance(result, bytes)
        parsed = json.loads(result.decode())
        assert parsed == original_dict

    @pytest.mark.parametrize(
        "invalid_json_string",
        [
            "not a json string",
            '{"invalid": }',
            '{"key": "value"',  # missing closing brace
            "{'key': 'value'}",  # single quotes not valid in JSON
            "{key: value}",  # unquoted keys
            "undefined",
            "{,}",
        ],
    )
    def test_convert_to_json_string_invalid(self, invalid_json_string: str):
        """Test JSON conversion rejects invalid JSON strings"""

        with pytest.raises(DataNotMatchException) as exc_info:
            entity_helper.convert_to_json(invalid_json_string)

        # Verify error message contains the invalid JSON string
        error_message = str(exc_info.value)
        assert "Invalid JSON string" in error_message
        # Verify the original input string is in the error message
        assert invalid_json_string in error_message or invalid_json_string[:50] in error_message

    def test_convert_to_json_string_with_non_string_keys(self):
        """Test JSON conversion rejects JSON strings with non-string keys in dict"""

        # This is actually not possible in standard JSON, as JSON object keys are always strings
        # But we can test that dict validation still works
        invalid_dict = {1: "value", 2: "another"}

        with pytest.raises(DataNotMatchException) as exc_info:
            entity_helper.convert_to_json(invalid_dict)

        error_message = str(exc_info.value)
        assert "JSON" in error_message

    def test_convert_to_json_long_invalid_string_truncated(self):
        """Test that long invalid JSON strings are truncated in error messages"""

        # Create a long invalid JSON string
        long_invalid_json = "invalid json " * 50  # > 200 characters

        with pytest.raises(DataNotMatchException) as exc_info:
            entity_helper.convert_to_json(long_invalid_json)

        error_message = str(exc_info.value)
        assert "Invalid JSON string" in error_message
        # Should contain truncated version with "..."
        assert "..." in error_message

    def test_convert_to_json_deeply_nested_dict(self):
        """Test JSON conversion with deeply nested dictionary to avoid recursion limit"""

        def create_deep_dict(depth):
            """Create a deeply nested dictionary with specified depth"""
            result = {"id_0": "-501"}
            for i in range(1, depth + 1):
                result = {"id_" + str(i): result}
            return result

        # Test with 512 levels (tests fallback from orjson to standard json library)
        # orjson fails at ~500 levels, so 512 will trigger fallback to json.dumps
        deep_dict = create_deep_dict(512)

        # Should not raise RecursionError
        result = entity_helper.convert_to_json(deep_dict)
        assert isinstance(result, bytes)

        # Verify the structure is preserved
        parsed = json.loads(result.decode())
        current = parsed
        for i in range(512, 0, -1):
            assert f"id_{i}" in current, f"Missing key id_{i} at depth {512 - i + 1}"
            current = current[f"id_{i}"]
        assert current == {"id_0": "-501"}, "Final nested value mismatch"

    def test_pack_field_value_to_field_data(self):
        """Test packing field values into field data protobuf"""
        # Test with scalar field
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.INT64
        field_info = {"name": "test_field"}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        value = 42
        entity_helper.pack_field_value_to_field_data(
            value, field_data, field_info, vector_bytes_cache
        )

        assert len(field_data.scalars.long_data.data) == 1
        assert field_data.scalars.long_data.data[0] == value

    def test_pack_field_value_to_field_data_vectors(self):
        """Test packing vector field values"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.FLOAT_VECTOR
        field_info = {"name": "vector_field"}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        value = [1.0, 2.0, 3.0, 4.0]

        entity_helper.pack_field_value_to_field_data(
            value, field_data, field_info, vector_bytes_cache
        )

        assert field_data.vectors.dim == 4
        assert list(field_data.vectors.float_vector.data) == value

    def test_extract_array_row_data(self):
        """Test extracting array data from protobuf"""
        # Create field data with array
        field_data = schema_pb2.FieldData()
        field_data.scalars.array_data.element_type = DataType.INT64

        # Add array data for index 0
        scalar_field = schema_pb2.ScalarField()
        scalar_field.long_data.data.extend([0, 1, 2])
        field_data.scalars.array_data.data.append(scalar_field)

        result = entity_helper.extract_array_row_data(field_data, 0)
        assert result == [0, 1, 2]

    def test_extract_array_row_data_string(self):
        """Test extracting string array data"""
        # Create field data with array
        field_data = schema_pb2.FieldData()
        field_data.scalars.array_data.element_type = DataType.VARCHAR

        # Add array data for index 0
        scalar_field = schema_pb2.ScalarField()
        scalar_field.string_data.data.extend(["str_0", "str_1"])
        field_data.scalars.array_data.data.append(scalar_field)

        result = entity_helper.extract_array_row_data(field_data, 0)
        assert result == ["str_0", "str_1"]

    def test_extract_array_row_data_bool(self):
        """Test extracting boolean array data"""
        # Create field data with array
        field_data = schema_pb2.FieldData()
        field_data.scalars.array_data.element_type = DataType.BOOL

        # Add array data for index 0
        scalar_field = schema_pb2.ScalarField()
        scalar_field.bool_data.data.extend([True, False, True])
        field_data.scalars.array_data.data.append(scalar_field)

        result = entity_helper.extract_array_row_data(field_data, 0)
        assert result == [True, False, True]

    def test_extract_array_row_data_float(self):
        """Test extracting float array data"""
        # Create field data with array
        field_data = schema_pb2.FieldData()
        field_data.scalars.array_data.element_type = DataType.FLOAT

        # Add array data for index 0
        scalar_field = schema_pb2.ScalarField()
        scalar_field.float_data.data.extend([1.1, 2.2, 3.3])
        field_data.scalars.array_data.data.append(scalar_field)

        result = entity_helper.extract_array_row_data(field_data, 0)
        assert result == pytest.approx([1.1, 2.2, 3.3])

    def test_extract_array_row_data_double(self):
        """Test extracting double array data"""
        # Create field data with array
        field_data = schema_pb2.FieldData()
        field_data.scalars.array_data.element_type = DataType.DOUBLE

        # Add array data for index 0
        scalar_field = schema_pb2.ScalarField()
        scalar_field.double_data.data.extend([1.11111, 2.22222, 3.33333])
        field_data.scalars.array_data.data.append(scalar_field)

        result = entity_helper.extract_array_row_data(field_data, 0)
        assert result == pytest.approx([1.11111, 2.22222, 3.33333])

    def test_extract_array_row_data_invalid_type(self):
        """Test extracting array data with invalid type"""
        # Create field data with array
        field_data = schema_pb2.FieldData()
        field_data.scalars.array_data.element_type = 999  # Invalid type

        # Add array data for index 0
        scalar_field = schema_pb2.ScalarField()
        field_data.scalars.array_data.data.append(scalar_field)

        result = entity_helper.extract_array_row_data(field_data, 0)
        assert result is None  # Returns None for unknown types


class TestEntityHelperExtended:
    def test_sparse_rows_to_proto_invalid_index(self):
        """Test error handling for invalid indices"""
        # Negative index
        with pytest.raises(ParamError, match="sparse vector index must be positive"):
            sparse_rows_to_proto([{-1: 0.5}])

        # Index too large
        with pytest.raises(ParamError, match="sparse vector index must be positive"):
            sparse_rows_to_proto([{2**32: 0.5}])

    def test_sparse_rows_to_proto_nan_value(self):
        """Test error handling for NaN values"""
        with pytest.raises(ParamError, match="sparse vector value must not be NaN"):
            sparse_rows_to_proto([{1: float("nan")}])

    def test_sparse_rows_to_proto_invalid_input(self):
        """Test error handling for invalid input"""
        with pytest.raises(ParamError, match="input must be a sparse matrix"):
            sparse_rows_to_proto("invalid")

    def test_sparse_proto_to_rows_invalid(self):
        """Test error handling for invalid proto"""
        with pytest.raises(ParamError, match="Vector must be a sparse float vector"):
            sparse_proto_to_rows("invalid")

    def test_entity_type_to_dtype(self):
        """Test converting entity type to dtype"""
        # Integer type
        assert entity_type_to_dtype(1) == 1
        assert entity_type_to_dtype(DataType.INT64) == DataType.INT64

        # We can't test string conversion without knowing exact protobuf enum names
        # Let's just test invalid type
        with pytest.raises(ParamError, match="invalid entity type"):
            entity_type_to_dtype([])

    def test_get_max_len_of_var_char(self):
        """Test getting max length of varchar field"""
        # With params
        field_info = {"params": {Config.MaxVarCharLengthKey: 100}}
        assert get_max_len_of_var_char(field_info) == 100

        # Without params - use default
        field_info = {}
        assert get_max_len_of_var_char(field_info) == Config.MaxVarCharLength

        # Partial params
        field_info = {"params": {}}
        assert get_max_len_of_var_char(field_info) == Config.MaxVarCharLength

    def test_convert_to_str_array(self):
        """Test converting to string array"""
        field_info = {"name": "test_field", "params": {Config.MaxVarCharLengthKey: 10}}

        # Valid strings
        result = convert_to_str_array(["hello", "world"], field_info)
        assert result == ["hello", "world"]

        # String exceeding max length
        with pytest.raises(ParamError, match="length of string exceeds max length"):
            convert_to_str_array(["this is too long"], field_info)

        # Non-string input
        with pytest.raises(ParamError, match="expects string input"):
            convert_to_str_array([123], field_info)

        # Without check
        result = convert_to_str_array([123, "test"], field_info, check=False)
        assert len(result) == 2

    @patch("pymilvus.client.entity_helper.Config")
    def test_convert_to_str_array_with_encoding(self, mock_config):
        """Test string array conversion with different encoding"""
        mock_config.EncodeProtocol = "latin-1"
        mock_config.MaxVarCharLengthKey = "max_length"
        mock_config.MaxVarCharLength = 65535

        field_info = {"name": "test_field"}
        strings = ["hello", "world"]
        result = convert_to_str_array(strings, field_info, check=False)
        assert len(result) == 2

    def test_entity_to_str_arr(self):
        """Test entity_to_str_arr wrapper function"""
        field_info = {"name": "test_field", "params": {Config.MaxVarCharLengthKey: 20}}
        result = entity_to_str_arr(["test1", "test2"], field_info)
        assert result == ["test1", "test2"]

    def test_extract_array_row_data(self):
        """Test extracting array row data"""
        # INT64 array
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.ARRAY
        array_data = schema_pb2.ArrayArray()

        # Add INT64 data
        int_array = schema_pb2.LongArray()
        int_array.data.extend([1, 2, 3])
        array_data.data.append(schema_pb2.ScalarField(long_data=int_array))
        int_array2 = schema_pb2.LongArray()
        int_array2.data.extend([4, 5])
        array_data.data.append(schema_pb2.ScalarField(long_data=int_array2))

        array_data.element_type = DataType.INT64
        field_data.scalars.array_data.CopyFrom(array_data)

        result = extract_array_row_data(field_data, 0)
        assert result == [1, 2, 3]
        result = extract_array_row_data(field_data, 1)
        assert result == [4, 5]

    def test_extract_array_row_data_string(self):
        """Test extracting string array data"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.ARRAY
        array_data = schema_pb2.ArrayArray()

        # Add string data
        str_array = schema_pb2.StringArray()
        str_array.data.extend(["hello", "world"])
        array_data.data.append(schema_pb2.ScalarField(string_data=str_array))

        array_data.element_type = DataType.VARCHAR
        field_data.scalars.array_data.CopyFrom(array_data)

        result = extract_array_row_data(field_data, 0)
        assert result == ["hello", "world"]

    def test_extract_array_row_data_bool(self):
        """Test extracting boolean array data"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.ARRAY
        array_data = schema_pb2.ArrayArray()

        # Add bool data
        bool_array = schema_pb2.BoolArray()
        bool_array.data.extend([True, False, True])
        array_data.data.append(schema_pb2.ScalarField(bool_data=bool_array))

        array_data.element_type = DataType.BOOL
        field_data.scalars.array_data.CopyFrom(array_data)

        result = extract_array_row_data(field_data, 0)
        assert result == [True, False, True]

    def test_extract_array_row_data_float(self):
        """Test extracting float array data"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.ARRAY
        array_data = schema_pb2.ArrayArray()

        # Add float data
        float_array = schema_pb2.FloatArray()
        float_array.data.extend([1.1, 2.2, 3.3])
        array_data.data.append(schema_pb2.ScalarField(float_data=float_array))

        array_data.element_type = DataType.FLOAT
        field_data.scalars.array_data.CopyFrom(array_data)

        result = extract_array_row_data(field_data, 0)
        assert len(result) == 3
        assert result[0] == pytest.approx(1.1)

    def test_extract_array_row_data_double(self):
        """Test extracting double array data"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.ARRAY
        array_data = schema_pb2.ArrayArray()

        # Add double data
        double_array = schema_pb2.DoubleArray()
        double_array.data.extend([1.11111, 2.22222])
        array_data.data.append(schema_pb2.ScalarField(double_data=double_array))

        array_data.element_type = DataType.DOUBLE
        field_data.scalars.array_data.CopyFrom(array_data)

        result = extract_array_row_data(field_data, 0)
        assert len(result) == 2

    def test_extract_array_row_data_invalid_type(self):
        """Test error handling for invalid array element type"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.ARRAY
        array_data = schema_pb2.ArrayArray()
        array_data.element_type = 999  # Invalid type
        # Add at least one empty element to avoid index error
        array_data.data.append(schema_pb2.ScalarField())
        field_data.scalars.array_data.CopyFrom(array_data)

        assert extract_array_row_data(field_data, 0) is None

    def test_entity_to_array_arr(self):
        """Test converting entity to array array"""
        field_info = {"name": "array_field", "element_type": DataType.INT64}

        # List of lists
        data = [[1, 2, 3], [4, 5], [6]]
        result = entity_to_array_arr(data, field_info)
        assert len(result) == 3
        assert result[0].long_data.data == [1, 2, 3]
        assert result[1].long_data.data == [4, 5]

    def test_entity_to_array_arr_string(self):
        """Test converting string arrays"""
        field_info = {"name": "array_field", "element_type": DataType.VARCHAR}

        data = [["hello", "world"], ["foo"]]
        result = entity_to_array_arr(data, field_info)
        assert len(result) == 2
        assert list(result[0].string_data.data) == ["hello", "world"]

    def test_entity_to_array_arr_invalid_type(self):
        """Test error handling for invalid element type"""
        field_info = {"name": "array_field", "element_type": 999}

        with pytest.raises(ParamError, match="Unsupported element type"):
            entity_to_array_arr([[1, 2]], field_info)

    def test_pack_field_value_to_field_data(self):
        """Test packing field values to field data"""
        # pack_field_value_to_field_data takes different parameters
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.FLOAT_VECTOR
        field_data.field_name = "vector_field"
        field_info = {"name": "vector_field"}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Pack a single vector
        pack_field_value_to_field_data(
            np.array([1.0, 2.0]), field_data, field_info, vector_bytes_cache
        )

        # Check the result
        assert field_data.type == DataType.FLOAT_VECTOR
        assert field_data.vectors.dim == 2

    def test_pack_field_value_to_field_data_sparse(self):
        """Test packing sparse vectors"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.SPARSE_FLOAT_VECTOR
        field_data.field_name = "sparse_field"
        field_info = {"name": "sparse_field"}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Pack a single sparse vector
        sparse_data = {1: 0.5, 10: 0.3}
        pack_field_value_to_field_data(sparse_data, field_data, field_info, vector_bytes_cache)

        # Check the result
        assert field_data.type == DataType.SPARSE_FLOAT_VECTOR
        assert len(field_data.vectors.sparse_float_vector.contents) == 1

    def test_pack_field_value_to_field_data_scalars(self):
        """Test packing scalar field values"""
        # Test INT64
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.INT64
        field_data.field_name = "int_field"
        field_info = {"name": "int_field"}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        pack_field_value_to_field_data(42, field_data, field_info, vector_bytes_cache)

        assert field_data.type == DataType.INT64
        assert field_data.scalars.long_data.data[0] == 42

    def test_extract_field_info(self):
        """Test extracting primary field from schema"""
        # Create schema with primary field
        fields_info = [{"name": "id", "is_primary": True}, {"name": "vector", "is_primary": False}]

        # Test field extraction logic
        # This tests the field extraction patterns used in entity_helper
        for field in fields_info:
            if field.get("is_primary"):
                assert field["name"] == "id"

    def test_extract_dynamic_field_from_result(self):
        """Test extracting dynamic field from result"""

        # Create actual result object with fields_data and output_fields attributes
        class ActualResult:
            def __init__(self, fields_data: List, output_fields: List):
                self.fields_data = fields_data
                self.output_fields = output_fields

        # Test with dynamic field
        dynamic_field_data = schema_pb2.FieldData()
        dynamic_field_data.is_dynamic = True
        dynamic_field_data.field_name = "$meta"

        regular_field_data = schema_pb2.FieldData()
        regular_field_data.is_dynamic = False
        regular_field_data.field_name = "id"

        # Create result with dynamic field - extra_field comes before $meta
        result = ActualResult(
            fields_data=[regular_field_data, dynamic_field_data],
            output_fields=["id", "extra_field", "another_extra", "$meta"],
        )

        dynamic_field_name, dynamic_fields = extract_dynamic_field_from_result(result)
        assert dynamic_field_name == "$meta"
        # When $meta is found in output_fields, dynamic_fields gets cleared
        assert len(dynamic_fields) == 0

        # Test with no dynamic field
        result_no_dynamic = ActualResult(
            fields_data=[regular_field_data], output_fields=["id", "extra_field"]
        )

        dynamic_field_name, dynamic_fields = extract_dynamic_field_from_result(result_no_dynamic)
        assert dynamic_field_name is None
        assert "extra_field" in dynamic_fields
        assert "id" not in dynamic_fields

        # Test with dynamic field NOT in output_fields (dynamic_fields preserved)
        result_meta_not_in_output = ActualResult(
            fields_data=[regular_field_data, dynamic_field_data],
            output_fields=["id", "extra_field", "another_extra"],
        )

        dynamic_field_name, dynamic_fields = extract_dynamic_field_from_result(
            result_meta_not_in_output
        )
        assert dynamic_field_name == "$meta"
        assert "extra_field" in dynamic_fields
        assert "another_extra" in dynamic_fields
        assert "id" not in dynamic_fields  # id is a regular field
        assert len(dynamic_fields) == 2

    def test_data_validation(self):
        """Test data validation patterns"""
        # Test length validation
        data1 = [1, 2, 3]
        data2 = [4, 5, 6]
        assert len(data1) == len(data2)  # Valid same length

        # Length mismatch should be caught
        data3 = [1, 2]
        data4 = [3, 4, 5]
        assert len(data3) != len(data4)  # Invalid different length

    def test_schema_validation(self):
        """Test schema validation patterns"""
        fields_info = [
            {"name": "id", "is_primary": True, "auto_id": False},
            {"name": "vector", "is_primary": False},
        ]

        # Valid data matches schema
        data = {"id": [1, 2], "vector": [[1.0, 2.0], [3.0, 4.0]]}
        for field in fields_info:
            if not field.get("auto_id"):
                assert field["name"] in data  # Field should be in data

        # Test auto_id logic
        fields_info[0]["auto_id"] = True
        # When auto_id is True, id field should not be required in data

    def test_convert_to_entity(self):
        """Test converting field data to entities"""
        # Create field data
        field1 = schema_pb2.FieldData()
        field1.field_name = "id"
        field1.type = DataType.INT64
        long_array = schema_pb2.LongArray()
        long_array.data.extend([1, 2, 3])
        field1.scalars.long_data.CopyFrom(long_array)

        field2 = schema_pb2.FieldData()
        field2.field_name = "name"
        field2.type = DataType.VARCHAR
        str_array = schema_pb2.StringArray()
        str_array.data.extend(["a", "b", "c"])
        field2.scalars.string_data.CopyFrom(str_array)

        # extract_row_data_from_fields_data_v2 takes a single field_data and list of entity dicts
        # We need to create empty entity rows first
        entity_rows = [{} for _ in range(3)]

        # Process each field separately
        convert_to_entity(field1, entity_rows)
        convert_to_entity(field2, entity_rows)

        assert len(entity_rows) == 3
        assert entity_rows[0] == {"id": 1, "name": "a"}
        assert entity_rows[1] == {"id": 2, "name": "b"}
        assert entity_rows[2] == {"id": 3, "name": "c"}

    def test_entity_to_field_data(self):
        """Test converting entity to field data"""
        # entity_to_field_data expects a dict with specific structure
        entity = {"name": "test_field", "type": DataType.INT64, "values": [1, 2, 3, 4, 5]}
        field_info = {"name": "test_field"}

        result = entity_to_field_data(entity, field_info, 5)

        assert result.field_name == "test_field"
        assert result.type == DataType.INT64
        assert list(result.scalars.long_data.data) == [1, 2, 3, 4, 5]


class TestNullableVectorSupport:
    """Test nullable vector support in entity_helper module for all 6 vector types"""

    # Vector type configurations for parametrized tests
    VECTOR_TYPE_CONFIGS: ClassVar = [
        {
            "dtype": DataType.FLOAT_VECTOR,
            "name": "float_vector",
            "values": [[1.0, 2.0, 3.0, 4.0], None, [5.0, 6.0, 7.0, 8.0]],
            "dim": 4,
            "get_data_len": lambda fd: len(fd.vectors.float_vector.data),
            "expected_data_len": 8,  # 2 valid vectors * 4 dim
        },
        {
            "dtype": DataType.BINARY_VECTOR,
            "name": "binary_vector",
            "values": [b"\x01\x02\x03\x04", None, b"\x05\x06\x07\x08"],
            "dim": 32,  # 4 bytes * 8 bits
            "get_data_len": lambda fd: len(fd.vectors.binary_vector),
            "expected_data_len": 8,  # 2 valid vectors * 4 bytes
        },
        {
            "dtype": DataType.FLOAT16_VECTOR,
            "name": "float16_vector",
            "values": [
                np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16).tobytes(),
                None,
                np.array([5.0, 6.0, 7.0, 8.0], dtype=np.float16).tobytes(),
            ],
            "dim": 4,
            "get_data_len": lambda fd: len(fd.vectors.float16_vector),
            "expected_data_len": 16,  # 2 valid vectors * 4 dim * 2 bytes
        },
        {
            "dtype": DataType.BFLOAT16_VECTOR,
            "name": "bfloat16_vector",
            "values": [
                b"\x00\x3f\x00\x40\x00\x40\x00\x40",
                None,
                b"\x00\x40\x00\x40\x00\x40\x00\x40",
            ],
            "dim": 4,
            "get_data_len": lambda fd: len(fd.vectors.bfloat16_vector),
            "expected_data_len": 16,  # 2 valid vectors * 4 dim * 2 bytes
        },
        {
            "dtype": DataType.SPARSE_FLOAT_VECTOR,
            "name": "sparse_float_vector",
            "values": [{1: 0.5, 10: 0.3}, None, {5: 0.8}],
            "dim": None,
            "get_data_len": lambda fd: len(fd.vectors.sparse_float_vector.contents),
            "expected_data_len": 2,  # 2 valid sparse vectors
        },
        {
            "dtype": DataType.INT8_VECTOR,
            "name": "int8_vector",
            "values": [
                np.array([1, 2, 3, 4], dtype=np.int8).tobytes(),
                None,
                np.array([5, 6, 7, 8], dtype=np.int8).tobytes(),
            ],
            "dim": 4,
            "get_data_len": lambda fd: len(fd.vectors.int8_vector),
            "expected_data_len": 8,  # 2 valid vectors * 4 dim * 1 byte
        },
    ]

    @pytest.mark.parametrize(
        "config", VECTOR_TYPE_CONFIGS, ids=[c["name"] for c in VECTOR_TYPE_CONFIGS]
    )
    def test_entity_to_field_data_nullable_vector(self, config):
        """Test entity_to_field_data with nullable vector containing None values"""
        entity = {
            "name": f"nullable_{config['name']}",
            "type": config["dtype"],
            "values": config["values"],
        }
        field_info = {"name": f"nullable_{config['name']}", "nullable": True}
        if config["dim"]:
            field_info["params"] = {"dim": config["dim"]}

        result = entity_to_field_data(entity, field_info, 3)

        assert result.field_name == f"nullable_{config['name']}"
        assert result.type == config["dtype"]
        assert list(result.valid_data) == [True, False, True]
        assert config["get_data_len"](result) == config["expected_data_len"]

    @pytest.mark.parametrize(
        "config", VECTOR_TYPE_CONFIGS, ids=[c["name"] for c in VECTOR_TYPE_CONFIGS]
    )
    def test_entity_to_field_data_nullable_vector_all_none(self, config):
        """Test entity_to_field_data with nullable vector where all values are None"""
        entity = {
            "name": f"nullable_{config['name']}",
            "type": config["dtype"],
            "values": [None, None, None],
        }
        field_info = {"name": f"nullable_{config['name']}", "nullable": True}
        if config["dim"]:
            field_info["params"] = {"dim": config["dim"]}

        result = entity_to_field_data(entity, field_info, 3)

        assert list(result.valid_data) == [False, False, False]
        assert config["get_data_len"](result) == 0

    @pytest.mark.parametrize(
        "dtype,name",
        [
            (DataType.FLOAT_VECTOR, "float_vector"),
            (DataType.BINARY_VECTOR, "binary_vector"),
            (DataType.FLOAT16_VECTOR, "float16_vector"),
            (DataType.BFLOAT16_VECTOR, "bfloat16_vector"),
            (DataType.SPARSE_FLOAT_VECTOR, "sparse_float_vector"),
            (DataType.INT8_VECTOR, "int8_vector"),
        ],
    )
    def test_pack_field_value_nullable_vector_none(self, dtype, name):
        """Test pack_field_value_to_field_data with None value for nullable vector"""
        field_data = schema_pb2.FieldData()
        field_data.type = dtype
        field_data.field_name = f"nullable_{name}"
        field_info = {"name": f"nullable_{name}", "nullable": True}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        pack_field_value_to_field_data(None, field_data, field_info, vector_bytes_cache)

        # Verify no data added for each type
        if dtype == DataType.FLOAT_VECTOR:
            assert len(field_data.vectors.float_vector.data) == 0
        elif dtype == DataType.BINARY_VECTOR:
            assert len(field_data.vectors.binary_vector) == 0
        elif dtype == DataType.FLOAT16_VECTOR:
            assert len(field_data.vectors.float16_vector) == 0
        elif dtype == DataType.BFLOAT16_VECTOR:
            assert len(field_data.vectors.bfloat16_vector) == 0
        elif dtype == DataType.SPARSE_FLOAT_VECTOR:
            assert len(field_data.vectors.sparse_float_vector.contents) == 0
        elif dtype == DataType.INT8_VECTOR:
            assert len(field_data.vectors.int8_vector) == 0

    def test_extract_row_data_nullable_float_vector(self):
        """Test extracting nullable float vector from field data"""

        field_data = schema_pb2.FieldData()
        field_data.type = DataType.FLOAT_VECTOR
        field_data.field_name = "nullable_vector"
        field_data.vectors.dim = 4
        field_data.vectors.float_vector.data.extend([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
        field_data.valid_data.extend([True, False, True])

        result0 = extract_row_data_from_fields_data([field_data], 0)
        assert result0["nullable_vector"] == [1.0, 2.0, 3.0, 4.0]

        result1 = extract_row_data_from_fields_data([field_data], 1)
        assert result1["nullable_vector"] is None

        result2 = extract_row_data_from_fields_data([field_data], 2)
        assert result2["nullable_vector"] == [5.0, 6.0, 7.0, 8.0]

    def test_extract_row_data_nullable_sparse_vector(self):
        """Test extracting nullable sparse vector from field data"""

        field_data = schema_pb2.FieldData()
        field_data.type = DataType.SPARSE_FLOAT_VECTOR
        field_data.field_name = "nullable_sparse"
        vec1_data = struct.pack("I", 1) + struct.pack("f", 0.5)
        vec2_data = struct.pack("I", 5) + struct.pack("f", 0.8)
        field_data.vectors.sparse_float_vector.contents.extend([vec1_data, vec2_data])
        field_data.valid_data.extend([True, False, True])

        result0 = extract_row_data_from_fields_data([field_data], 0)
        assert result0["nullable_sparse"] == pytest.approx({1: 0.5})

        result1 = extract_row_data_from_fields_data([field_data], 1)
        assert result1["nullable_sparse"] is None

        result2 = extract_row_data_from_fields_data([field_data], 2)
        assert result2["nullable_sparse"] == pytest.approx({5: 0.8})

    def test_extract_row_data_non_nullable_vector_uses_logical_index(self):
        """Test that non-nullable vectors still use logical index directly"""

        field_data = schema_pb2.FieldData()
        field_data.type = DataType.FLOAT_VECTOR
        field_data.field_name = "regular_vector"
        field_data.vectors.dim = 4
        field_data.vectors.float_vector.data.extend(
            [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]
        )

        result0 = extract_row_data_from_fields_data([field_data], 0)
        assert result0["regular_vector"] == [1.0, 2.0, 3.0, 4.0]

        result1 = extract_row_data_from_fields_data([field_data], 1)
        assert result1["regular_vector"] == [5.0, 6.0, 7.0, 8.0]

        result2 = extract_row_data_from_fields_data([field_data], 2)
        assert result2["regular_vector"] == [9.0, 10.0, 11.0, 12.0]

    def test_pack_field_value_to_field_data_int8_vector(self):
        """Test packing int8 vector field values"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.INT8_VECTOR
        field_data.field_name = "int8_vector_field"
        field_info = {"name": "int8_vector_field", "params": {"dim": 768}}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Pack a single int8 vector
        vector = np.array([i % 128 - 64 for i in range(768)], dtype=np.int8)
        pack_field_value_to_field_data(vector, field_data, field_info, vector_bytes_cache)

        # Flush to merge collected bytes
        flush_vector_bytes(field_data, vector_bytes_cache)

        # Check the result
        assert field_data.type == DataType.INT8_VECTOR
        assert field_data.vectors.dim == 768
        assert len(field_data.vectors.int8_vector) == 768

        # Verify data correctness
        expected_bytes = vector.tobytes()
        assert field_data.vectors.int8_vector == expected_bytes

    def test_pack_field_value_to_field_data_int8_vector_multiple(self):
        """Test packing multiple int8 vectors to verify memory optimization"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.INT8_VECTOR
        field_data.field_name = "int8_vector_field"
        field_info = {"name": "int8_vector_field", "params": {"dim": 768}}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Pack multiple vectors (simulating batch insert)
        num_vectors = 1000
        vectors = []
        for i in range(num_vectors):
            vector = np.array([(i + j) % 128 - 64 for j in range(768)], dtype=np.int8)
            vectors.append(vector)
            pack_field_value_to_field_data(vector, field_data, field_info, vector_bytes_cache)

        # Flush to merge all collected bytes
        flush_vector_bytes(field_data, vector_bytes_cache)

        # Verify final result
        assert field_data.vectors.dim == 768
        expected_total_size = num_vectors * 768
        assert len(field_data.vectors.int8_vector) == expected_total_size

        # Verify data correctness for sample vectors
        for idx in [0, 100, 500, 999]:
            expected_bytes = vectors[idx].tobytes()
            actual_bytes = field_data.vectors.int8_vector[idx * 768 : (idx + 1) * 768]
            assert expected_bytes == actual_bytes, f"Vector {idx} data mismatch"

    def test_flush_int8_vector_bytes(self):
        """Test flush_vector_bytes function for int8 vectors"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.INT8_VECTOR
        field_data.field_name = "int8_vector_field"
        field_info = {"name": "int8_vector_field", "params": {"dim": 128}}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Pack some vectors
        vectors = []
        for i in range(10):
            vector = np.array([i % 128 - 64 for _ in range(128)], dtype=np.int8)
            vectors.append(vector)
            pack_field_value_to_field_data(vector, field_data, field_info, vector_bytes_cache)

        # Before flush, data might be in cache
        # Flush to merge all bytes
        flush_vector_bytes(field_data, vector_bytes_cache)

        # Verify all data is merged
        expected_size = 10 * 128
        assert len(field_data.vectors.int8_vector) == expected_size

        # Verify data correctness
        for idx in range(10):
            expected_bytes = vectors[idx].tobytes()
            actual_bytes = field_data.vectors.int8_vector[idx * 128 : (idx + 1) * 128]
            assert expected_bytes == actual_bytes

    def test_pack_field_value_to_field_data_int8_vector_large_batch(self):
        """Test packing large batch of int8 vectors to verify O(n) performance"""

        field_data = schema_pb2.FieldData()
        field_data.type = DataType.INT8_VECTOR
        field_data.field_name = "int8_vector_field"
        field_info = {"name": "int8_vector_field", "params": {"dim": 768}}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Pack a large number of vectors (similar to the bug scenario)
        num_vectors = 10000
        vectors = []

        start_time = time.time()
        for i in range(num_vectors):
            vector = np.array([(i + j) % 128 - 64 for j in range(768)], dtype=np.int8)
            vectors.append(vector)
            pack_field_value_to_field_data(vector, field_data, field_info, vector_bytes_cache)

        # Flush to merge all collected bytes
        flush_vector_bytes(field_data, vector_bytes_cache)
        elapsed_time = time.time() - start_time

        # Verify performance: should complete in reasonable time (< 10 seconds)
        assert (
            elapsed_time < 10.0
        ), f"Operation took {elapsed_time:.2f} seconds, expected < 10 seconds"

        # Verify data correctness
        expected_total_size = num_vectors * 768
        assert len(field_data.vectors.int8_vector) == expected_total_size

        # Sample verification
        sample_indices = [0, 1000, 5000, 9999]
        for idx in sample_indices:
            expected_bytes = vectors[idx].tobytes()
            actual_bytes = field_data.vectors.int8_vector[idx * 768 : (idx + 1) * 768]
            assert expected_bytes == actual_bytes, f"Vector {idx} data mismatch"

    def test_pack_field_value_to_field_data_int8_vector_invalid_dtype(self):
        """Test error handling for invalid int8 vector dtype"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.INT8_VECTOR
        field_data.field_name = "int8_vector_field"
        field_info = {"name": "int8_vector_field", "params": {"dim": 768}}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Try to pack with wrong dtype
        vector = np.array([1, 2, 3], dtype=np.int32)
        with pytest.raises(ParamError, match="invalid input for int8 vector"):
            pack_field_value_to_field_data(vector, field_data, field_info, vector_bytes_cache)

    def test_pack_field_value_to_field_data_int8_vector_none(self):
        """Test handling None value for int8 vector"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.INT8_VECTOR
        field_data.field_name = "int8_vector_field"
        field_info = {"name": "int8_vector_field", "params": {"dim": 768}}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Pack None value
        pack_field_value_to_field_data(None, field_data, field_info, vector_bytes_cache)

        # Dimension should be set from params
        assert field_data.vectors.dim == 768

    def test_pack_field_value_to_field_data_binary_vector_multiple(self):
        """Test packing multiple binary vectors to verify memory optimization"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.BINARY_VECTOR
        field_data.field_name = "binary_vector_field"
        field_info = {"name": "binary_vector_field", "params": {"dim": 128}}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Pack multiple vectors
        num_vectors = 1000
        vectors = []
        for i in range(num_vectors):
            vector = bytes([(i + j) % 256 for j in range(16)])  # 128 bits = 16 bytes
            vectors.append(vector)
            pack_field_value_to_field_data(vector, field_data, field_info, vector_bytes_cache)

        # Flush to merge all collected bytes
        flush_vector_bytes(field_data, vector_bytes_cache)

        # Verify final result
        assert field_data.vectors.dim == 128
        expected_total_size = num_vectors * 16
        assert len(field_data.vectors.binary_vector) == expected_total_size

        # Verify data correctness for sample vectors
        for idx in [0, 100, 500, 999]:
            expected_bytes = vectors[idx]
            actual_bytes = field_data.vectors.binary_vector[idx * 16 : (idx + 1) * 16]
            assert expected_bytes == actual_bytes, f"Vector {idx} data mismatch"

    def test_pack_field_value_to_field_data_float16_vector_multiple(self):
        """Test packing multiple float16 vectors to verify memory optimization"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.FLOAT16_VECTOR
        field_data.field_name = "float16_vector_field"
        field_info = {"name": "float16_vector_field", "params": {"dim": 128}}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Pack multiple vectors
        num_vectors = 1000
        vectors = []
        for i in range(num_vectors):
            vector = np.array([float(i + j) for j in range(128)], dtype=np.float16)
            vectors.append(vector)
            pack_field_value_to_field_data(vector, field_data, field_info, vector_bytes_cache)

        # Flush to merge all collected bytes
        flush_vector_bytes(field_data, vector_bytes_cache)

        # Verify final result
        assert field_data.vectors.dim == 128
        expected_total_size = num_vectors * 128 * 2  # float16 = 2 bytes per element
        assert len(field_data.vectors.float16_vector) == expected_total_size

        # Verify data correctness for sample vectors
        for idx in [0, 100, 500, 999]:
            expected_bytes = vectors[idx].view(np.uint8).tobytes()
            actual_bytes = field_data.vectors.float16_vector[idx * 128 * 2 : (idx + 1) * 128 * 2]
            assert expected_bytes == actual_bytes, f"Vector {idx} data mismatch"

    def test_pack_field_value_to_field_data_bfloat16_vector_multiple(self):
        """Test packing multiple bfloat16 vectors to verify memory optimization"""
        field_data = schema_pb2.FieldData()
        field_data.type = DataType.BFLOAT16_VECTOR
        field_data.field_name = "bfloat16_vector_field"
        field_info = {"name": "bfloat16_vector_field", "params": {"dim": 128}}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Pack multiple vectors using bytes format (since bfloat16 dtype may not be available)
        num_vectors = 1000
        vectors = []
        for i in range(num_vectors):
            # Create bytes directly for bfloat16 (2 bytes per element)
            vector_bytes = bytes([(i + j) % 256 for j in range(128 * 2)])
            vectors.append(vector_bytes)
            pack_field_value_to_field_data(vector_bytes, field_data, field_info, vector_bytes_cache)

        # Flush to merge all collected bytes
        flush_vector_bytes(field_data, vector_bytes_cache)

        # Verify final result
        assert field_data.vectors.dim == 128
        expected_total_size = num_vectors * 128 * 2  # bfloat16 = 2 bytes per element
        assert len(field_data.vectors.bfloat16_vector) == expected_total_size

        # Verify data correctness for sample vectors
        for idx in [0, 100, 500, 999]:
            expected_bytes = vectors[idx]
            actual_bytes = field_data.vectors.bfloat16_vector[idx * 128 * 2 : (idx + 1) * 128 * 2]
            assert expected_bytes == actual_bytes, f"Vector {idx} data mismatch"

    def test_flush_vector_bytes_all_types(self):
        """Test flush_vector_bytes function for all bytes vector types"""
        vector_types = [
            (DataType.INT8_VECTOR, "int8_vector", 768),
            (DataType.BINARY_VECTOR, "binary_vector", 128),
            (DataType.FLOAT16_VECTOR, "float16_vector", 128),
            (DataType.BFLOAT16_VECTOR, "bfloat16_vector", 128),
        ]

        for vector_type, vector_attr, dim in vector_types:
            field_data = schema_pb2.FieldData()
            field_data.type = vector_type
            field_data.field_name = f"{vector_attr}_field"
            field_info = {"name": f"{vector_attr}_field", "params": {"dim": dim}}
            vector_bytes_cache: Dict[int, List[bytes]] = {}

            # Pack some vectors
            num_vectors = 10
            for i in range(num_vectors):
                if vector_type == DataType.INT8_VECTOR:
                    vector = np.array([i % 128 - 64 for _ in range(dim)], dtype=np.int8)
                elif vector_type == DataType.BINARY_VECTOR:
                    vector = bytes([i % 256 for _ in range(dim // 8)])
                elif vector_type == DataType.FLOAT16_VECTOR:
                    vector = np.array([float(i) for _ in range(dim)], dtype=np.float16)
                else:  # BFLOAT16_VECTOR
                    # Use bytes format for bfloat16 since dtype may not be available
                    vector = bytes([i % 256 for _ in range(dim * 2)])

                pack_field_value_to_field_data(vector, field_data, field_info, vector_bytes_cache)

            # Flush to merge all bytes
            flush_vector_bytes(field_data, vector_bytes_cache)

            # Verify all data is merged
            vector_data = getattr(field_data.vectors, vector_attr)
            if vector_type == DataType.INT8_VECTOR:
                expected_size = num_vectors * dim
            elif vector_type == DataType.BINARY_VECTOR:
                expected_size = num_vectors * (dim // 8)
            else:  # FLOAT16_VECTOR or BFLOAT16_VECTOR
                expected_size = num_vectors * dim * 2

            assert len(vector_data) == expected_size, f"{vector_attr} size mismatch"


class TestPackFieldValueExtendedTypes:
    """Test pack_field_value_to_field_data for extended field types"""

    def test_pack_timestamptz_field(self):
        """Test packing TIMESTAMPTZ field value"""
        field_data = schema_types.FieldData()
        field_data.type = DataType.TIMESTAMPTZ
        field_data.field_name = "ts_field"
        field_info = {"name": "ts_field"}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Test with valid timestamp string
        pack_field_value_to_field_data(
            "2024-01-01T12:00:00Z", field_data, field_info, vector_bytes_cache
        )
        assert len(field_data.scalars.string_data.data) == 1
        assert field_data.scalars.string_data.data[0] == "2024-01-01T12:00:00Z"

    def test_pack_timestamptz_field_none(self):
        """Test packing None value for TIMESTAMPTZ field"""
        field_data = schema_types.FieldData()
        field_data.type = DataType.TIMESTAMPTZ
        field_data.field_name = "ts_field"
        field_info = {"name": "ts_field"}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        pack_field_value_to_field_data(None, field_data, field_info, vector_bytes_cache)
        assert len(field_data.scalars.string_data.data) == 0

    def test_pack_geometry_field(self):
        """Test packing GEOMETRY field value"""
        field_data = schema_types.FieldData()
        field_data.type = DataType.GEOMETRY
        field_data.field_name = "geo_field"
        field_info = {"name": "geo_field", "params": {"max_length": 1000}}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        # Test with WKT geometry string - pack expects string directly
        wkt = "POINT(1.0 2.0)"
        pack_field_value_to_field_data(wkt, field_data, field_info, vector_bytes_cache)
        assert len(field_data.scalars.geometry_wkt_data.data) == 1

    def test_pack_geometry_field_none(self):
        """Test packing None value for GEOMETRY field"""
        field_data = schema_types.FieldData()
        field_data.type = DataType.GEOMETRY
        field_data.field_name = "geo_field"
        field_info = {"name": "geo_field"}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        pack_field_value_to_field_data(None, field_data, field_info, vector_bytes_cache)
        assert len(field_data.scalars.geometry_wkt_data.data) == 0

    def test_pack_unsupported_type_raises_error(self):
        """Test that unsupported field type raises ParamError"""
        field_data = schema_types.FieldData()
        field_data.type = 999  # Invalid type
        field_data.field_name = "bad_field"
        field_info = {"name": "bad_field"}
        vector_bytes_cache: Dict[int, List[bytes]] = {}

        with pytest.raises(ParamError, match="Unsupported data type"):
            pack_field_value_to_field_data("value", field_data, field_info, vector_bytes_cache)


class TestConvertToArrayExtended:
    """Test convert_to_array and convert_to_array_arr for various element types"""

    def test_convert_to_array_int8(self):
        """Test converting INT8 array"""
        field_info = {"name": "arr_field", "element_type": DataType.INT8}
        result = convert_to_array([1, 2, 3], field_info)
        assert list(result.int_data.data) == [1, 2, 3]

    def test_convert_to_array_int16(self):
        """Test converting INT16 array"""
        field_info = {"name": "arr_field", "element_type": DataType.INT16}
        result = convert_to_array([100, 200, 300], field_info)
        assert list(result.int_data.data) == [100, 200, 300]

    def test_convert_to_array_int32(self):
        """Test converting INT32 array"""
        field_info = {"name": "arr_field", "element_type": DataType.INT32}
        result = convert_to_array([1000, 2000], field_info)
        assert list(result.int_data.data) == [1000, 2000]

    def test_convert_to_array_string_type(self):
        """Test converting STRING (alias) array"""
        field_info = {"name": "arr_field", "element_type": DataType.STRING}
        result = convert_to_array(["a", "b"], field_info)
        assert list(result.string_data.data) == ["a", "b"]

    def test_convert_to_array_numpy_input(self):
        """Test converting numpy array"""
        field_info = {"name": "arr_field", "element_type": DataType.INT64}
        arr = np.array([1, 2, 3], dtype=np.int64)
        result = convert_to_array(arr, field_info)
        assert list(result.long_data.data) == [1, 2, 3]


class TestConvertToArrayOfVector:
    """Test convert_to_array_of_vector function"""

    def test_convert_empty_array_of_vector(self):
        """Test converting empty array of vectors"""
        field_info = {
            "name": "vec_arr_field",
            "element_type": DataType.FLOAT_VECTOR,
            "params": {"dim": "4"},
        }
        result = convert_to_array_of_vector([], field_info)
        assert result.dim == 4
        assert len(result.float_vector.data) == 0

    def test_convert_array_of_float_vectors(self):
        """Test converting array of float vectors"""
        field_info = {
            "name": "vec_arr_field",
            "element_type": DataType.FLOAT_VECTOR,
            "params": {"dim": 2},
        }
        vectors = [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]
        result = convert_to_array_of_vector(vectors, field_info)
        assert result.dim == 2
        assert list(result.float_vector.data) == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]

    def test_convert_array_of_float_vectors_numpy(self):
        """Test converting numpy array of vectors"""
        field_info = {
            "name": "vec_arr_field",
            "element_type": DataType.FLOAT_VECTOR,
            "params": {"dim": 2},
        }
        vectors = [np.array([1.0, 2.0], dtype=np.float32), np.array([3.0, 4.0], dtype=np.float32)]
        result = convert_to_array_of_vector(vectors, field_info)
        assert list(result.float_vector.data) == [1.0, 2.0, 3.0, 4.0]

    def test_convert_array_of_vector_unsupported_type(self):
        """Test unsupported element type raises error"""
        field_info = {
            "name": "vec_arr_field",
            "element_type": DataType.BINARY_VECTOR,
            "params": {"dim": 8},
        }
        with pytest.raises(ParamError, match="Unsupported element type"):
            convert_to_array_of_vector([[1, 2]], field_info)

    def test_convert_array_of_float_vectors_invalid_dtype(self):
        """Test numpy array with invalid dtype raises error"""
        field_info = {
            "name": "vec_arr_field",
            "element_type": DataType.FLOAT_VECTOR,
            "params": {"dim": 2},
        }
        # int64 dtype should raise error
        vectors = [np.array([1, 2], dtype=np.int64)]
        with pytest.raises(ParamError, match="invalid input for float32 vector"):
            convert_to_array_of_vector(vectors, field_info)


class TestExtractRowDataV2:
    """Test extract_row_data_from_fields_data_v2 function"""

    def test_extract_bool_with_validity(self):
        """Test extracting bool data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.BOOL
        field_data.field_name = "bool_field"
        field_data.scalars.bool_data.data.extend([True, False, True])
        field_data.valid_data.extend([True, False, True])

        entity_rows = [{}, {}, {}]
        extract_row_data_from_fields_data_v2(field_data, entity_rows)

        assert entity_rows[0]["bool_field"] is True
        assert entity_rows[1]["bool_field"] is None
        assert entity_rows[2]["bool_field"] is True

    def test_extract_int_types_with_validity(self):
        """Test extracting INT8/16/32 data with validity mask"""

        for dtype in [DataType.INT8, DataType.INT16, DataType.INT32]:
            field_data = schema_types.FieldData()
            field_data.type = dtype
            field_data.field_name = "int_field"
            field_data.scalars.int_data.data.extend([10, 20, 30])
            field_data.valid_data.extend([True, False, True])

            entity_rows = [{}, {}, {}]
            extract_row_data_from_fields_data_v2(field_data, entity_rows)

            assert entity_rows[0]["int_field"] == 10
            assert entity_rows[1]["int_field"] is None
            assert entity_rows[2]["int_field"] == 30

    def test_extract_int64_with_validity(self):
        """Test extracting INT64 data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.INT64
        field_data.field_name = "int64_field"
        field_data.scalars.long_data.data.extend([100, 200])
        field_data.valid_data.extend([False, True])

        entity_rows = [{}, {}]
        extract_row_data_from_fields_data_v2(field_data, entity_rows)

        assert entity_rows[0]["int64_field"] is None
        assert entity_rows[1]["int64_field"] == 200

    def test_extract_float_with_validity(self):
        """Test extracting FLOAT data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.FLOAT
        field_data.field_name = "float_field"
        field_data.scalars.float_data.data.extend([1.5, 2.5])
        field_data.valid_data.extend([True, False])

        entity_rows = [{}, {}]
        extract_row_data_from_fields_data_v2(field_data, entity_rows)

        assert entity_rows[0]["float_field"] == pytest.approx(1.5)
        assert entity_rows[1]["float_field"] is None

    def test_extract_double_with_validity(self):
        """Test extracting DOUBLE data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.DOUBLE
        field_data.field_name = "double_field"
        field_data.scalars.double_data.data.extend([1.111, 2.222])
        field_data.valid_data.extend([True, True])

        entity_rows = [{}, {}]
        extract_row_data_from_fields_data_v2(field_data, entity_rows)

        assert entity_rows[0]["double_field"] == pytest.approx(1.111)
        assert entity_rows[1]["double_field"] == pytest.approx(2.222)

    def test_extract_timestamptz_with_validity(self):
        """Test extracting TIMESTAMPTZ data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.TIMESTAMPTZ
        field_data.field_name = "ts_field"
        field_data.scalars.string_data.data.extend(["2024-01-01T00:00:00Z", "2024-12-31T23:59:59Z"])
        field_data.valid_data.extend([True, True])

        entity_rows = [{}, {}]
        extract_row_data_from_fields_data_v2(field_data, entity_rows)

        assert entity_rows[0]["ts_field"] == "2024-01-01T00:00:00Z"
        assert entity_rows[1]["ts_field"] == "2024-12-31T23:59:59Z"

    def test_extract_varchar_with_validity(self):
        """Test extracting VARCHAR data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.VARCHAR
        field_data.field_name = "str_field"
        field_data.scalars.string_data.data.extend(["hello", "world"])

        entity_rows = [{}, {}]
        extract_row_data_from_fields_data_v2(field_data, entity_rows)

        assert entity_rows[0]["str_field"] == "hello"
        assert entity_rows[1]["str_field"] == "world"

    def test_extract_geometry_with_validity(self):
        """Test extracting GEOMETRY data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.GEOMETRY
        field_data.field_name = "geo_field"
        field_data.scalars.geometry_wkt_data.data.extend(["POINT(1 2)", "POINT(3 4)"])
        field_data.valid_data.extend([True, False])

        entity_rows = [{}, {}]
        extract_row_data_from_fields_data_v2(field_data, entity_rows)

        assert entity_rows[0]["geo_field"] == "POINT(1 2)"
        assert entity_rows[1]["geo_field"] is None

    def test_extract_json_returns_true(self):
        """Test extracting JSON data returns True for lazy processing"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.JSON
        field_data.field_name = "json_field"

        entity_rows = [{}]
        result = extract_row_data_from_fields_data_v2(field_data, entity_rows)
        assert result is True

    def test_extract_vector_types_return_true(self):
        """Test vector data types return True for lazy processing"""

        vector_types = [
            DataType.FLOAT_VECTOR,
            DataType.FLOAT16_VECTOR,
            DataType.BFLOAT16_VECTOR,
            DataType.BINARY_VECTOR,
            DataType.SPARSE_FLOAT_VECTOR,
            DataType.INT8_VECTOR,
        ]

        for vtype in vector_types:
            field_data = schema_types.FieldData()
            field_data.type = vtype
            field_data.field_name = "vec_field"

            entity_rows = [{}]
            result = extract_row_data_from_fields_data_v2(field_data, entity_rows)
            assert result is True, f"Expected True for {vtype}"

    def test_extract_array_of_struct_returns_true(self):
        """Test _ARRAY_OF_STRUCT returns True for lazy processing"""

        field_data = schema_types.FieldData()
        field_data.type = DataType._ARRAY_OF_STRUCT
        field_data.field_name = "struct_field"

        entity_rows = [{}]
        result = extract_row_data_from_fields_data_v2(field_data, entity_rows)
        assert result is True

    def test_extract_array_of_vector_returns_true(self):
        """Test _ARRAY_OF_VECTOR returns True for lazy processing"""

        field_data = schema_types.FieldData()
        field_data.type = DataType._ARRAY_OF_VECTOR
        field_data.field_name = "vec_arr_field"

        entity_rows = [{}]
        result = extract_row_data_from_fields_data_v2(field_data, entity_rows)
        assert result is True


class TestExtractRowDataFromFieldsData:
    """Test extract_row_data_from_fields_data function"""

    @pytest.mark.parametrize("empty_data", [None, []])
    def test_empty_fields_data(self, empty_data):
        """Test with empty fields data returns empty dict"""
        result = extract_row_data_from_fields_data(empty_data, 0)
        assert result == {}

    def test_extract_bool_field(self):
        """Test extracting bool field"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.BOOL
        field_data.field_name = "bool_field"
        field_data.scalars.bool_data.data.extend([True, False, True])

        result = extract_row_data_from_fields_data([field_data], 1)
        assert result["bool_field"] is False

    def test_extract_int_field(self):
        """Test extracting int field"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.INT32
        field_data.field_name = "int_field"
        field_data.scalars.int_data.data.extend([10, 20, 30])

        result = extract_row_data_from_fields_data([field_data], 2)
        assert result["int_field"] == 30

    def test_extract_with_validity(self):
        """Test extracting with validity mask"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.INT64
        field_data.field_name = "int64_field"
        field_data.scalars.long_data.data.extend([100, 200, 300])
        field_data.valid_data.extend([True, False, True])

        result = extract_row_data_from_fields_data([field_data], 1)
        assert result["int64_field"] is None

    def test_extract_geometry_field(self):
        """Test extracting geometry field"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.GEOMETRY
        field_data.field_name = "geo_field"
        field_data.scalars.geometry_wkt_data.data.extend(["POINT(0 0)", "POINT(1 1)"])

        result = extract_row_data_from_fields_data([field_data], 0)
        assert result["geo_field"] == "POINT(0 0)"


class TestGetArrayHelpers:
    """Test get_array_length and get_array_value_at_index functions"""

    def test_get_array_length_int64(self):
        """Test get_array_length for INT64 array"""

        array_item = schema_types.ScalarField()
        array_item.long_data.data.extend([1, 2, 3, 4, 5])

        result = get_array_length(array_item)
        assert result == 5

    def test_get_array_length_int_types(self):
        """Test get_array_length for INT8/16/32 array"""

        array_item = schema_types.ScalarField()
        array_item.int_data.data.extend([10, 20, 30])

        result = get_array_length(array_item)
        assert result == 3

    def test_get_array_length_float(self):
        """Test get_array_length for FLOAT array"""

        array_item = schema_types.ScalarField()
        array_item.float_data.data.extend([1.0, 2.0])

        result = get_array_length(array_item)
        assert result == 2

    def test_get_array_length_double(self):
        """Test get_array_length for DOUBLE array"""

        array_item = schema_types.ScalarField()
        array_item.double_data.data.extend([1.1, 2.2, 3.3])

        result = get_array_length(array_item)
        assert result == 3

    def test_get_array_length_string(self):
        """Test get_array_length for string array"""

        array_item = schema_types.ScalarField()
        array_item.string_data.data.extend(["a", "b", "c", "d"])

        result = get_array_length(array_item)
        assert result == 4

    def test_get_array_length_bool(self):
        """Test get_array_length for bool array"""

        array_item = schema_types.ScalarField()
        array_item.bool_data.data.extend([True, False])

        result = get_array_length(array_item)
        assert result == 2

    def test_get_array_value_at_index_int64(self):
        """Test get_array_value_at_index for INT64 array"""

        array_item = schema_types.ScalarField()
        array_item.long_data.data.extend([100, 200, 300])

        assert get_array_value_at_index(array_item, 0) == 100
        assert get_array_value_at_index(array_item, 1) == 200
        assert get_array_value_at_index(array_item, 2) == 300

    def test_get_array_value_at_index_float(self):
        """Test get_array_value_at_index for FLOAT array"""

        array_item = schema_types.ScalarField()
        array_item.float_data.data.extend([1.5, 2.5])

        assert get_array_value_at_index(array_item, 0) == pytest.approx(1.5)
        assert get_array_value_at_index(array_item, 1) == pytest.approx(2.5)

    def test_get_array_value_at_index_string(self):
        """Test get_array_value_at_index for string array"""

        array_item = schema_types.ScalarField()
        array_item.string_data.data.extend(["hello", "world"])

        assert get_array_value_at_index(array_item, 0) == "hello"
        assert get_array_value_at_index(array_item, 1) == "world"


class TestEntityToFieldDataExtended:
    """Test entity_to_field_data for various entity types"""

    def test_entity_to_field_data_timestamptz(self):
        """Test entity_to_field_data for TIMESTAMPTZ type"""

        entity = {
            "type": DataType.TIMESTAMPTZ,
            "name": "ts_field",
            "values": ["2024-01-01T00:00:00Z", "2024-06-15T12:30:00Z"],
        }
        field_info = {"name": "ts_field"}

        result = entity_to_field_data(entity, field_info, 2)
        assert result.field_name == "ts_field"
        assert list(result.scalars.string_data.data) == [
            "2024-01-01T00:00:00Z",
            "2024-06-15T12:30:00Z",
        ]

    def test_entity_to_field_data_geometry(self):
        """Test entity_to_field_data for GEOMETRY type"""

        entity = {
            "type": DataType.GEOMETRY,
            "name": "geo_field",
            "values": ["POINT(0 0)", "LINESTRING(0 0, 1 1)"],
        }
        field_info = {"name": "geo_field", "params": {"max_length": 1000}}

        result = entity_to_field_data(entity, field_info, 2)
        assert result.field_name == "geo_field"
        assert len(result.scalars.geometry_wkt_data.data) == 2


class TestExtractArrayRowDataFunctions:
    """Test extract_array_row_data_with_validity and extract_array_row_data_no_validity"""

    def test_extract_array_row_data_int8_no_validity(self):
        """Test extracting INT8 array data without validity"""

        field_data = schema_types.FieldData()
        field_data.field_name = "arr_field"
        field_data.scalars.array_data.element_type = DataType.INT8

        # Add array data
        arr1 = schema_types.ScalarField()
        arr1.int_data.data.extend([1, 2, 3])
        field_data.scalars.array_data.data.append(arr1)

        arr2 = schema_types.ScalarField()
        arr2.int_data.data.extend([4, 5])
        field_data.scalars.array_data.data.append(arr2)

        entity_rows = [{}, {}]
        extract_array_row_data_no_validity(field_data, entity_rows, 2)

        assert entity_rows[0]["arr_field"] == [1, 2, 3]
        assert entity_rows[1]["arr_field"] == [4, 5]

    def test_extract_array_row_data_int16_no_validity(self):
        """Test extracting INT16 array data without validity"""

        field_data = schema_types.FieldData()
        field_data.field_name = "arr_field"
        field_data.scalars.array_data.element_type = DataType.INT16

        arr1 = schema_types.ScalarField()
        arr1.int_data.data.extend([100, 200])
        field_data.scalars.array_data.data.append(arr1)

        entity_rows = [{}]
        extract_array_row_data_no_validity(field_data, entity_rows, 1)

        assert entity_rows[0]["arr_field"] == [100, 200]

    def test_extract_array_row_data_int32_no_validity(self):
        """Test extracting INT32 array data without validity"""

        field_data = schema_types.FieldData()
        field_data.field_name = "arr_field"
        field_data.scalars.array_data.element_type = DataType.INT32

        arr1 = schema_types.ScalarField()
        arr1.int_data.data.extend([1000, 2000, 3000])
        field_data.scalars.array_data.data.append(arr1)

        entity_rows = [{}]
        extract_array_row_data_no_validity(field_data, entity_rows, 1)

        assert entity_rows[0]["arr_field"] == [1000, 2000, 3000]

    def test_extract_array_row_data_with_validity_int64(self):
        """Test extracting INT64 array data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.field_name = "arr_field"
        field_data.scalars.array_data.element_type = DataType.INT64
        field_data.valid_data.extend([True, False, True])

        arr1 = schema_types.ScalarField()
        arr1.long_data.data.extend([1, 2, 3])
        field_data.scalars.array_data.data.append(arr1)

        arr2 = schema_types.ScalarField()  # Will be None due to valid_data[1] = False
        field_data.scalars.array_data.data.append(arr2)

        arr3 = schema_types.ScalarField()
        arr3.long_data.data.extend([7, 8, 9])
        field_data.scalars.array_data.data.append(arr3)

        entity_rows = [{}, {}, {}]
        extract_array_row_data_with_validity(field_data, entity_rows, 3)

        assert entity_rows[0]["arr_field"] == [1, 2, 3]
        assert entity_rows[1]["arr_field"] is None
        assert entity_rows[2]["arr_field"] == [7, 8, 9]

    def test_extract_array_row_data_with_validity_bool(self):
        """Test extracting BOOL array data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.field_name = "bool_arr"
        field_data.scalars.array_data.element_type = DataType.BOOL
        field_data.valid_data.extend([True, True])

        arr1 = schema_types.ScalarField()
        arr1.bool_data.data.extend([True, False])
        field_data.scalars.array_data.data.append(arr1)

        arr2 = schema_types.ScalarField()
        arr2.bool_data.data.extend([False, True, False])
        field_data.scalars.array_data.data.append(arr2)

        entity_rows = [{}, {}]
        extract_array_row_data_with_validity(field_data, entity_rows, 2)

        assert entity_rows[0]["bool_arr"] == [True, False]
        assert entity_rows[1]["bool_arr"] == [False, True, False]

    def test_extract_array_row_data_with_validity_varchar(self):
        """Test extracting VARCHAR array data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.field_name = "str_arr"
        field_data.scalars.array_data.element_type = DataType.VARCHAR
        field_data.valid_data.extend([True, False])

        arr1 = schema_types.ScalarField()
        arr1.string_data.data.extend(["hello", "world"])
        field_data.scalars.array_data.data.append(arr1)

        arr2 = schema_types.ScalarField()
        field_data.scalars.array_data.data.append(arr2)

        entity_rows = [{}, {}]
        extract_array_row_data_with_validity(field_data, entity_rows, 2)

        assert entity_rows[0]["str_arr"] == ["hello", "world"]
        assert entity_rows[1]["str_arr"] is None

    def test_extract_array_row_data_with_validity_float(self):
        """Test extracting FLOAT array data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.field_name = "float_arr"
        field_data.scalars.array_data.element_type = DataType.FLOAT
        field_data.valid_data.extend([True])

        arr1 = schema_types.ScalarField()
        arr1.float_data.data.extend([1.5, 2.5, 3.5])
        field_data.scalars.array_data.data.append(arr1)

        entity_rows = [{}]
        extract_array_row_data_with_validity(field_data, entity_rows, 1)

        assert entity_rows[0]["float_arr"][0] == pytest.approx(1.5)

    def test_extract_array_row_data_with_validity_double(self):
        """Test extracting DOUBLE array data with validity mask"""

        field_data = schema_types.FieldData()
        field_data.field_name = "double_arr"
        field_data.scalars.array_data.element_type = DataType.DOUBLE
        field_data.valid_data.extend([True])

        arr1 = schema_types.ScalarField()
        arr1.double_data.data.extend([1.111, 2.222])
        field_data.scalars.array_data.data.append(arr1)

        entity_rows = [{}]
        extract_array_row_data_with_validity(field_data, entity_rows, 1)

        assert entity_rows[0]["double_arr"][0] == pytest.approx(1.111)


class TestExtractVectorArrayRowData:
    """Test extract_vector_array_row_data function"""

    def test_extract_float_vector_array(self):
        """Test extracting FLOAT_VECTOR array"""

        field_data = schema_types.FieldData()
        field_data.type = DataType._ARRAY_OF_VECTOR

        # Create vector array data
        vector_data = schema_types.VectorField()
        vector_data.float_vector.data.extend([1.0, 2.0, 3.0, 4.0])
        field_data.vectors.vector_array.element_type = DataType.FLOAT_VECTOR
        field_data.vectors.vector_array.data.append(vector_data)

        result = extract_vector_array_row_data(field_data, 0)
        assert len(result) == 4

    def test_extract_binary_vector_array(self):
        """Test extracting BINARY_VECTOR array"""

        field_data = schema_types.FieldData()
        field_data.type = DataType._ARRAY_OF_VECTOR

        vector_data = schema_types.VectorField()
        vector_data.binary_vector = b"\x01\x02"
        field_data.vectors.vector_array.element_type = DataType.BINARY_VECTOR
        field_data.vectors.vector_array.data.append(vector_data)

        result = extract_vector_array_row_data(field_data, 0)
        assert result == [b"\x01\x02"]


class TestSparseEdgeCases:
    """Test edge cases for sparse vector handling"""

    def test_sparse_proto_to_rows_empty_range(self):
        """Test sparse_proto_to_rows with same start and end"""
        proto = schema_types.SparseFloatArray(dim=100)
        proto.contents.append(struct.pack("I", 0) + struct.pack("f", 1.0))

        rows = entity_helper.sparse_proto_to_rows(proto, 0, 0)
        assert len(rows) == 0

    def test_entity_is_sparse_matrix_exception_handling(self):
        """Test entity_is_sparse_matrix exception branch"""

        # Create an object that raises exception when iterated
        class BadIterator:
            def __iter__(self):
                raise RuntimeError("Cannot iterate")

            def __len__(self):
                return 1

        result = entity_helper.entity_is_sparse_matrix(BadIterator())
        assert result is False


class TestExtractStructArray:
    """Test extract_struct_array_from_column_data function"""

    def test_extract_basic_struct_array(self):
        """Test extracting struct array with scalar fields"""

        # Create struct array field structure
        struct_field = schema_types.FieldData()
        struct_field.type = DataType._ARRAY_OF_STRUCT

        # Add sub-field: name (string)
        name_field = schema_types.FieldData()
        name_field.field_name = "name"
        name_field.type = DataType.ARRAY

        # Helper to create array data for a row
        def create_string_array(values):
            arr = schema_types.ScalarField()
            arr.string_data.data.extend(values)
            return arr

        name_field.scalars.array_data.data.append(create_string_array(["Alice", "Bob"]))
        struct_field.struct_arrays.fields.append(name_field)

        # Add sub-field: age (int)
        age_field = schema_types.FieldData()
        age_field.field_name = "age"
        age_field.type = DataType.ARRAY

        def create_int_array(values):
            arr = schema_types.ScalarField()
            arr.int_data.data.extend(values)
            return arr

        age_field.scalars.array_data.data.append(create_int_array([30, 25]))
        struct_field.struct_arrays.fields.append(age_field)

        # Extract row 0
        result = extract_struct_array_from_column_data(struct_field.struct_arrays, 0)

        assert len(result) == 2
        assert result[0] == {"name": "Alice", "age": 30}
        assert result[1] == {"name": "Bob", "age": 25}

    def test_extract_struct_array_invalid_input(self):
        """Test extracting with invalid input"""

        assert extract_struct_array_from_column_data(None, 0) == []

        empty_field = schema_types.FieldData()
        assert extract_struct_array_from_column_data(empty_field, 0) == []

    def test_extract_struct_array_with_vectors(self):
        """Test extracting struct array containing vectors"""

        struct_field = schema_types.FieldData()
        struct_field.type = DataType._ARRAY_OF_STRUCT

        # Add sub-field: vector (FLOAT_VECTOR)
        vec_field = schema_types.FieldData()
        vec_field.field_name = "vector"
        vec_field.type = DataType._ARRAY_OF_VECTOR

        # Create vector array data for row 0
        # 2 vectors of dim 2: [1.0, 2.0] and [3.0, 4.0]
        vec_data = schema_types.VectorField()
        vec_data.dim = 2
        vec_data.float_vector.data.extend([1.0, 2.0, 3.0, 4.0])

        vec_field.vectors.vector_array.element_type = DataType.FLOAT_VECTOR
        vec_field.vectors.vector_array.data.append(vec_data)
        struct_field.struct_arrays.fields.append(vec_field)

        result = extract_struct_array_from_column_data(struct_field.struct_arrays, 0)

        assert len(result) == 2
        assert result[0]["vector"] == [1.0, 2.0]
        assert result[1]["vector"] == [3.0, 4.0]

    def test_extract_struct_array_all_vector_types(self):
        """Test extracting struct array with all vector types"""

        # FLOAT16
        struct_field = schema_types.FieldData()
        f16_field = schema_types.FieldData()
        f16_field.field_name = "f16"
        f16_field.type = DataType._ARRAY_OF_VECTOR

        vec_data = schema_types.VectorField()
        vec_data.dim = 2
        # 1.0 -> 0x3c00, 2.0 -> 0x4000 (roughly, float16 logic)
        vec_data.float16_vector = np.array([1.0, 2.0], dtype=np.float16).tobytes()

        f16_field.vectors.vector_array.element_type = DataType.FLOAT16_VECTOR
        f16_field.vectors.vector_array.data.append(vec_data)
        struct_field.struct_arrays.fields.append(f16_field)

        result = extract_struct_array_from_column_data(struct_field.struct_arrays, 0)
        assert len(result) == 1
        assert len(result[0]["f16"]) == 2

        # BINARY
        struct_field = schema_types.FieldData()
        bin_field = schema_types.FieldData()
        bin_field.field_name = "bin"
        bin_field.type = DataType._ARRAY_OF_VECTOR

        vec_data = schema_types.VectorField()
        vec_data.dim = 8
        vec_data.binary_vector = b"\x01"

        bin_field.vectors.vector_array.element_type = DataType.BINARY_VECTOR
        bin_field.vectors.vector_array.data.append(vec_data)
        struct_field.struct_arrays.fields.append(bin_field)

        result = extract_struct_array_from_column_data(struct_field.struct_arrays, 0)
        assert len(result) == 1
        assert result[0]["bin"] == [b"\x01"]

        # INT8
        struct_field = schema_types.FieldData()
        int8_field = schema_types.FieldData()
        int8_field.field_name = "int8"
        int8_field.type = DataType._ARRAY_OF_VECTOR

        vec_data = schema_types.VectorField()
        vec_data.dim = 2
        vec_data.int8_vector = b"\x01\x02"

        int8_field.vectors.vector_array.element_type = DataType.INT8_VECTOR
        int8_field.vectors.vector_array.data.append(vec_data)
        struct_field.struct_arrays.fields.append(int8_field)

        result = extract_struct_array_from_column_data(struct_field.struct_arrays, 0)
        assert len(result) == 1
        assert result[0]["int8"] == [1, 2]

    def test_extract_struct_array_unexpected_field_type(self):
        """Test error when struct contains unexpected field type"""

        struct_field = schema_types.FieldData()

        # Invalid sub-field type (INT64 is not ARRAY or ARRAY_OF_VECTOR)
        bad_field = schema_types.FieldData()
        bad_field.field_name = "bad"
        bad_field.type = DataType.INT64

        struct_field.struct_arrays.fields.append(bad_field)

        # Should detect 0 structs because type is not one of checked ones for counting
        result = extract_struct_array_from_column_data(struct_field.struct_arrays, 0)
        assert len(result) == 0

        # Force a count > 0 to test the inner loop error
        # Add a valid field first so num_structs > 0
        valid_field = schema_types.FieldData()
        valid_field.field_name = "valid"
        valid_field.type = DataType.ARRAY
        arr = schema_types.ScalarField()
        arr.int_data.data.extend([1])
        valid_field.scalars.array_data.data.append(arr)
        struct_field.struct_arrays.fields.insert(0, valid_field)

        with pytest.raises(ParamError, match="Unexpected field type"):
            extract_struct_array_from_column_data(struct_field.struct_arrays, 0)


class TestMockedSparseMatrix:
    """Test sparse matrix functions using mocks to avoid environment issues"""

    def test_entity_is_sparse_matrix_mocked(self):
        """Test entity_is_sparse_matrix with mocked scipy sparse matrix"""

        # Create a mock that acts like a scipy sparse matrix
        mock_sparse = MagicMock()
        mock_sparse.shape = (1, 10)

        def is_sparse_side_effect(arg):
            return isinstance(arg, MagicMock)

        # Patch the name in entity_helper namespace to be sure
        with patch(
            "pymilvus.client.entity_helper.SciPyHelper.is_scipy_sparse",
            side_effect=is_sparse_side_effect,
        ):
            assert entity_helper.entity_is_sparse_matrix(mock_sparse) is True

            # Test list of sparse matrices
            mock_sparse2 = MagicMock()
            mock_sparse2.shape = (1, 10)
            assert entity_helper.entity_is_sparse_matrix([mock_sparse, mock_sparse2]) is True

    def test_pack_field_value_to_field_data_sparse_mocked(self):
        """Test packing sparse vector using mocks"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.SPARSE_FLOAT_VECTOR
        field_data.field_name = "sparse"
        field_info = {"name": "sparse"}
        vector_bytes_cache = {}

        # Mock sparse matrix (scipy.sparse.csr_matrix style)
        mock_sparse = MagicMock()
        mock_sparse.shape = (1, 10)
        mock_sparse.tocsr.return_value = mock_sparse
        mock_sparse.indptr = [0, 2]
        mock_sparse.indices = [1, 5]
        mock_sparse.data = [0.5, 0.8]

        def is_sparse_side_effect(arg):
            return isinstance(arg, MagicMock)

        with patch(
            "pymilvus.client.entity_helper.SciPyHelper.is_scipy_sparse",
            side_effect=is_sparse_side_effect,
        ):
            entity_helper.pack_field_value_to_field_data(
                mock_sparse, field_data, field_info, vector_bytes_cache
            )

        assert len(field_data.vectors.sparse_float_vector.contents) == 1
        # content is bytes, we can parse it back to check
        row = entity_helper.sparse_parse_single_row(
            field_data.vectors.sparse_float_vector.contents[0]
        )
        assert len(row) == 2
        assert row[1] == pytest.approx(0.5)
        assert row[5] == pytest.approx(0.8)

    def test_sparse_rows_to_proto_with_list_of_mocks(self):
        """Test sparse_rows_to_proto with list of mocked sparse matrices"""

        mock_sparse1 = MagicMock()
        mock_sparse1.shape = (1, 10)
        mock_sparse1.indices = [1]
        mock_sparse1.data = [0.5]

        mock_sparse2 = MagicMock()
        mock_sparse2.shape = (1, 10)
        mock_sparse2.indices = [2]
        mock_sparse2.data = [0.6]

        def is_sparse_side_effect(arg):
            return isinstance(arg, MagicMock)

        with patch(
            "pymilvus.client.entity_helper.SciPyHelper.is_scipy_sparse",
            side_effect=is_sparse_side_effect,
        ):
            # Directly test the function that handles batch of sparse matrices
            proto = entity_helper.sparse_rows_to_proto([mock_sparse1, mock_sparse2])

        assert len(proto.contents) == 2
        row1 = entity_helper.sparse_parse_single_row(proto.contents[0])
        row2 = entity_helper.sparse_parse_single_row(proto.contents[1])
        assert row1[1] == pytest.approx(0.5)
        assert row2[2] == pytest.approx(0.6)


class TestExtractVectorArrayRowDataExtended:
    """Test extract_vector_array_row_data for remaining types"""

    def test_extract_float16_vector_array(self):
        """Test extracting FLOAT16_VECTOR array"""

        field_data = schema_types.FieldData()
        field_data.type = DataType._ARRAY_OF_VECTOR

        vector_data = schema_types.VectorField()
        # 2 elements: 1.0, 2.0
        vector_data.float16_vector = np.array([1.0, 2.0], dtype=np.float16).tobytes()
        field_data.vectors.vector_array.element_type = DataType.FLOAT16_VECTOR
        field_data.vectors.vector_array.data.append(vector_data)

        result = extract_vector_array_row_data(field_data, 0)
        assert result == pytest.approx([1.0, 2.0])

    def test_extract_bfloat16_vector_array(self):
        """Test extracting BFLOAT16_VECTOR array"""

        field_data = schema_types.FieldData()
        field_data.type = DataType._ARRAY_OF_VECTOR

        vector_data = schema_types.VectorField()
        # Mock bfloat16 bytes
        vector_data.bfloat16_vector = b"\x01\x02\x03\x04"
        field_data.vectors.vector_array.element_type = DataType.BFLOAT16_VECTOR
        field_data.vectors.vector_array.data.append(vector_data)

        result = extract_vector_array_row_data(field_data, 0)
        assert isinstance(result, list)
        assert len(result) > 0

    def test_extract_int8_vector_array(self):
        """Test extracting INT8_VECTOR array"""

        field_data = schema_types.FieldData()
        field_data.type = DataType._ARRAY_OF_VECTOR

        vector_data = schema_types.VectorField()
        vector_data.int8_vector = b"\x01\x02"
        field_data.vectors.vector_array.element_type = DataType.INT8_VECTOR
        field_data.vectors.vector_array.data.append(vector_data)

        result = extract_vector_array_row_data(field_data, 0)
        assert result == [1, 2]

    def test_extract_vector_array_unimplemented(self):
        """Test extracting unsupported vector type"""

        field_data = schema_types.FieldData()
        field_data.type = DataType._ARRAY_OF_VECTOR

        vector_data = schema_types.VectorField()
        field_data.vectors.vector_array.element_type = 999
        field_data.vectors.vector_array.data.append(vector_data)

        with pytest.raises(ParamError, match="Unimplemented type"):
            extract_vector_array_row_data(field_data, 0)


class TestPackFieldValueErrors:
    """Test error handling in pack_field_value_to_field_data"""

    def test_pack_bool_invalid(self):
        field_data = schema_types.FieldData()
        field_data.type = DataType.BOOL
        field_data.field_name = "f"
        field_info = {"name": "f"}

        class NoBool:
            def __bool__(self):
                raise ValueError("no")

        with pytest.raises(DataNotMatchException):
            entity_helper.pack_field_value_to_field_data(NoBool(), field_data, field_info, {})

    def test_pack_int_invalid(self):
        field_data = schema_types.FieldData()
        field_data.type = DataType.INT64
        field_data.field_name = "f"
        field_info = {"name": "f"}

        with pytest.raises(DataNotMatchException):
            entity_helper.pack_field_value_to_field_data("not int", field_data, field_info, {})

    def test_pack_float_invalid(self):
        field_data = schema_types.FieldData()
        field_data.type = DataType.FLOAT
        field_data.field_name = "f"
        field_info = {"name": "f"}

        with pytest.raises(DataNotMatchException):
            entity_helper.pack_field_value_to_field_data("not float", field_data, field_info, {})

    def test_pack_float_vector_invalid(self):
        field_data = schema_types.FieldData()
        field_data.type = DataType.FLOAT_VECTOR
        field_data.field_name = "f"
        field_info = {"name": "f"}

        # This raises TypeError (extend with list of strings) -> caught -> DataNotMatchException
        with pytest.raises(DataNotMatchException):
            entity_helper.pack_field_value_to_field_data(["not float"], field_data, field_info, {})

    def test_pack_binary_vector_invalid(self):
        field_data = schema_types.FieldData()
        field_data.type = DataType.BINARY_VECTOR
        field_data.field_name = "f"
        field_info = {"name": "f"}

        with pytest.raises(DataNotMatchException):
            entity_helper.pack_field_value_to_field_data(123, field_data, field_info, {})

    def test_pack_float16_vector_invalid(self):
        field_data = schema_types.FieldData()
        field_data.type = DataType.FLOAT16_VECTOR
        field_data.field_name = "f"
        field_info = {"name": "f"}

        # raises ParamError explicit check
        with pytest.raises(ParamError):
            entity_helper.pack_field_value_to_field_data("bad", field_data, field_info, {})

    def test_pack_bfloat16_vector_invalid(self):
        field_data = schema_types.FieldData()
        field_data.type = DataType.BFLOAT16_VECTOR
        field_data.field_name = "f"
        field_info = {"name": "f"}

        with pytest.raises(ParamError):
            entity_helper.pack_field_value_to_field_data("bad", field_data, field_info, {})

    def test_pack_int8_vector_invalid(self):
        field_data = schema_types.FieldData()
        field_data.type = DataType.INT8_VECTOR
        field_data.field_name = "f"
        field_info = {"name": "f"}

        with pytest.raises(ParamError):
            entity_helper.pack_field_value_to_field_data("bad", field_data, field_info, {})

    def test_pack_varchar_invalid(self):
        field_data = schema_types.FieldData()
        field_data.type = DataType.VARCHAR
        field_data.field_name = "f"
        field_info = {"name": "f", "params": {Config.MaxVarCharLengthKey: 10}}

        with pytest.raises(DataNotMatchException):
            entity_helper.pack_field_value_to_field_data(123, field_data, field_info, {})

    def test_pack_json_invalid(self):
        field_data = schema_types.FieldData()
        field_data.type = DataType.JSON
        field_data.field_name = "f"
        field_info = {"name": "f"}

        class Unserializable:
            pass

        with pytest.raises(DataNotMatchException):
            entity_helper.pack_field_value_to_field_data(
                Unserializable(), field_data, field_info, {}
            )


class TestPackFieldValueNone:
    """Test packing None values for all types"""

    def test_pack_none_values(self):

        types_to_test = [
            (DataType.BOOL, "bool_data"),
            (DataType.INT8, "int_data"),
            (DataType.INT16, "int_data"),
            (DataType.INT32, "int_data"),
            (DataType.INT64, "long_data"),
            (DataType.FLOAT, "float_data"),
            (DataType.DOUBLE, "double_data"),
            (DataType.VARCHAR, "string_data"),
            (DataType.TIMESTAMPTZ, "string_data"),
            (DataType.JSON, "json_data"),
        ]

        for dtype, field_attr in types_to_test:
            field_data = schema_types.FieldData()
            field_data.type = dtype
            field_data.field_name = "f"
            field_info = {"name": "f"}

            # Pack None
            pack_field_value_to_field_data(None, field_data, field_info, {})

            # Pack another None (coverage for extend/append logic)
            pack_field_value_to_field_data(None, field_data, field_info, {})

            field_container = getattr(field_data.scalars, field_attr)
            assert len(field_container.data) == 0

    def test_pack_none_vectors(self):

        vector_types = [
            DataType.FLOAT_VECTOR,
            DataType.BINARY_VECTOR,
            DataType.FLOAT16_VECTOR,
            DataType.BFLOAT16_VECTOR,
            DataType.SPARSE_FLOAT_VECTOR,
            DataType.INT8_VECTOR,
        ]

        for dtype in vector_types:
            field_data = schema_types.FieldData()
            field_data.type = dtype
            field_data.field_name = "v"
            field_info = {"name": "v", "params": {"dim": 8}}

            pack_field_value_to_field_data(None, field_data, field_info, {})

            # For None, it handles dim
            if dtype != DataType.SPARSE_FLOAT_VECTOR:
                assert field_data.vectors.dim == 8
            else:
                assert field_data.vectors.dim == 0


class TestExtractRowDataV1Extended:
    """Test extract_row_data_from_fields_data (V1) for missing coverage"""

    def test_extract_v1_vectors(self):
        """Test extracting vectors via V1 API"""

        # Float Vector
        field_data = schema_types.FieldData()
        field_data.type = DataType.FLOAT_VECTOR
        field_data.field_name = "fv"
        field_data.vectors.dim = 2
        field_data.vectors.float_vector.data.extend([1.0, 2.0, 3.0, 4.0])  # 2 vectors

        # It extracts one row at a time key-value pairs
        row0 = extract_row_data_from_fields_data([field_data], 0)
        row1 = extract_row_data_from_fields_data([field_data], 1)

        assert row0["fv"] == [1.0, 2.0]
        assert row1["fv"] == [3.0, 4.0]

    def test_extract_v1_binary_vector(self):
        """Test binary vector V1 extraction"""

        field_data = schema_types.FieldData()
        field_data.type = DataType.BINARY_VECTOR
        field_data.field_name = "bv"
        field_data.vectors.dim = 8
        field_data.vectors.binary_vector = b"\x01\x02"  # 2 vectors (1 byte each)

        row0 = extract_row_data_from_fields_data([field_data], 0)
        row1 = extract_row_data_from_fields_data([field_data], 1)

        # Binary vectors in V1 are returned as a list containing the bytes element
        assert row0["bv"] == [b"\x01"]
        assert row1["bv"] == [b"\x02"]

    def test_extract_v1_float16(self):

        field_data = schema_types.FieldData()
        field_data.type = DataType.FLOAT16_VECTOR
        field_data.field_name = "f16"
        field_data.vectors.dim = 1
        # 2 elements
        field_data.vectors.float16_vector = b"\x00\x3c\x00\x3c"  # 1.0, 1.0 (approx)

        row0 = extract_row_data_from_fields_data([field_data], 0)
        row1 = extract_row_data_from_fields_data([field_data], 1)

        assert len(row0["f16"]) == 1
        assert len(row1["f16"]) == 1
        assert isinstance(row0["f16"], list)


class TestJsonEdgeCases:
    """Test JSON conversion edge cases"""

    def test_json_dict_with_int_keys(self):

        data = {1: "value"}
        with pytest.raises(DataNotMatchException, match="JSON key must be str"):
            convert_to_json(data)

    def test_json_arr_with_none(self):

        with pytest.raises(ParamError):
            entity_to_json_arr([None], {"name": "json_field"})


class TestArrayConversion:
    """Test convert_to_array for scalar types"""

    def test_convert_bool_array(self):

        result = convert_to_array([True, False], {"element_type": DataType.BOOL})
        assert list(result.bool_data.data) == [True, False]

    def test_convert_float_array(self):

        result = convert_to_array([1.5, 2.5], {"element_type": DataType.FLOAT})
        assert list(result.float_data.data) == pytest.approx([1.5, 2.5])

    def test_convert_double_array(self):

        result = convert_to_array([1.11, 2.22], {"element_type": DataType.DOUBLE})
        assert list(result.double_data.data) == pytest.approx([1.11, 2.22])

    def test_get_array_value_at_index_extended(self):

        # Test Double
        arr_data = convert_to_array([1.1, 2.2], {"element_type": DataType.DOUBLE})
        assert get_array_value_at_index(arr_data, 0) == pytest.approx(1.1)
        assert get_array_value_at_index(arr_data, 1) == pytest.approx(2.2)
        assert get_array_value_at_index(arr_data, 2) is None

        # Test Bool
        arr_data = convert_to_array([True, False], {"element_type": DataType.BOOL})
        assert get_array_value_at_index(arr_data, 0) is True
        assert get_array_value_at_index(arr_data, 1) is False


class TestExtractRowDataV1Scalar:
    """Test V1 extraction for scalars to cover lines 1047-1093 etc."""

    def test_extract_v1_scalars(self):

        # Create field data for various scalar types
        types = [
            (DataType.BOOL, "bool_data", [True, False], "bool_f"),
            (DataType.INT32, "int_data", [1, 2], "int_f"),
            (DataType.INT64, "long_data", [10, 20], "long_f"),
            (DataType.FLOAT, "float_data", [1.1, 2.2], "float_f"),
            (DataType.DOUBLE, "double_data", [3.3, 4.4], "double_f"),
            (DataType.VARCHAR, "string_data", ["a", "b"], "str_f"),
            (DataType.JSON, "json_data", [b'{"x": 1}', b'{"x": 2}'], "json_f"),
        ]

        fields_data_list = []
        for dtype, attr, val, fname in types:
            fd = schema_types.FieldData()
            fd.type = dtype
            fd.field_name = fname
            getattr(fd.scalars, attr).data.extend(val)
            fields_data_list.append(fd)

        row0 = extract_row_data_from_fields_data(fields_data_list, 0)

        assert row0["bool_f"] is True
        assert row0["int_f"] == 1
        assert row0["long_f"] == 10
        assert row0["float_f"] == pytest.approx(1.1, rel=1e-4)  # float32 precision
        assert row0["double_f"] == pytest.approx(3.3)
        assert row0["str_f"] == "a"
        assert row0["json_f"] == {"x": 1}

    def test_extract_v1_json_dynamic(self):

        fd = schema_types.FieldData()
        fd.type = DataType.JSON
        fd.field_name = "meta"
        fd.is_dynamic = True
        fd.scalars.json_data.data.append(b'{"dy": 100, "other": 200}')

        # Extract, specifying 'dy' as dynamic field
        # dynamic_output_fields argument
        row0 = extract_row_data_from_fields_data([fd], 0, dynamic_output_fields=["dy"])

        assert row0["dy"] == 100
        assert "other" not in row0
        # Wait, logic: row_data.update({k: v for k,v in json_dict.items() if k in dynamic_fields})
        # So only 'dy' is updated into row_data.

        # Test without dynamic_output_fields (None or empty) -> updates all
        row0_all = extract_row_data_from_fields_data([fd], 0)
        assert row0_all["dy"] == 100
        assert row0_all["other"] == 200


class TestPackExceptionsMock:
    """Test exception handling in pack_field_value_to_field_data using mocks"""

    def test_pack_exceptions(self):

        # List of types and attributes to fail
        fail_targets = [
            (DataType.BOOL, "bool_data", True),
            (DataType.INT32, "int_data", 1),
            (DataType.INT64, "long_data", 1),
            (DataType.FLOAT, "float_data", 1.0),
            (DataType.DOUBLE, "double_data", 1.0),
            (DataType.VARCHAR, "string_data", "s"),
            (DataType.TIMESTAMPTZ, "string_data", "2024-01-01"),
            (DataType.JSON, "json_data", {"x": 1}),
            (DataType.ARRAY, "array_data", [1, 2]),
            (DataType.GEOMETRY, "geometry_wkt_data", "POINT(0 0)"),
            (DataType.FLOAT_VECTOR, "float_vector", [1.0] * 8),
            (DataType.BINARY_VECTOR, "binary_vector", b"\x00"),
            (DataType.FLOAT16_VECTOR, "float16_vector", b"\x00\x00"),
            (DataType.BFLOAT16_VECTOR, "bfloat16_vector", b"\x00\x00"),
            (DataType.SPARSE_FLOAT_VECTOR, "sparse_float_vector", {0: 1.0}),
            (DataType.INT8_VECTOR, "int8_vector", b"\x00"),
        ]

        for dtype, attr_name, val in fail_targets:
            field_data_mock = MagicMock()
            field_data_mock.type = dtype
            # Set dim for vectors if accessed
            field_data_mock.vectors.dim = 0
            field_info = {
                "name": "f",
                "element_type": DataType.INT64,
            }  # element_type for ARRAY check

            valid_val = val

            # Configure mock to raise TypeError on access/append which happens LAST
            if "vector" in attr_name:
                if attr_name == "float_vector":
                    getattr(field_data_mock.vectors, attr_name).data.extend.side_effect = TypeError(
                        "Mock Error"
                    )
                elif attr_name == "sparse_float_vector":
                    getattr(field_data_mock.vectors, attr_name).contents.append.side_effect = (
                        TypeError("Mock Error")
                    )
            else:
                getattr(field_data_mock.scalars, attr_name).data.append.side_effect = TypeError(
                    "Mock Error"
                )

            # Skip types where generic exception block is not easily reached/mocked via field_data
            if dtype in (
                DataType.BINARY_VECTOR,
                DataType.FLOAT16_VECTOR,
                DataType.BFLOAT16_VECTOR,
                DataType.INT8_VECTOR,
            ):
                continue

            # For sparse, we need is_scipy_sparse to work
            if dtype == DataType.SPARSE_FLOAT_VECTOR:
                with patch(
                    "pymilvus.client.entity_helper.SciPyHelper.is_scipy_sparse", return_value=False
                ):
                    # Relax match to just expect DataNotMatchException
                    with pytest.raises(DataNotMatchException):
                        pack_field_value_to_field_data(valid_val, field_data_mock, field_info, {})
            else:
                with pytest.raises(DataNotMatchException):
                    pack_field_value_to_field_data(valid_val, field_data_mock, field_info, {})

    def test_pack_vector_exceptions_value_trigger(self):
        """Trigger exceptions in vector packing using object that raises error"""

        class ErrorObj:
            def __len__(self):
                raise TypeError("Len Error")

            def __bytes__(self):
                raise TypeError("Bytes Error")

            def __iter__(self):
                raise TypeError("Iter Error")

        field_info = {"name": "f"}

        fd = schema_types.FieldData()
        fd.type = DataType.BINARY_VECTOR
        with pytest.raises(DataNotMatchException):
            pack_field_value_to_field_data(ErrorObj(), fd, field_info, {})


class TestExtractRowDataV1Validity:
    """Test extract_row_data_from_fields_data (V1) validity and error handling"""

    def test_json_validity_v1(self):

        fd = schema_types.FieldData()
        fd.type = DataType.JSON
        fd.field_name = "j"
        fd.scalars.json_data.data.append(b"{}")
        fd.valid_data.append(False)  # Invalid

        row0 = extract_row_data_from_fields_data([fd], 0)
        assert row0["j"] is None

    def test_json_invalid_bytes_v1(self):

        fd = schema_types.FieldData()
        fd.type = DataType.JSON
        fd.field_name = "j"
        fd.scalars.json_data.data.append(b"{invalid")

        with pytest.raises(orjson.JSONDecodeError):
            extract_row_data_from_fields_data([fd], 0)

    def test_array_validity_v1(self):

        fd = schema_types.FieldData()
        fd.type = DataType.ARRAY
        fd.field_name = "arr"
        # We need to populate array_data.data to satisfy length check
        fd.scalars.array_data.data.append(schema_types.ScalarField())
        fd.valid_data.append(False)

        row0 = extract_row_data_from_fields_data([fd], 0)
        assert row0["arr"] is None

    def test_vector_validity_v1(self):

        # Test Float Vector validity
        fd = schema_types.FieldData()
        fd.type = DataType.FLOAT_VECTOR
        fd.field_name = "fv"
        fd.valid_data.append(False)

        row0 = extract_row_data_from_fields_data([fd], 0)
        assert row0["fv"] is None

        # Test Binary Vector validity
        fd = schema_types.FieldData()
        fd.type = DataType.BINARY_VECTOR
        fd.field_name = "bv"
        fd.valid_data.append(False)
        row0 = extract_row_data_from_fields_data([fd], 0)
        assert row0["bv"] is None

        # Test Float16 Vector validity
        fd = schema_types.FieldData()
        fd.type = DataType.FLOAT16_VECTOR
        fd.field_name = "f16"
        fd.valid_data.append(False)
        row0 = extract_row_data_from_fields_data([fd], 0)
        assert row0["f16"] is None


class TestStructArrayTruncated:
    """Test extract_struct_array_from_column_data with truncated/invalid vector data"""

    def test_struct_array_truncated_vectors_all_types(self):

        # Types to test: FLOAT, FLOAT16, BFLOAT16, INT8, BINARY

        def set_float(vec_data):
            vec_data.float_vector.data.extend([1.0, 2.0])

        def set_float16(vec_data):
            vec_data.float16_vector = b"\x00\x00\x00\x00"

        def set_bfloat16(vec_data):
            vec_data.bfloat16_vector = b"\x00\x00\x00\x00"

        def set_int8(vec_data):
            vec_data.int8_vector = b"\x00\x00"

        def set_binary(vec_data):
            vec_data.binary_vector = b"\x00"

        types_setup = [
            (DataType.FLOAT_VECTOR, "float_vector", set_float, 2),
            (DataType.FLOAT16_VECTOR, "float16_vector", set_float16, 2),  # 4 bytes = 1 vec of dim 2
            (DataType.BFLOAT16_VECTOR, "bfloat16_vector", set_bfloat16, 2),
            (DataType.INT8_VECTOR, "int8_vector", set_int8, 2),  # 2 bytes = 1 vec dim 2
            (DataType.BINARY_VECTOR, "binary_vector", set_binary, 8),  # 1 byte = 1 vec dim 8
        ]

        for dtype, _attr_name, fill_func, dim in types_setup:
            sa_field = schema_types.FieldData()
            sa_field.type = DataType._ARRAY_OF_STRUCT

            vec_field = schema_types.FieldData()
            vec_field.type = DataType._ARRAY_OF_VECTOR
            vec_field.field_name = "vec"
            vec_field.vectors.vector_array.element_type = dtype
            vec_f = schema_types.VectorField()
            vec_field.vectors.vector_array.data.append(vec_f)

            vector_data = vec_field.vectors.vector_array.data[0]
            vector_data.dim = dim

            # Fill with enough data for ONE struct
            fill_func(vector_data)  # e.g. add [1.0, 2.0]

            # Array field to force num_structs = 2
            arr_field = schema_types.FieldData()
            arr_field.type = DataType.ARRAY
            arr_field.field_name = "ids"
            arr_scalar = schema_types.ScalarField()
            arr_field.scalars.array_data.data.append(arr_scalar)
            arr_field.scalars.array_data.data[0].int_data.data.extend([1, 2])

            # Add fields carefully to sa_field.struct_arrays.fields
            # Since fields is repeated, we must add() to it
            sa_field.struct_arrays.fields.add().CopyFrom(arr_field)
            sa_field.struct_arrays.fields.add().CopyFrom(vec_field)

            result = extract_struct_array_from_column_data(sa_field.struct_arrays, 0)
            assert len(result) == 2
            # Second struct should have None for vector because data ran out
            assert result[1]["vec"] is None


class TestExtractRowDataV1Int8:
    """Test extraction of INT8 vectors in V1"""

    def test_extract_v1_int8(self):

        fd = schema_types.FieldData()
        fd.type = DataType.INT8_VECTOR
        fd.field_name = "i8"
        fd.vectors.dim = 2

        # 4 bytes = 2 vectors of dim 2
        fd.vectors.int8_vector = b"\x01\x02\x03\x04"

        row0 = extract_row_data_from_fields_data([fd], 0)

        assert row0["i8"] == [b"\x01\x02"]


class TestJsonStrInput:
    """Test string input for convert_to_json"""

    def test_json_str_valid(self):
        res = convert_to_json('{"key": 1}')
        assert res == b'{"key": 1}'


class TestEntityHelperCoverage:
    """Test class for remaining coverage gaps in entity_helper.py"""

    def test_entity_is_sparse_matrix_inner_functions(self):
        """Cover inner functions of entity_is_sparse_matrix"""
        # Test case where row is valid list/dict but elements fail type check
        # Non-string fallback in is_type_in_str
        data = [[(1, object())]]
        assert entity_helper.entity_is_sparse_matrix(data) is False

        # inner is_type_in_str: passed string that is not parseable
        data = [[(1, "not_a_number")]]
        assert entity_helper.entity_is_sparse_matrix(data) is False

    def test_sparse_rows_to_proto_errors(self):
        """Cover sparse_rows_to_proto error paths"""
        # Length mismatch - hard to trigger via public API due to zip usage/dict conversion
        # Skipping length mismatch test

        with pytest.raises(ParamError, match=r"positive and less than 2\^32-1"):
            # Index -1
            entity_helper.sparse_rows_to_proto([{-1: 0.5}])

        # NaN value
        with pytest.raises(ParamError, match="NaN"):
            entity_helper.sparse_rows_to_proto([{1: float("nan")}])

    def test_sparse_proto_to_rows_defaults(self):
        """Cover sparse_proto_to_rows default arguments"""
        proto = schema_types.SparseFloatArray()
        # Add dummy content
        proto.contents.append(b"dummy")

        with mock.patch("pymilvus.client.entity_helper.sparse_parse_single_row") as mock_parse:
            entity_helper.sparse_proto_to_rows(proto)  # defaults start=None, end=None
            mock_parse.assert_called()

        with pytest.raises(ParamError, match="Vector must be a sparse float vector"):
            entity_helper.sparse_proto_to_rows("not_a_proto")

    def test_get_input_num_rows_mock(self):
        """Cover get_input_num_rows with scipy sparse"""
        with mock.patch(
            "pymilvus.client.entity_helper.SciPyHelper.is_scipy_sparse", return_value=True
        ):
            mock_entity = mock.MagicMock()
            mock_entity.shape = (5, 10)
            assert entity_helper.get_input_num_rows(mock_entity) == 5

    def test_entity_type_to_dtype(self):
        """Cover entity_type_to_dtype"""
        assert entity_helper.entity_type_to_dtype(100) == 100  # int pass through
        assert entity_helper.entity_type_to_dtype("Int64") == DataType.INT64
        with pytest.raises(ParamError):
            entity_helper.entity_type_to_dtype(None)

    def test_convert_to_str_array_errors(self):
        """Cover convert_to_str_array errors"""
        field_info = {"name": "test", "params": {"max_length": 5}}

        # Not a string
        with pytest.raises(ParamError, match="expects string input"):
            entity_helper.convert_to_str_array([123], field_info)

        # Too long
        with pytest.raises(ParamError, match="exceeds max length"):
            entity_helper.convert_to_str_array(["123456"], field_info)

    def test_convert_to_json_errors(self):
        """Cover convert_to_json errors"""
        # non-string key in dict
        with pytest.raises(DataNotMatchException, match="JSON key must be str"):
            entity_helper.convert_to_json({1: "v"})

        # non-string key in parsed json string
        # Mock orjson.loads to return a dict with non-string keys from a string input
        with mock.patch("pymilvus.client.entity_helper.orjson.loads", return_value={1: "v"}):
            with pytest.raises(DataNotMatchException, match="Invalid JSON string"):
                entity_helper.convert_to_json('{"1": "v"}')

    def test_convert_to_json_arr_none(self):
        """Cover convert_to_json_arr None check"""
        field_info = {"name": "test"}
        with pytest.raises(ParamError, match="expects a non-None input"):
            entity_helper.convert_to_json_arr([None], field_info)

    def test_convert_to_array_unsupported(self):
        """Cover convert_to_array unsupported type"""
        field_info = {"name": "test", "element_type": DataType.UNKNOWN}
        with pytest.raises(ParamError, match="Unsupported element type"):
            entity_helper.convert_to_array([], field_info)

    def test_convert_to_array_of_vector_errors(self):
        """Cover convert_to_array_of_vector errors"""
        field_info = {"name": "test", "element_type": DataType.FLOAT_VECTOR, "params": {"dim": 4}}

        # Invalid dtype
        arr = [np.array([1, 2], dtype="int")]
        with pytest.raises(ParamError, match="dtype=float32"):
            entity_helper.convert_to_array_of_vector(arr, field_info)

        # Unsupported element type
        field_info["element_type"] = DataType.INT64
        with pytest.raises(ParamError, match="Unsupported element type"):
            entity_helper.convert_to_array_of_vector([], field_info)

    def test_flush_vector_bytes(self):
        """Cover flush_vector_bytes"""
        field_data = schema_types.FieldData()
        field_data.type = DataType.INT8_VECTOR
        cache = {id(field_data): [b"\x01"]}

        entity_helper.flush_vector_bytes(field_data, cache)
        assert field_data.vectors.int8_vector == b"\x01"
        assert id(field_data) not in cache

        # Empty cache
        entity_helper.flush_vector_bytes(field_data, {})
        # Should do nothing
