Skip to content

BigQuery: support CREATE VECTOR INDEX #2437

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
moshap-firebolt wants to merge 1 commit into
base: main
Choose a base branch
from
Open

BigQuery: support CREATE VECTOR INDEX #2437

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion src/ast/ddl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2827,6 +2827,10 @@ pub struct CreateIndex {
pub using: Option<IndexType>,
/// columns included in the index
pub columns: Vec<IndexColumn>,
/// whether this is a BigQuery `CREATE VECTOR INDEX`
pub vector: bool,
/// whether the statement is `CREATE OR REPLACE`
pub or_replace: bool,
/// whether the index is unique
pub unique: bool,
/// whether the index is created concurrently
Expand All @@ -2843,6 +2847,8 @@ pub struct CreateIndex {
pub nulls_distinct: Option<bool>,
/// WITH clause: <https://www.postgresql.org/docs/current/sql-createindex.html>
pub with: Vec<Expr>,
/// BigQuery `OPTIONS(...)` clause, e.g. on `CREATE VECTOR INDEX`
pub options: Vec<SqlOption>,
/// WHERE clause: <https://www.postgresql.org/docs/current/sql-createindex.html>
pub predicate: Option<Expr>,
/// Index options: <https://www.postgresql.org/docs/current/sql-createindex.html>
Expand All @@ -2860,8 +2866,10 @@ impl fmt::Display for CreateIndex {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
"CREATE {unique}INDEX {concurrently}{async_}{if_not_exists}",
"CREATE {or_replace}{unique}{vector}INDEX {concurrently}{async_}{if_not_exists}",
or_replace = if self.or_replace { "OR REPLACE " } else { "" },
unique = if self.unique { "UNIQUE " } else { "" },
vector = if self.vector { "VECTOR " } else { "" },
concurrently = if self.concurrently {
"CONCURRENTLY "
} else {
Expand All @@ -2882,6 +2890,9 @@ impl fmt::Display for CreateIndex {
write!(f, " USING {value} ")?;
}
write!(f, "({})", display_comma_separated(&self.columns))?;
if !self.options.is_empty() {
write!(f, " OPTIONS({})", display_comma_separated(&self.options))?;
}
if !self.include.is_empty() {
write!(f, " INCLUDE ({})", display_comma_separated(&self.include))?;
}
Expand Down
4 changes: 4 additions & 0 deletions src/ast/spans.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -699,13 +699,16 @@ impl Spanned for CreateIndex {
table_name,
using: _,
columns,
vector: _, // bool
or_replace: _, // bool
unique: _, // bool
concurrently: _, // bool
r#async: _, // bool
if_not_exists: _, // bool
include,
nulls_distinct: _, // bool
with,
options,
predicate,
index_options: _,
alter_options,
Expand All @@ -718,6 +721,7 @@ impl Spanned for CreateIndex {
.chain(columns.iter().map(|i| i.column.span()))
.chain(include.iter().map(|i| i.span))
.chain(with.iter().map(|i| i.span()))
.chain(options.iter().map(|i| i.span()))
.chain(predicate.iter().map(|i| i.span()))
.chain(alter_options.iter().map(|i| i.span())),
)
Expand Down
40 changes: 40 additions & 0 deletions src/parser/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5288,6 +5288,13 @@ impl<'a> Parser<'a> {
self.parse_create_schema(or_replace)
} else if self.parse_keyword(Keyword::WAREHOUSE) {
self.parse_create_warehouse(or_replace).map(Into::into)
} else if matches!(
&self.peek_token_ref().token,
Token::Word(w) if w.keyword == Keyword::NoKeyword && w.value.eq_ignore_ascii_case("VECTOR")
) {
// BigQuery `CREATE [OR REPLACE] VECTOR INDEX ...`; VECTOR is not a keyword.
self.next_token();
self.parse_create_vector_index(or_replace).map(Into::into)
} else if or_replace {
self.expected_ref(
"[EXTERNAL] TABLE or [MATERIALIZED] VIEW or FUNCTION or SCHEMA or WAREHOUSE after CREATE OR REPLACE",
Expand Down Expand Up @@ -8245,6 +8252,36 @@ impl<'a> Parser<'a> {
Ok(Statement::Discard { object_type })
}

/// Parse a BigQuery `CREATE [OR REPLACE] VECTOR INDEX` statement (`VECTOR` already consumed).
fn parse_create_vector_index(&mut self, or_replace: bool) -> Result<CreateIndex, ParserError> {
self.expect_keyword_is(Keyword::INDEX)?;
let if_not_exists = self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
let name = self.parse_object_name(false)?;
self.expect_keyword_is(Keyword::ON)?;
let table_name = self.parse_object_name(false)?;
let columns = self.parse_parenthesized_index_column_list()?;
let options = self.parse_options(Keyword::OPTIONS)?;
Ok(CreateIndex {
name: Some(name),
table_name,
using: None,
columns,
vector: true,
or_replace,
unique: false,
concurrently: false,
r#async: false,
if_not_exists,
include: vec![],
nulls_distinct: None,
with: vec![],
options,
predicate: None,
index_options: vec![],
alter_options: vec![],
})
}

/// Parse a `CREATE INDEX` statement.
pub fn parse_create_index(&mut self, unique: bool) -> Result<CreateIndex, ParserError> {
let concurrently = self.parse_keyword(Keyword::CONCURRENTLY);
Expand Down Expand Up @@ -8326,13 +8363,16 @@ impl<'a> Parser<'a> {
table_name,
using,
columns,
vector: false,
or_replace: false,
unique,
concurrently,
r#async,
if_not_exists,
include,
nulls_distinct,
with,
options: vec![],
predicate,
index_options,
alter_options,
Expand Down
77 changes: 77 additions & 0 deletions tests/sqlparser_bigquery.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2950,3 +2950,80 @@ fn test_create_snapshot_table() {
"CREATE SNAPSHOT TABLE IF NOT EXISTS dataset_id.table1 CLONE dataset_id.table2 FOR SYSTEM_TIME AS OF TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 1 HOUR) OPTIONS(expiration_timestamp = TIMESTAMP '2025-01-01 00:00:00 UTC')",
);
}
#[test]
fn parse_bigquery_create_vector_index() {
// `CREATE VECTOR INDEX ... OPTIONS(...)` is a CreateIndex flagged `vector`.
let sql =
"CREATE VECTOR INDEX emb ON t(embedding) OPTIONS(distance_type = 'COSINE', dimension = 4)";
match bigquery().verified_stmt(sql) {
Statement::CreateIndex(CreateIndex {
name,
table_name,
using,
columns,
vector,
or_replace,
unique,
if_not_exists,
with,
options,
..
}) => {
assert!(vector);
assert!(!or_replace);
assert!(!unique);
assert!(!if_not_exists);
assert_eq!(name.unwrap().to_string(), "emb");
assert_eq!(table_name.to_string(), "t");
assert_eq!(using, None);
assert!(with.is_empty());
assert_eq!(
columns.iter().map(|c| c.to_string()).collect::<Vec<_>>(),
vec!["embedding"]
);
assert_eq!(
options,
vec![
SqlOption::KeyValue {
key: Ident::new("distance_type"),
value: Expr::Value(
Value::SingleQuotedString("COSINE".to_string()).with_empty_span()
),
},
SqlOption::KeyValue {
key: Ident::new("dimension"),
value: Expr::value(number("4")),
},
]
);
}
other => panic!("expected CreateIndex, got {other:?}"),
}

// OR REPLACE, IF NOT EXISTS, multi-part names, and the OPTIONS-less form all round-trip.
match bigquery().verified_stmt(
"CREATE OR REPLACE VECTOR INDEX emb ON t(embedding) OPTIONS(distance_type = 'COSINE')",
) {
Statement::CreateIndex(CreateIndex {
vector, or_replace, ..
}) => {
assert!(vector);
assert!(or_replace);
}
other => panic!("expected CreateIndex, got {other:?}"),
}
match bigquery().verified_stmt(
"CREATE VECTOR INDEX IF NOT EXISTS mydataset.emb ON mydataset.t(embedding) OPTIONS(distance_type = 'EUCLIDEAN')",
) {
Statement::CreateIndex(CreateIndex {
vector,
if_not_exists,
..
}) => {
assert!(vector);
assert!(if_not_exists);
}
other => panic!("expected CreateIndex, got {other:?}"),
}
bigquery().verified_stmt("CREATE VECTOR INDEX emb ON t(embedding)");
}
6 changes: 6 additions & 0 deletions tests/sqlparser_common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9728,6 +9728,9 @@ fn test_create_index_with_using_function() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq!("idx_name", name.to_string());
assert_eq!("test", table_name.to_string());
Expand Down Expand Up @@ -9785,6 +9788,9 @@ fn test_create_index_with_with_clause() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
pretty_assertions::assert_eq!("title_idx", name.to_string());
pretty_assertions::assert_eq!("films", table_name.to_string());
Expand Down
33 changes: 33 additions & 0 deletions tests/sqlparser_postgres.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2993,6 +2993,9 @@ fn parse_create_index() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq_vec(&["my_index"], &name);
assert_eq_vec(&["my_table"], &table_name);
Expand Down Expand Up @@ -3030,6 +3033,9 @@ fn parse_create_anonymous_index() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq!(None, name);
assert_eq_vec(&["my_table"], &table_name);
Expand Down Expand Up @@ -3150,6 +3156,9 @@ fn parse_create_indices_with_operator_classes() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq_vec(&["the_index_name"], &name);
assert_eq_vec(&["users"], &table_name);
Expand Down Expand Up @@ -3179,6 +3188,9 @@ fn parse_create_indices_with_operator_classes() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq_vec(&["the_index_name"], &name);
assert_eq_vec(&["users"], &table_name);
Expand Down Expand Up @@ -3263,6 +3275,9 @@ fn parse_create_bloom() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq_vec(&["bloomidx"], &name);
assert_eq_vec(&["tbloom"], &table_name);
Expand Down Expand Up @@ -3320,6 +3335,9 @@ fn parse_create_brin() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq_vec(&["brin_sensor_data_recorded_at"], &name);
assert_eq_vec(&["sensor_data"], &table_name);
Expand Down Expand Up @@ -3388,6 +3406,9 @@ fn parse_create_index_concurrently() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq_vec(&["my_index"], &name);
assert_eq_vec(&["my_table"], &table_name);
Expand Down Expand Up @@ -3425,6 +3446,9 @@ fn parse_create_index_with_predicate() {
predicate: Some(_),
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq_vec(&["my_index"], &name);
assert_eq_vec(&["my_table"], &table_name);
Expand Down Expand Up @@ -3462,6 +3486,9 @@ fn parse_create_index_with_include() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq_vec(&["my_index"], &name);
assert_eq_vec(&["my_table"], &table_name);
Expand Down Expand Up @@ -3499,6 +3526,9 @@ fn parse_create_index_with_nulls_distinct() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq_vec(&["my_index"], &name);
assert_eq_vec(&["my_table"], &table_name);
Expand Down Expand Up @@ -3534,6 +3564,9 @@ fn parse_create_index_with_nulls_distinct() {
predicate: None,
index_options,
alter_options,
vector: _,
or_replace: _,
options: _,
}) => {
assert_eq_vec(&["my_index"], &name);
assert_eq_vec(&["my_table"], &table_name);
Expand Down
Loading