Skip to content

Commit

Permalink
fix(cubesql): Correctly type pg_index fields
Browse files Browse the repository at this point in the history
  • Loading branch information
@MazterQyou
MazterQyou authored Jul 3, 2024
1 parent f7efec9 commit edadaf1
Show file tree
Hide file tree
Showing 4 changed files with 139 additions and 19 deletions.
38 changes: 25 additions & 13 deletions rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/pg_index.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use async_trait::async_trait;
use datafusion::{
arrow::{
array::{
Array, ArrayRef, BooleanBuilder, Int64Builder, ListBuilder, StringBuilder,
Array, ArrayRef, BooleanBuilder, Int16Builder, ListBuilder, StringBuilder,
UInt16Builder, UInt32Builder,
},
datatypes::{DataType, Field, Schema, SchemaRef},
Expand All @@ -31,10 +31,10 @@ struct PgCatalogIndexBuilder {
indisready: BooleanBuilder,
indislive: BooleanBuilder,
indisreplident: BooleanBuilder,
indkey: ListBuilder<Int64Builder>,
indcollation: StringBuilder,
indclass: StringBuilder,
indoption: StringBuilder,
indkey: ListBuilder<Int16Builder>,
indcollation: ListBuilder<UInt32Builder>,
indclass: ListBuilder<UInt32Builder>,
indoption: ListBuilder<Int16Builder>,
indexprs: StringBuilder,
indpred: StringBuilder,
}
Expand All @@ -58,10 +58,10 @@ impl PgCatalogIndexBuilder {
indisready: BooleanBuilder::new(capacity),
indislive: BooleanBuilder::new(capacity),
indisreplident: BooleanBuilder::new(capacity),
indkey: ListBuilder::new(Int64Builder::new(capacity)),
indcollation: StringBuilder::new(capacity),
indclass: StringBuilder::new(capacity),
indoption: StringBuilder::new(capacity),
indkey: ListBuilder::new(Int16Builder::new(capacity)),
indcollation: ListBuilder::new(UInt32Builder::new(capacity)),
indclass: ListBuilder::new(UInt32Builder::new(capacity)),
indoption: ListBuilder::new(Int16Builder::new(capacity)),
indexprs: StringBuilder::new(capacity),
indpred: StringBuilder::new(capacity),
}
Expand Down Expand Up @@ -136,12 +136,24 @@ impl TableProvider for PgCatalogIndexProvider {
Field::new("indisreplident", DataType::Boolean, false),
Field::new(
"indkey",
DataType::List(Box::new(Field::new("item", DataType::Int64, true))),
DataType::List(Box::new(Field::new("item", DataType::Int16, true))),
false,
),
Field::new(
"indcollation",
DataType::List(Box::new(Field::new("item", DataType::UInt32, true))),
false,
),
Field::new(
"indclass",
DataType::List(Box::new(Field::new("item", DataType::UInt32, true))),
false,
),
Field::new(
"indoption",
DataType::List(Box::new(Field::new("item", DataType::Int16, true))),
false,
),
Field::new("indcollation", DataType::Utf8, false),
Field::new("indclass", DataType::Utf8, false),
Field::new("indoption", DataType::Utf8, false),
Field::new("indexprs", DataType::Utf8, true),
Field::new("indpred", DataType::Utf8, true),
]))
Expand Down
28 changes: 22 additions & 6 deletions rust/cubesql/cubesql/src/compile/engine/udf/common.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2659,6 +2659,13 @@ pub fn create_pg_expandarray_udtf() -> TableUDF {

for i in 0..arr.len() {
let values = arr.value(i);
let values = match values.data_type() {
DataType::Int64 => values,
DataType::Int16 => cast(&values, &DataType::Int64)?,
_ => {
return Err(DataFusionError::Internal("information_schema._pg_expandarray only supports inputs of type List<Int16> or List<Int64>".to_string()))
}
};
let values_arr = values.as_any().downcast_ref::<Int64Array>();
if values_arr.is_none() {
return Err(DataFusionError::Execution(format!("Unsupported type")));
Expand Down Expand Up @@ -2688,16 +2695,25 @@ pub fn create_pg_expandarray_udtf() -> TableUDF {
});

let return_type: ReturnTypeFunction =
// FIXME: it is likely more correct to analyze List type and match type of `x`.
// For now, in order to avoid unexpected breakages, this is left as `Int64`.
Arc::new(move |_| Ok(Arc::new(DataType::Struct(fields()))));

TableUDF::new(
"information_schema._pg_expandarray",
&Signature::exact(
vec![DataType::List(Box::new(Field::new(
"item",
DataType::Int64,
true,
)))],
&Signature::one_of(
vec![
TypeSignature::Exact(vec![DataType::List(Box::new(Field::new(
"item",
DataType::Int64,
true,
)))]),
TypeSignature::Exact(vec![DataType::List(Box::new(Field::new(
"item",
DataType::Int16,
true,
)))]),
],
Volatility::Immutable,
),
&return_type,
Expand Down
84 changes: 84 additions & 0 deletions rust/cubesql/cubesql/src/compile/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24111,4 +24111,88 @@ LIMIT {{ limit }}{% endif %}"#.to_string(),
}
)
}

#[tokio::test]
async fn test_metabase_introspection_indoption() -> Result<(), CubeError> {
init_logger();

insta::assert_snapshot!(
"metabase_introspection_indoption",
execute_query(
r#"
SELECT
tmp.TABLE_CAT,
tmp.TABLE_SCHEM,
tmp.TABLE_NAME,
tmp.NON_UNIQUE,
tmp.INDEX_QUALIFIER,
tmp.INDEX_NAME,
tmp.TYPE,
tmp.ORDINAL_POSITION,
trim(
both '"'
from
pg_catalog.pg_get_indexdef(tmp.CI_OID, tmp.ORDINAL_POSITION, false)
) AS COLUMN_NAME,
CASE
tmp.AM_NAME
WHEN 'btree' THEN CASE
tmp.I_INDOPTION [tmp.ORDINAL_POSITION - 1] & 1 :: smallint
WHEN 1 THEN 'D'
ELSE 'A'
END
ELSE NULL
END AS ASC_OR_DESC,
tmp.CARDINALITY,
tmp.PAGES,
tmp.FILTER_CONDITION
FROM
(
SELECT
NULL AS TABLE_CAT,
n.nspname AS TABLE_SCHEM,
ct.relname AS TABLE_NAME,
NOT i.indisunique AS NON_UNIQUE,
NULL AS INDEX_QUALIFIER,
ci.relname AS INDEX_NAME,
CASE
i.indisclustered
WHEN true THEN 1
ELSE CASE
am.amname
WHEN 'hash' THEN 2
ELSE 3
END
END AS TYPE,
(information_schema._pg_expandarray(i.indkey)).n AS ORDINAL_POSITION,
ci.reltuples AS CARDINALITY,
ci.relpages AS PAGES,
pg_catalog.pg_get_expr(i.indpred, i.indrelid) AS FILTER_CONDITION,
ci.oid AS CI_OID,
i.indoption AS I_INDOPTION,
am.amname AS AM_NAME
FROM
pg_catalog.pg_class ct
JOIN pg_catalog.pg_namespace n ON (ct.relnamespace = n.oid)
JOIN pg_catalog.pg_index i ON (ct.oid = i.indrelid)
JOIN pg_catalog.pg_class ci ON (ci.oid = i.indexrelid)
JOIN pg_catalog.pg_am am ON (ci.relam = am.oid)
WHERE
true
AND n.nspname = 'public'
AND ct.relname = 'IT_Assistance_Needed'
) AS tmp
ORDER BY
NON_UNIQUE,
TYPE,
INDEX_NAME,
ORDINAL_POSITION
"#
.to_string(),
DatabaseProtocol::PostgreSQL
)
.await?
);
Ok(())
}
}
8 changes: 8 additions & 0 deletions ...esql/src/compile/snapshots/cubesql__compile__tests__metabase_introspection_indoption.snap
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
---
source: cubesql/src/compile/mod.rs
expression: "execute_query(r#\"\n SELECT\n tmp.TABLE_CAT,\n tmp.TABLE_SCHEM,\n tmp.TABLE_NAME,\n tmp.NON_UNIQUE,\n tmp.INDEX_QUALIFIER,\n tmp.INDEX_NAME,\n tmp.TYPE,\n tmp.ORDINAL_POSITION,\n trim(\n both '\"'\n from\n pg_catalog.pg_get_indexdef(tmp.CI_OID, tmp.ORDINAL_POSITION, false)\n ) AS COLUMN_NAME,\n CASE\n tmp.AM_NAME\n WHEN 'btree' THEN CASE\n tmp.I_INDOPTION [tmp.ORDINAL_POSITION - 1] & 1 :: smallint\n WHEN 1 THEN 'D'\n ELSE 'A'\n END\n ELSE NULL\n END AS ASC_OR_DESC,\n tmp.CARDINALITY,\n tmp.PAGES,\n tmp.FILTER_CONDITION\n FROM\n (\n SELECT\n NULL AS TABLE_CAT,\n n.nspname AS TABLE_SCHEM,\n ct.relname AS TABLE_NAME,\n NOT i.indisunique AS NON_UNIQUE,\n NULL AS INDEX_QUALIFIER,\n ci.relname AS INDEX_NAME,\n CASE\n i.indisclustered\n WHEN true THEN 1\n ELSE CASE\n am.amname\n WHEN 'hash' THEN 2\n ELSE 3\n END\n END AS TYPE,\n (information_schema._pg_expandarray(i.indkey)).n AS ORDINAL_POSITION,\n ci.reltuples AS CARDINALITY,\n ci.relpages AS PAGES,\n pg_catalog.pg_get_expr(i.indpred, i.indrelid) AS FILTER_CONDITION,\n ci.oid AS CI_OID,\n i.indoption AS I_INDOPTION,\n am.amname AS AM_NAME\n FROM\n pg_catalog.pg_class ct\n JOIN pg_catalog.pg_namespace n ON (ct.relnamespace = n.oid)\n JOIN pg_catalog.pg_index i ON (ct.oid = i.indrelid)\n JOIN pg_catalog.pg_class ci ON (ci.oid = i.indexrelid)\n JOIN pg_catalog.pg_am am ON (ci.relam = am.oid)\n WHERE\n true\n AND n.nspname = 'public'\n AND ct.relname = 'IT_Assistance_Needed'\n ) AS tmp\n ORDER BY\n NON_UNIQUE,\n TYPE,\n INDEX_NAME,\n ORDINAL_POSITION\n \"#.to_string(),\n DatabaseProtocol::PostgreSQL).await?"
---
+-----------+-------------+------------+------------+-----------------+------------+------+------------------+-------------+-------------+-------------+-------+------------------+
| TABLE_CAT | TABLE_SCHEM | TABLE_NAME | NON_UNIQUE | INDEX_QUALIFIER | INDEX_NAME | TYPE | ORDINAL_POSITION | COLUMN_NAME | ASC_OR_DESC | CARDINALITY | PAGES | FILTER_CONDITION |
+-----------+-------------+------------+------------+-----------------+------------+------+------------------+-------------+-------------+-------------+-------+------------------+
+-----------+-------------+------------+------------+-----------------+------------+------+------------------+-------------+-------------+-------------+-------+------------------+

0 comments on commit edadaf1

Please sign in to comment.