CoolFace
Datasetpublic

giahy2507/AptMQL-Bench

AptMQL-Bench 📄 Paper: AptMQL-Bench: From Text-to-SQL to Text-to-MQL via Access-Pattern Schema Design and Data-Preserving Migration  ·  arXiv: coming soon AptMQL-Bench is a benchmark for text-to-MQL — the task of translating human-readable requests into executable MongoDB Query Language (MQL) aggregation pipelines. It contains 21 document-oriented databases, 3,181 natural-language requests, and their associated gold MQL queries. Most existing text-to-MQL resources are… See the full description on the dataset page: https://huggingface.co/datasets/giahy2507/AptMQL-Bench.

sourceHugging Facecc-by-sa-4.0updated 10d agoView on Hugging Face
0likes44downloads
eval_script.py650 linesDownload Raw Back to root
1# turn off warning
2import warnings
3warnings.filterwarnings("ignore")
4
5import json
6from typing import Dict, List, TypeAlias, Any, cast, TypedDict
7import traceback
8
9import pandas as pd
10from pandas.api.types import is_datetime64_any_dtype, is_numeric_dtype, is_object_dtype
11from pandas.testing import assert_frame_equal, assert_series_equal
12
13from pymongo import MongoClient
14from bson import json_util as ejson
15import bson.decimal128
16
17
18class FuzzyMatchingResults(TypedDict):
19    is_match: bool
20    message: str
21
22JSONDict: TypeAlias = Dict[str, "AnyJSON"]
23JSONList: TypeAlias = List["AnyJSON"]
24JSONPrimitive: TypeAlias = str | int | float | bool | None
25AnyJSON: TypeAlias = JSONDict | JSONList | JSONPrimitive
26
27
28class EjsonParseException(Exception):
29    pass
30
31
32def parse_to_ejson(string_result: str) -> list[JSONList | JSONDict]:
33    try:
34        parsed_result: AnyJSON = ejson.loads(string_result)
35    except json.decoder.JSONDecodeError:
36        raise EjsonParseException(string_result)
37    if isinstance(parsed_result, list):
38        return [transform_to_dict(entry) for entry in parsed_result]
39    return [transform_to_dict(parsed_result)]
40
41
42def transform_to_dict(input_obj: AnyJSON) -> JSONList | JSONDict:
43    if isinstance(input_obj, dict):
44        return {key: handle_special_types(val) for key, val in input_obj.items()}
45    if isinstance(input_obj, list):
46        return input_obj
47    # Must be JSONPrimitive
48    return dict(no_name_col=input_obj)
49
50
51def handle_special_types(value: Any) -> Any:
52    if isinstance(value, bson.decimal128.Decimal128):
53        return value.to_decimal()
54    return value
55
56
57def flatten_json_list(
58    value: JSONList, parent_key: str, separator: str = "_"
59) -> Dict[str, JSONPrimitive]:
60    flattened: Dict[str, JSONPrimitive] = {}
61
62    # Iterate over the elements in the list
63    for i, item in enumerate(value):
64        # Create a new key for each element in the list
65        new_key_i = f"{parent_key}{separator}{i}"
66
67        # Check if the element is a dictionary or list for further recursion
68        if isinstance(item, dict):
69            # Recursively flatten nested dictionaries or lists within the list
70            flattened.update(flatten_json(item, new_key_i, separator=separator))
71        elif isinstance(item, list):
72            flattened.update(flatten_json_list(item, new_key_i, separator=separator))
73        else:
74            # Base case: update the flattened dictionary with non-dict and non-list values
75            flattened[new_key_i] = item
76
77    # Return the flattened list as a dictionary
78    return flattened
79
80
81def flatten_json(
82    json_data: JSONDict | JSONList, parent_key: str = "", separator: str = "_"
83) -> Dict[str, JSONPrimitive]:
84    """
85    Recursively flattens a nested JSON structure.
86    Parameters:
87    :param: json_data: The input JSON data to be flattened.
88    :param: parent_key: The key prefix for the current level of recursion.
89    :param: separator: The character used to separate nested keys in the flattened structure.
90    :return: A flattened dictionary representation of the input JSON data.
91
92    Example:
93    >>> json_data = {"user": {"name": {"first": "John", "last": "Doe"}, "age": 25}}
94    >>> flattened = flatten_json(json_data)
95    >>> print(flattened)
96    {'user_name_first': 'John', 'user_name_last': 'Doe', 'user_age': 25}
97    """
98    if isinstance(json_data, list):
99        return flatten_json_list(json_data, parent_key=parent_key, separator=separator)
100
101    # Initialize an empty dictionary to store flattened key-value pairs
102    flattened: Dict[str, JSONPrimitive] = {}
103
104    # Iterate over key-value pairs in the input JSON data
105    for key, value in json_data.items():
106        # Create a new key by concatenating the parent_key, separator, and current key
107        new_key = f"{parent_key}{separator}{key}" if parent_key else key
108
109        # Check if the value is a dictionary for further recursion
110        if isinstance(value, dict):
111            # Recursively flatten nested dictionaries and update the flattened dictionary
112            flattened.update(flatten_json(value, new_key, separator=separator))
113        # Check if the value is a list for handling nested dictionaries or lists within the list
114        elif isinstance(value, list):
115            flattened.update(flatten_json_list(value, new_key, separator=separator))
116        else:
117            # Base case: update the flattened dictionary with non-dict and non-list values
118            flattened[new_key] = value
119
120    # Return the flattened dictionary
121    return flattened
122
123
124def json_to_dataframe(json_data: list[JSONDict | JSONList]) -> pd.DataFrame:
125    """
126    Converts a list of JSON objects into a pandas DataFrame.
127    :param: json_data: A list of dictionaries representing JSON objects.
128    :return: A DataFrame where each JSON object is flattened into columns,
129                     allowing for easier analysis and manipulation.
130
131    Example:
132    >>> json_data = [
133    ...     {"id": 1, "name": {"first": "John", "last": "Doe"}, "age": 25},
134    ...     {"id": 2, "name": {"first": "Jane", "last": "Smith"}, "age": 30}
135    ... ]
136    >>> df = json_to_dataframe(json_data)
137    >>> print(df)
138       id  name_first name_last  age
139    0   1        John       Doe   25
140    1   2        Jane     Smith   30
141    """
142    # Flatten each JSON object in the list
143    flattened_data = [flatten_json(entry) for entry in json_data]
144    # Create a DataFrame from the flattened data
145    df = pd.DataFrame(flattened_data)
146    return df
147
148
149def has_list(value: Any) -> bool:
150    """
151    Check if the input value is a list.
152    :param: value: the input value to be checked.
153    :return: True if the input value is a list, False otherwise.
154    """
155    return isinstance(value, list)
156
157
158def check_for_nested_arrays(mql_str: list):
159    """
160    Check if the input object contains an embedded array
161    :param mql_str:
162    :return:
163    """
164    # Create a temporary dataframe from mql_str
165    tmp_df = pd.json_normalize(mql_str)
166
167    # Loop through all columns and identify columns with and without lists
168    columns_with_lists = [
169        col for col in tmp_df.columns if any(tmp_df[col].apply(has_list))
170    ]
171
172    # Check to see if there are columns with nested lists. If so, return True
173    return True if columns_with_lists else False
174
175
176def json_normalize_with_nested_lists(
177    mql_data: list[JSONDict | JSONList],
178    max_iterations: int = 10,
179    max_rows: int = 10000
180) -> pd.DataFrame:
181    """
182    Creates a tidy table using json_normalize and pd.explode. Used to unpack nested arrays inside the dictionary.
183    :param: mql_data: the input dictionary to be normalized.
184    :param: max_iterations: Maximum number of explosion iterations to prevent infinite loops (default: 10)
185    :param: max_rows: Maximum rows to prevent memory exhaustion (default: 10,000)
186    :return: Tidy DataFrame obtained by normalizing the input dictionary.
187    """
188    import gc
189    
190    mdb_df = pd.json_normalize(cast(list[dict], mql_data), sep="_")
191    
192    # Pre-identify columns with lists more efficiently
193    columns_with_lists = []
194    for col in mdb_df.columns:
195        # Check first non-null value only for performance
196        sample = mdb_df[col].dropna()
197        if len(sample) > 0 and isinstance(sample.iloc[0], list):
198            columns_with_lists.append(col)
199    
200    iteration = 0
201    exceeded_limit = False
202    
203    while len(columns_with_lists) > 0 and iteration < max_iterations:
204        iteration += 1
205        
206        col = columns_with_lists[0]
207        
208        # Estimate row growth before exploding
209        max_list_len = mdb_df[col].dropna().apply(lambda x: len(x) if isinstance(x, list) else 1).max()
210        estimated_rows = len(mdb_df) * max_list_len
211        
212        if estimated_rows > max_rows:
213            print(f"Warning: Exploding '{col}' would create ~{estimated_rows} rows (max: {max_rows}). Stopping to prevent memory exhaustion.")
214            exceeded_limit = True
215            break
216        
217        # Explode one column at a time to control memory
218        mdb_df = mdb_df.explode(column=col, ignore_index=True)
219        
220        # Re-check for nested lists
221        columns_with_lists = []
222        for col in mdb_df.columns:
223            sample = mdb_df[col].dropna()
224            if len(sample) > 0 and isinstance(sample.iloc[0], list):
225                columns_with_lists.append(col)
226        
227        # Force garbage collection every iteration
228        if iteration % 2 == 0:
229            gc.collect()
230    
231    # Skip expensive re-normalization if we hit limits
232    if exceeded_limit:
233        print(f"Returning partially exploded DataFrame with {len(mdb_df)} rows.")
234        return mdb_df
235    
236    # Only re-normalize if we actually have dict/list values and DataFrame is reasonably sized
237    needs_renormalize = False
238    if len(mdb_df) < 5000:  # Only check if reasonably sized
239        for col in mdb_df.columns:
240            sample = mdb_df[col].dropna()
241            if len(sample) > 0 and isinstance(sample.iloc[0], (dict, list)):
242                needs_renormalize = True
243                break
244    
245    if needs_renormalize:
246        res = pd.json_normalize(mdb_df.to_dict("records"), sep="_")
247        del mdb_df
248        gc.collect()
249        return res
250    
251    return mdb_df
252
253
254def normalize_table(df: pd.DataFrame, sort_values: bool) -> pd.DataFrame:
255    """
256    Normalize a DataFrame by sorting columns in alphabetical order and resetting the index.
257    :param df: The input DataFrame to be normalized.
258    :param sort_values (bool): Whether or not to sort the values in the dataframe.
259    :return: The normalized DataFrame.
260    """
261    df = normalize_cols(df)
262
263    # Sort columns in alphabetical order
264    sorted_df = df.reindex(sorted(df.columns), axis=1)
265
266    if sort_values:
267        # Sort the values in the columns to ensure consistent comparison
268        sorted_df = sorted_df.sort_values(list(sorted_df.columns))
269
270    # Reset index
271    sorted_df = sorted_df.reset_index(drop=True)
272
273    return sorted_df
274
275
276def is_date_series(series: pd.Series) -> bool:
277    """
278    Check if a pandas series contains date like values
279    """
280    if not is_object_dtype(series):
281        return False
282
283    try:
284        pd.to_datetime(series)
285        return True
286    except (ValueError, TypeError):
287        return False
288
289
290def is_numeric_series(series: pd.Series) -> bool:
291    """
292    Check if a pandas series contains numeric values
293    """
294    if is_numeric_dtype(series):
295        return True
296    if not is_object_dtype(series):
297        return False
298
299    try:
300        pd.to_numeric(series.astype(str))
301        return True
302    except (ValueError, TypeError):
303        return False
304
305
306def normalize_cols(df: pd.DataFrame) -> pd.DataFrame:
307    """
308    Normalises date and numeric fields in a dataframe so that type mismatch does not impact the comparison result
309    """
310    for col in df.columns:
311        if is_date_series(df[col]) or is_datetime64_any_dtype(df[col]):
312            try:
313                df[col] = pd.to_datetime(df[col])
314            except ValueError:
315                pass
316        if is_numeric_series(df[col]):
317            try:
318                df[col] = pd.to_numeric(df[col].astype(str))
319            except ValueError:
320                pass
321    return df
322
323
324def compare_df(
325    df_A: pd.DataFrame,
326    df_B: pd.DataFrame,
327    sort_values: bool = True,
328    rtol: float = 1e-5,
329    atol: float = 1e-8,
330    allow_column_subset: bool = False,
331    verbose: bool = True,
332) -> bool:
333    """
334    Compare two DataFrames by performing fuzzy column matching and value comparison.
335    
336    This function compares DataFrames by matching columns based on their values rather
337    than column names, making it robust to column naming differences. It normalizes
338    data types (dates and numeric fields) before comparison and uses tolerances for
339    numeric comparisons.
340    
341    The comparison process:
342    1. Validates that df_A is not empty and column counts match
343    2. Normalizes date and numeric columns in both DataFrames
344    3. Matches columns from df_A to df_B by comparing sorted values
345    4. Normalizes both DataFrames (sort columns alphabetically, optionally sort values)
346    5. Compares the normalized DataFrames with specified tolerances
347    
348    Args:
349        df_A (pd.DataFrame): The first DataFrame to compare (reference DataFrame)
350        df_B (pd.DataFrame): The second DataFrame to compare (target DataFrame)
351        sort_values (bool, optional): Whether to sort DataFrame values before comparison.
352            Defaults to True.
353        rtol (float, optional): Relative tolerance for numeric comparisons. Defaults to 1e-5.
354        atol (float, optional): Absolute tolerance for numeric comparisons. Defaults to 1e-8.
355        verbose (bool, optional): Whether to print comparison messages. Defaults to True.
356        allow_column_subset (bool, optional): Whether to allow df_B to have additional columns not in df_A. Defaults to False.
357    Returns:
358        tuple[bool, str]: A tuple containing:
359            - bool: True if DataFrames match, False otherwise
360            - str: Description of the comparison result:
361                - Success: "Found all matched columns, and all the values are also matched"
362                - Failure: Describes empty DataFrame, column count mismatch, unmatched
363                  columns, or data value mismatches
364    
365    Examples:
366        >>> df1 = pd.DataFrame({"id": [1, 2], "name": ["John", "Jane"]})
367        >>> df2 = pd.DataFrame({"user_id": [1, 2], "user_name": ["John", "Jane"]})
368        >>> is_match, message = compare_df(df1, df2)
369        >>> print(is_match)
370        True
371        >>> print(message)
372        Found all matched columns, and all the values are also matched
373    
374    Note:
375        - Column matching is based on sorted values, not column names
376        - Date columns are normalized to datetime format
377        - Numeric columns (including string representations) are converted to numeric
378        - Uses pandas assert_frame_equal with check_dtype=False and check_exact=False
379        - If all columns can't be matched, returns False with details of matched columns
380        - If values don't match, returns the first line of the pandas assertion error
381    """
382    if df_A.empty:
383        message = "df_A is empty; considered a subset by definition."
384        if verbose:
385            print(message)
386        return False, message
387
388    if len(df_A) != len(df_B):
389        message = f"Row count mismatch: df_A has {len(df_A)} rows, df_B has {len(df_B)} rows."
390        if verbose:
391            print(message)
392        return False, message
393
394    if len(df_A.columns) != len(df_B.columns) and not allow_column_subset:
395        message = f"Column count mismatch: df_A has {len(df_A.columns)} columns, df_B has {len(df_B.columns)} columns."
396        if verbose:
397            print(message)
398        return False, message
399
400    # Normalize date columns
401    df_A = normalize_cols(df_A.copy())
402    df_B = normalize_cols(df_B.copy())
403
404    # Find matched columns by value comparison
405    col_mappings = dict()
406    for A_column in df_A.columns:
407        for B_column in df_B.columns:
408            if B_column in col_mappings.values():
409                continue
410            try:
411                assert_series_equal(
412                    df_A[A_column].sort_values().reset_index(drop=True),
413                    df_B[B_column].sort_values().reset_index(drop=True),
414                    check_dtype=False,
415                    check_names=False,
416                    check_exact=False,
417                    rtol=rtol,
418                    atol=atol,
419                )
420                col_mappings[A_column] = B_column
421                break
422            except AssertionError:
423                continue
424    
425    if len(col_mappings) == 0:
426        message = f"Could not find any matched columns between df_A and df_B."
427        if verbose:
428            print(message)
429        return False, message
430
431    if len(col_mappings) != len(df_A.columns) and not allow_column_subset:
432        message = f"Only found {len(col_mappings)} matched columns out of {len(df_A.columns)} columns. Found columns: {col_mappings}"
433        if verbose:
434            print(message)
435        return False, message
436
437    # Normalize df_A
438    df_A_norm = normalize_table(df_A, sort_values=sort_values)
439
440    # Get matched columns from df_B then normalize
441    df_mB = df_B[list(col_mappings.values())].rename(
442        columns={v: k for k, v in col_mappings.items()}
443    )
444    df_mB_norm = normalize_table(df_mB, sort_values=sort_values)
445    
446    # Compare normalized dataframes
447    try:
448        assert_frame_equal(df_A_norm, df_mB_norm, check_dtype=False, check_exact=False, rtol=rtol, atol=atol)
449        return True, f"Found all matched columns, and all the values are also matched"
450    except AssertionError:
451        error = traceback.format_exc()
452        # search this pattern "Series values are different (100.0 %)" and get the number
453        pd_error = error.rsplit("\n\n", 1)[-1]
454        fline_pd_error = pd_error.split("\n")[0]
455        return False, f"Found all matched columns, but data mismatch with error: {fline_pd_error}"
456
457def is_valid(df: pd.DataFrame) -> bool:
458    """
459    Validates the MQL dataframe making sure the values of all columns are consistent in terms of data type
460    """
461    for col in df.columns:
462        try:
463            df[col].sort_values().reset_index(drop=True)
464        except TypeError as ex:
465            print("MQL results are invalid:", ex)
466            return False
467    return True
468
469def compare_fuzzy(
470    mdb_A_output: str,
471    mdb_B_output: str,
472    sort_values: bool = False,
473    allow_column_subset: bool = True,
474    explode_nested_arrays: bool = True,
475    rtol: float = 0.01,
476    atol: float = 0.01,
477    verbose: bool = False,
478) -> FuzzyMatchingResults:
479    """
480    Compare two MongoDB query results using fuzzy matching techniques.
481    
482    This function performs a comprehensive comparison between two MongoDB (MQL) 
483    query results by parsing Extended JSON (EJSON) strings, converting them to 
484    pandas DataFrames, normalizing data structures, and comparing values with 
485    tolerance for numeric differences.
486    
487    The comparison process:
488    1. Parses EJSON strings to JSON objects
489    2. Converts JSON to pandas DataFrames with flattened structure
490    3. Validates both DataFrames for data type consistency
491    4. Detects and normalizes nested arrays in MongoDB results if explode_nested_arrays is True
492    5. Performs fuzzy DataFrame comparison with column matching and value tolerance
493    
494    Args:
495        mdb_A_output (str): First MongoDB query results as an EJSON-formatted string. Should 
496            represent a list of dictionaries or a single dictionary.
497        mdb_B_output (str): Second MongoDB query results as an EJSON-formatted string. Should 
498            represent a list of dictionaries or a single dictionary.
499        sort_values (bool, optional): Whether to sort DataFrame values before 
500            comparison. Set to False if order matters. Defaults to False.
501        allow_column_subset (bool, optional): Whether to allow MongoDB B results to have additional columns not present in MongoDB A results. Defaults to False.
502        verbose (bool, optional): Whether to print detailed comparison messages 
503            during execution. Defaults to False.
504    
505    Returns:
506        FuzzyMatchingResults: A TypedDict containing:
507            - is_match (bool): True if results match within tolerances, False otherwise
508            - message (str): Detailed description of comparison result including:
509                - Success: "Found all matched columns, and all the values are also matched"
510                - Failure: Specific error such as data type inconsistencies, column 
511                  mismatches, or value differences
512    
513    Examples:
514        >>> mdb_A_output = '[{"id": 1, "name": "John"}, {"id": 2, "name": "Jane"}]'
515        >>> mdb_B_output = '[{"_id": 1, "name": "John"}, {"_id": 2, "name": "Jane"}]'
516        >>> result = compare_fuzzy(mdb_A_output, mdb_B_output)
517        >>> print(result['is_match'])
518        True
519        
520        >>> mdb_A_output = '[{"Nationality": "England"}]'
521        >>> mdb_B_output = '["England"]'
522        >>> result = compare_fuzzy(mdb_A_output, mdb_B_output)
523        >>> print(result['message'])
524        Found all matched columns, and all the values are also matched
525    
526    Note:
527        - Uses relative tolerance (rtol=0.01) and absolute tolerance (atol=0.01) 
528          for numeric comparisons
529        - Handles BSON special types (e.g., Decimal128) via EJSON parsing
530        - Automatically flattens nested JSON structures using underscore separators
531        - Normalizes date and numeric columns before comparison
532        - Column matching is based on sorted values, not column names
533        - MongoDB results with nested arrays are exploded into separate rows
534    
535    Raises:
536        EjsonParseException: If input strings cannot be parsed as valid EJSON
537    """
538    # Parse ejson strings to JSON objects
539    mdb_A_result = parse_to_ejson(mdb_A_output)
540    mdb_B_result = parse_to_ejson(mdb_B_output)
541
542    if len(mdb_A_result) != len(mdb_B_result):
543        message = f"Row count mismatch: mdb_A result has {len(mdb_A_result)} rows, mdb_B result has {len(mdb_B_result)} rows."
544        if verbose:
545            print(message)
546        return FuzzyMatchingResults(is_match=False, message=message)
547
548    # Convert JSON to DataFrames
549    mdb_A_df = json_to_dataframe(mdb_A_result)
550    mdb_B_df = json_to_dataframe(mdb_B_result)
551    if not is_valid(mdb_B_df):
552        return FuzzyMatchingResults(is_match=False, message="mdb_B results contain inconsistent data types within columns.")
553
554    # 1st check
555    match_non_explode, message_non_explode = compare_df(mdb_A_df, mdb_B_df, sort_values=sort_values, rtol=rtol, atol=atol, allow_column_subset=allow_column_subset, verbose=verbose)
556    if match_non_explode:
557        return FuzzyMatchingResults(is_match=match_non_explode, message=message_non_explode)
558
559    # check for nested arrays, and 
560    # normalize if found (creates a tidy table using pd.json_normalize and pd.explode)
561    if (check_for_nested_arrays(mdb_A_result) or check_for_nested_arrays(mdb_B_result)) and explode_nested_arrays:
562        mdb_A_df = json_normalize_with_nested_lists(mdb_A_result)
563        mdb_B_df = json_normalize_with_nested_lists(mdb_B_result)
564
565        # 2nd check after normalizing nested arrays
566        match_explode, message_explode = compare_df(mdb_A_df, mdb_B_df, sort_values=sort_values, rtol=rtol, atol=atol, allow_column_subset=allow_column_subset, verbose=verbose)
567        if match_explode:
568            return FuzzyMatchingResults(is_match=match_explode, message=f"[2nd check - exploded nested array]\n{message_explode}")
569        else:
570            return FuzzyMatchingResults(is_match=match_explode, message=f"[1st check - before exploding nested array]\n{message_non_explode}\n[2nd check - exploded nested array]\n{message_explode}")
571    else:
572        return FuzzyMatchingResults(is_match=match_non_explode, message=message_non_explode)
573
574
575def query_mongodb(mongo_uri: str, 
576                  mongo_db_name: str, 
577                  collection_name: str,
578                  aggregation_pipeline: list, 
579                  exclude_fields: List[str],
580                  timeout=60, 
581                  verbose: bool=False) -> list[Dict[str, Any]]:
582    """
583    Execute a query against a MongoDB collection and return the results as a list of dictionaries.
584    
585    This function supports two query modes:
586    1. If aggregation_pipeline is empty: performs a find() query on all documents,
587       excluding specified fields
588    2. If aggregation_pipeline is provided: runs the aggregation pipeline,
589       ignoring exclude_fields parameter
590    
591    Args:
592        mongo_uri (str): MongoDB connection URI string
593        mongo_db_name (str): Name of the MongoDB database
594        collection_name (str): Name of the collection to query
595        aggregation_pipeline (list): List of aggregation pipeline stages. If empty,
596            performs a simple find() query instead.
597        exclude_fields (List[str]): List of field names to exclude from find() results.
598            Ignored when aggregation_pipeline is provided.
599    
600    Returns:
601        list[Dict[str, Any]]: List of documents as dictionaries, or empty list if error occurs
602    
603    Note:
604        - Connection is automatically closed after query execution
605        - Errors are logged to stdout and an empty list is returned
606        - For find() queries, _id field can be excluded via exclude_fields
607        - Aggregation pipelines have full control over returned fields
608    """
609    client = None
610    try:
611        client = MongoClient(mongo_uri)
612        db = client[mongo_db_name]
613        collection = db[collection_name]
614        
615        if aggregation_pipeline:
616            # Run aggregation pipeline, ignore exclude_fields
617            results = list(collection.aggregate(aggregation_pipeline, maxTimeMS=timeout*1000))
618        else:
619            # Perform find query with field exclusions
620            projection = {field: 0 for field in exclude_fields} if exclude_fields else None
621            results = list(collection.find({}, projection))
622        
623        if verbose:
624            print(f"MongoDB query returned {len(results)} documents.")
625        if client:
626            client.close()
627        return results
628        
629    except Exception as e:
630        if verbose:
631            print(f"Error querying MongoDB: {e}")
632        if client:
633            client.close()
634        return None
635
636if __name__ == "__main__":
637    # correct case
638    mdb_A_output = '[{"latitude": 2.76083, "longitude": 101.738}]'
639    mdb_B_output = '[{"location": "Kuala Lumpur", "lat": 2.76083, "lng": 101.738}]'
640    result = compare_fuzzy(mdb_A_output, mdb_B_output, allow_column_subset=True)
641    print(result['is_match'])
642    print(result['message'])
643    
644    # incorrect case, latitude is different more than atol=0.01
645    mdb_A_output = '[{"latitude": 2.76083, "longitude": 101.738}]'
646    mdb_B_output = '[{"location": "Kuala Lumpur", "lat": 2.8, "lng": 101.738}]'
647    result = compare_fuzzy(mdb_A_output, mdb_B_output, allow_column_subset=True)
648    print(result['is_match'])
649    print(result['message'])
650