diff --git a/adbc_drivers_validation/tests/ingest.py b/adbc_drivers_validation/tests/ingest.py index d4b342c..f0d87e4 100644 --- a/adbc_drivers_validation/tests/ingest.py +++ b/adbc_drivers_validation/tests/ingest.py @@ -116,6 +116,14 @@ def make_table_name(prefix: str, query: Query | str) -> str: return _SANITIZE_TABLE_NAME.sub("_", f"{prefix}_{suffix}") +def _sort_table_by_first_column(table: pyarrow.Table) -> pyarrow.Table: + # Sort locally because some vendors, such as Cassandra, do not support + # ORDER BY here. + if table.num_rows <= 1 or len(table.schema) == 0: + return table + return table.sort_by([(table.schema[0].name, "ascending")]) + + class TestIngest: def test_create( self, @@ -142,14 +150,17 @@ def test_create( fields = [] for field in data.schema: fields.append(driver.quote_identifier(field.name)) - select = f"SELECT {', '.join(fields)} FROM {driver.quote_identifier(table_name)} ORDER BY {fields[0]} ASC" + select = ( + f"SELECT {', '.join(fields)} FROM {driver.quote_identifier(table_name)}" + ) with conn.cursor() as cursor: with driver.setup_statement(query, cursor): result = execute_query_without_prepare(cursor, select) # TODO: we should also inspect the type name and make sure it matches the # metadata - expected = subquery.expected() + result = _sort_table_by_first_column(result) + expected = _sort_table_by_first_column(subquery.expected()) compare.compare_tables(expected, result, query.metadata()) subschema = pyarrow.schema(list(result.schema)[1:]) utils.assert_field_type_name(driver, "query", query, subschema) @@ -190,7 +201,7 @@ def test_append( idx = driver.quote_identifier("idx") value = driver.quote_identifier("value") - select = f"SELECT {idx}, {value} FROM {driver.quote_identifier(table_name)} ORDER BY {idx} ASC" + select = f"SELECT {idx}, {value} FROM {driver.quote_identifier(table_name)}" with conn.cursor() as cursor: result = execute_query_without_prepare(cursor, select) @@ -202,8 +213,12 @@ def test_append( }, schema=expected.schema, ) + result = _sort_table_by_first_column(result) + expected = _sort_table_by_first_column( + pyarrow.concat_tables([expected, expected2]) + ) compare.compare_tables( - pyarrow.concat_tables([expected, expected2]), + expected, result, query.metadata(), ) @@ -263,7 +278,7 @@ def test_createappend( idx = driver.quote_identifier("idx") value = driver.quote_identifier("value") - select = f"SELECT {idx}, {value} FROM {driver.quote_identifier(table_name)} ORDER BY {idx} ASC" + select = f"SELECT {idx}, {value} FROM {driver.quote_identifier(table_name)}" with conn.cursor() as cursor: result = execute_query_without_prepare(cursor, select) @@ -277,8 +292,12 @@ def test_createappend( }, schema=expected.schema, ) + result = _sort_table_by_first_column(result) + expected = _sort_table_by_first_column( + pyarrow.concat_tables([expected, expected2]) + ) compare.compare_tables( - pyarrow.concat_tables([expected, expected2]), + expected, result, query.metadata(), ) @@ -373,7 +392,8 @@ def test_replace( idx = driver.quote_identifier("idx") value = driver.quote_identifier("value") - select = f"SELECT {idx}, {value} FROM {driver.quote_identifier(table_name)} ORDER BY {idx} ASC" + # The final replacement contains one row, so ordering is unnecessary. + select = f"SELECT {idx}, {value} FROM {driver.quote_identifier(table_name)}" with conn.cursor() as cursor: result = execute_query_without_prepare(cursor, select) @@ -403,11 +423,12 @@ def test_replace_noop( idx = driver.quote_identifier("idx") value = driver.quote_identifier("value") - select = f"SELECT {idx}, {value} FROM {driver.quote_identifier(table_name)} ORDER BY {idx} ASC" + select = f"SELECT {idx}, {value} FROM {driver.quote_identifier(table_name)}" with conn.cursor() as cursor: result = execute_query_without_prepare(cursor, select) - expected = subquery.expected() + result = _sort_table_by_first_column(result) + expected = _sort_table_by_first_column(subquery.expected()) compare.compare_tables(expected, result, query.metadata()) def test_not_null( @@ -818,7 +839,9 @@ def test_create_multiple_batches( fields = [] for field in data.schema: fields.append(driver.quote_identifier(field.name)) - select = f"SELECT {', '.join(fields)} FROM {driver.quote_identifier(table_name)} ORDER BY {fields[0]} ASC" + select = ( + f"SELECT {', '.join(fields)} FROM {driver.quote_identifier(table_name)}" + ) with conn.cursor() as cursor: result = execute_query_without_prepare(cursor, select) @@ -837,8 +860,10 @@ def test_create_multiple_batches( ) expected_tables.append(expected_i) + result = _sort_table_by_first_column(result) + expected = _sort_table_by_first_column(pyarrow.concat_tables(expected_tables)) compare.compare_tables( - pyarrow.concat_tables(expected_tables), + expected, result, query.metadata(), ) @@ -881,7 +906,9 @@ def test_create_large_batch( fields = [] for field in data.schema: fields.append(driver.quote_identifier(field.name)) - select = f"SELECT {', '.join(fields)} FROM {driver.quote_identifier(table_name)} ORDER BY {fields[0]} ASC" + select = ( + f"SELECT {', '.join(fields)} FROM {driver.quote_identifier(table_name)}" + ) with conn.cursor() as cursor: result = execute_query_without_prepare(cursor, select) @@ -900,8 +927,10 @@ def test_create_large_batch( ) expected_tables.append(expected_i) + result = _sort_table_by_first_column(result) + expected = _sort_table_by_first_column(pyarrow.concat_tables(expected_tables)) compare.compare_tables( - pyarrow.concat_tables(expected_tables), + expected, result, query.metadata(), ) @@ -979,7 +1008,9 @@ def test_ingest_then_query( fields = [] for field in data.schema: fields.append(driver.quote_identifier(field.name)) - select = f"SELECT {', '.join(fields)} FROM {driver.quote_identifier(table_name)} ORDER BY {fields[0]} ASC" + select = ( + f"SELECT {', '.join(fields)} FROM {driver.quote_identifier(table_name)}" + ) with conn.cursor() as cursor: driver.try_drop_table(cursor, table_name=table_name) @@ -992,4 +1023,6 @@ def test_ingest_then_query( result = execute_query_without_prepare(cursor, select) + result = _sort_table_by_first_column(result) + expected = _sort_table_by_first_column(expected) compare.compare_tables(expected, result, query.metadata()) diff --git a/adbc_drivers_validation/tests/statement.py b/adbc_drivers_validation/tests/statement.py index 8e3e07f..8b0fc68 100644 --- a/adbc_drivers_validation/tests/statement.py +++ b/adbc_drivers_validation/tests/statement.py @@ -245,7 +245,10 @@ def test_rows_affected( assert rows_affected == -1 cursor.adbc_statement.set_sql_query( - f"INSERT INTO {quoted_name} (id) VALUES (1)" + driver.query_override( + "TestStatement.test_rows_affected.insert", + f"INSERT INTO {quoted_name} (id) VALUES (1)", + ) ) rows_affected = cursor.adbc_statement.execute_update() if driver.features.statement_rows_affected: @@ -257,7 +260,10 @@ def test_rows_affected( assert rows_affected == -1 cursor.adbc_statement.set_sql_query( - f"UPDATE {quoted_name} SET id = id + 1 WHERE id = 1" + driver.query_override( + "TestStatement.test_rows_affected.update", + f"UPDATE {quoted_name} SET id = id + 1 WHERE id = 1", + ) ) rows_affected = cursor.adbc_statement.execute_update() if driver.features.statement_rows_affected: @@ -269,7 +275,10 @@ def test_rows_affected( assert rows_affected == -1 cursor.adbc_statement.set_sql_query( - f"DELETE FROM {quoted_name} WHERE id = 2" + driver.query_override( + "TestStatement.test_rows_affected.delete", + f"DELETE FROM {quoted_name} WHERE id = 2", + ) ) rows_affected = cursor.adbc_statement.execute_update() if driver.features.statement_rows_affected: