Skip to content

Commit

Permalink
feat(engine): handle target.* (#933)
Browse files Browse the repository at this point in the history
  • Loading branch information
@grieve54706
grieve54706 authored Nov 21, 2024
1 parent 5d9db4c commit e23469d
Show file tree
Hide file tree
Showing 12 changed files with 174 additions and 53 deletions.
56 changes: 56 additions & 0 deletions ibis-server/tests/routers/v2/connector/test_postgres.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,36 @@
],
"primaryKey": "orderkey",
},
{
"name": "Customer",
"refSql": "SELECT * FROM public.customer",
"columns": [
{
"name": "custkey",
"type": "integer",
"expression": "c_custkey",
},
{
"name": "orders",
"type": "Orders",
"relationship": "CustomerOrders",
},
{
"name": "orders_key",
"type": "varchar",
"isCalculated": True,
"expression": "orders.orderkey",
},
],
},
],
"relationships": [
{
"name": "CustomerOrders",
"models": ["Customer", "Orders"],
"joinType": "ONE_TO_MANY",
"condition": "Customer.custkey = Orders.custkey",
}
],
}

Expand All @@ -83,6 +113,9 @@ def postgres(request) -> PostgresContainer:
pd.read_parquet(file_path("resource/tpch/data/orders.parquet")).to_sql(
"orders", engine, index=False
)
pd.read_parquet(file_path("resource/tpch/data/customer.parquet")).to_sql(
"customer", engine, index=False
)
with engine.begin() as conn:
conn.execute(text("COMMENT ON TABLE orders IS 'This is a table comment'"))
conn.execute(text("COMMENT ON COLUMN orders.o_comment IS 'This is a comment'"))
Expand Down Expand Up @@ -148,6 +181,29 @@ def test_query_with_connection_url(manifest_str, postgres: PostgresContainer):
assert result["data"][0][0] == 1
assert result["dtypes"] is not None

def test_query_with_dot_all(manifest_str, postgres: PostgresContainer):
connection_info = _to_connection_info(postgres)
test_sqls = [
'SELECT "Customer".* FROM "Customer"',
'SELECT c.* FROM "Customer" AS c',
'SELECT c.* FROM "Customer" AS c JOIN "Orders" AS o ON c.custkey = o.custkey',
]
for sql in test_sqls:
response = client.post(
url=f"{base_url}/query",
params={"limit": 1},
json={
"connectionInfo": connection_info,
"manifestStr": manifest_str,
"sql": sql,
},
)
assert response.status_code == 200
result = response.json()
assert len(result["columns"]) == 1 # Not include calculated column
assert len(result["data"]) == 1
assert result["dtypes"] is not None

def test_dry_run_with_connection_url_and_password_with_bracket_should_not_raise_value_error(
manifest_str, postgres: PostgresContainer
):
Expand Down
2 changes: 1 addition & 1 deletion wren-base/src/main/java/io/wren/base/dto/Column.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ public static Column relationshipColumn(String name, String type, String relatio
return new Column(name, type, relationship, false, false, null, null);
}

public static Column caluclatedColumn(String name, String type, String expression)
public static Column calculatedColumn(String name, String type, String expression)
{
return new Column(name, type, null, true, false, expression, null);
}
Expand Down
32 changes: 19 additions & 13 deletions wren-base/src/main/java/io/wren/base/sqlrewrite/analyzer/StatementAnalyzer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Stream;

import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.ImmutableList.toImmutableList;
Expand Down Expand Up @@ -374,21 +375,26 @@ else if (item instanceof SingleColumn) {

private void analyzeSelectAllColumns(AllColumns allColumns, Scope scope, ImmutableList.Builder<Expression> outputExpressions)
{
List<Field> fields = scope.getRelationType().getFields();
Stream<Field> collectedColumns = fields.stream().filter(f -> f.getSourceColumn().map(c -> !c.isCalculated()).orElse(true));
Stream<Field> outputExpressionStream = fields.stream();

if (allColumns.getTarget().isPresent()) {
// TODO handle target.*
}
else {
List<Field> fields = scope.getRelationType()
.getFields()
.stream()
.filter(f -> f.getSourceColumn().map(c -> !c.isCalculated()).orElse(true))
.collect(toImmutableList());
analysis.addCollectedColumns(fields);
scope.getRelationType().getFields().stream().map(field ->
field.getRelationAlias().map(DereferenceExpression::from)
.orElse(DereferenceExpression.from(QualifiedName.of(field.getTableName().getSchemaTableName().getTableName(), field.getColumnName()))))
.forEach(outputExpressions::add);
QualifiedName target = QualifiedName.of(((Identifier) allColumns.getTarget().get()).getValue());
collectedColumns = collectedColumns
.filter(field ->
field.getRelationAlias().map(ra -> ra.equals(target)).orElse(false)
|| field.getTableName().getSchemaTableName().getTableName().equals(target.getParts().getFirst()));
outputExpressionStream = outputExpressionStream
.filter(field -> field.getRelationAlias().map(ra -> ra.equals(target)).orElse(false));
}

analysis.addCollectedColumns(collectedColumns.collect(toImmutableList()));
outputExpressionStream
.map(field ->
field.getRelationAlias().map(DereferenceExpression::from)
.orElse(DereferenceExpression.from(QualifiedName.of(field.getTableName().getSchemaTableName().getTableName(), field.getColumnName()))))
.forEach(outputExpressions::add);
}

private void analyzeSelectSingleColumn(SingleColumn singleColumn, Scope scope, ImmutableList.Builder<Expression> outputExpressions)
Expand Down
34 changes: 17 additions & 17 deletions wren-base/src/test/java/io/wren/base/sqlrewrite/AbstractTestModel.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,10 +113,10 @@ public void testToManyCalculated()
Model newCustomer = addColumnsToModel(
customer,
Column.column("orders", "Orders", "OrdersCustomer", true),
Column.caluclatedColumn("totalprice", WrenTypes.BIGINT, "sum(orders.totalprice)"),
Column.caluclatedColumn("buy_item_count", WrenTypes.BIGINT, "count(distinct orders.lineitem.orderkey_linenumber)"),
Column.caluclatedColumn("lineitem_totalprice", WrenTypes.BIGINT, "sum(orders.lineitem.discount * orders.lineitem.extendedprice)"),
Column.caluclatedColumn("test_col", WrenTypes.BIGINT, "sum(orders.lineitem.discount * nationkey)"));
Column.calculatedColumn("totalprice", WrenTypes.BIGINT, "sum(orders.totalprice)"),
Column.calculatedColumn("buy_item_count", WrenTypes.BIGINT, "count(distinct orders.lineitem.orderkey_linenumber)"),
Column.calculatedColumn("lineitem_totalprice", WrenTypes.BIGINT, "sum(orders.lineitem.discount * orders.lineitem.extendedprice)"),
Column.calculatedColumn("test_col", WrenTypes.BIGINT, "sum(orders.lineitem.discount * nationkey)"));
Manifest manifest = withDefaultCatalogSchema()
.setModels(List.of(newCustomer, orders, lineitem))
.setRelationships(List.of(ordersCustomer, ordersLineitem))
Expand Down Expand Up @@ -156,8 +156,8 @@ public void testToOneCalculated()
Model newLineitem = addColumnsToModel(
lineitem,
Column.column("orders", "Orders", "OrdersLineitem", true),
Column.caluclatedColumn("col_1", WrenTypes.BIGINT, "orders.totalprice + orders.totalprice"),
Column.caluclatedColumn("col_2", WrenTypes.BIGINT, "concat(orders.orderkey, '#', orders.customer.custkey)"));
Column.calculatedColumn("col_1", WrenTypes.BIGINT, "orders.totalprice + orders.totalprice"),
Column.calculatedColumn("col_2", WrenTypes.BIGINT, "concat(orders.orderkey, '#', orders.customer.custkey)"));
Model newOrders = addColumnsToModel(
orders,
Column.column("customer", "Customer", "OrdersCustomer", true));
Expand Down Expand Up @@ -221,11 +221,11 @@ public void testModelWithCycle()
Model newCustomer = addColumnsToModel(
customer,
Column.column("orders", "Orders", "OrdersCustomer", true),
Column.caluclatedColumn("total_price", WrenTypes.BIGINT, "sum(orders.totalprice)"));
Column.calculatedColumn("total_price", WrenTypes.BIGINT, "sum(orders.totalprice)"));
Model newOrders = addColumnsToModel(
orders,
Column.column("customer", "Customer", "OrdersCustomer", true),
Column.caluclatedColumn("customer_name", WrenTypes.BIGINT, "customer.name"));
Column.calculatedColumn("customer_name", WrenTypes.BIGINT, "customer.name"));
Manifest manifest = withDefaultCatalogSchema()
.setModels(List.of(newCustomer, newOrders))
.setRelationships(List.of(ordersCustomer, ordersLineitem))
Expand All @@ -252,7 +252,7 @@ public void testModelOnModel()
Model newCustomer = addColumnsToModel(
customer,
Column.column("orders", "Orders", "OrdersCustomer", true),
Column.caluclatedColumn("totalprice", WrenTypes.BIGINT, "sum(orders.totalprice)"));
Column.calculatedColumn("totalprice", WrenTypes.BIGINT, "sum(orders.totalprice)"));
Model onCustomer = Model.onBaseObject(
"OnCustomer",
"Customer",
Expand Down Expand Up @@ -282,17 +282,17 @@ public void testCalculatedUseAnotherCalculated()
Model newCustomer = addColumnsToModel(
customer,
Column.column("orders", "Orders", "OrdersCustomer", true),
Column.caluclatedColumn("total_price", WrenTypes.BIGINT, "sum(orders.totalprice)"));
Column.calculatedColumn("total_price", WrenTypes.BIGINT, "sum(orders.totalprice)"));
Model newOrders = addColumnsToModel(
orders,
Column.column("customer", "Customer", "OrdersCustomer", true),
Column.column("lineitem", "Lineitem", "OrdersLineitem", true),
Column.caluclatedColumn("customer_name", WrenTypes.BIGINT, "customer.name"),
Column.caluclatedColumn("extended_price", WrenTypes.BIGINT, "sum(lineitem.extendedprice)"));
Column.calculatedColumn("customer_name", WrenTypes.BIGINT, "customer.name"),
Column.calculatedColumn("extended_price", WrenTypes.BIGINT, "sum(lineitem.extendedprice)"));
Model newLineitem = addColumnsToModel(
lineitem,
Column.column("orders", "Orders", "OrdersLineitem", true),
Column.caluclatedColumn("test_column", WrenTypes.BIGINT, "orders.customer.total_price + extendedprice"));
Column.calculatedColumn("test_column", WrenTypes.BIGINT, "orders.customer.total_price + extendedprice"));
Manifest manifest = withDefaultCatalogSchema()
.setModels(List.of(newCustomer, newOrders, newLineitem))
.setRelationships(List.of(ordersCustomer, ordersLineitem))
Expand All @@ -309,17 +309,17 @@ public void testSelectEmpty()
Model newCustomer = addColumnsToModel(
customer,
Column.column("orders", "Orders", "OrdersCustomer", true),
Column.caluclatedColumn("total_price", WrenTypes.BIGINT, "sum(orders.totalprice)"));
Column.calculatedColumn("total_price", WrenTypes.BIGINT, "sum(orders.totalprice)"));
Model newOrders = addColumnsToModel(
orders,
Column.column("customer", "Customer", "OrdersCustomer", true),
Column.column("lineitem", "Lineitem", "OrdersLineitem", true),
Column.caluclatedColumn("customer_name", WrenTypes.BIGINT, "customer.name"),
Column.caluclatedColumn("extended_price", WrenTypes.BIGINT, "sum(lineitem.extendedprice)"));
Column.calculatedColumn("customer_name", WrenTypes.BIGINT, "customer.name"),
Column.calculatedColumn("extended_price", WrenTypes.BIGINT, "sum(lineitem.extendedprice)"));
Model newLineitem = addColumnsToModel(
lineitem,
Column.column("orders", "Orders", "OrdersLineitem", true),
Column.caluclatedColumn("test_column", WrenTypes.BIGINT, "orders.customer.total_price + extendedprice"));
Column.calculatedColumn("test_column", WrenTypes.BIGINT, "orders.customer.total_price + extendedprice"));
Manifest manifest = withDefaultCatalogSchema()
.setModels(List.of(newCustomer, newOrders, newLineitem))
.setRelationships(List.of(ordersCustomer, ordersLineitem))
Expand Down
2 changes: 1 addition & 1 deletion wren-base/src/test/java/io/wren/base/sqlrewrite/TestAllRulesRewrite.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ public TestAllRulesRewrite()
Column.relationshipColumn("band", "Band", "AlbumBand"),
Column.column("price", WrenTypes.INTEGER, null, true),
Column.column("bandId", WrenTypes.INTEGER, null, true),
Column.caluclatedColumn("bandName", WrenTypes.VARCHAR, "band.name"),
Column.calculatedColumn("bandName", WrenTypes.VARCHAR, "band.name"),
Column.column("status", "Inventory", null, true),
Column.column("statusA", "InventoryA", null, true),
Column.relationshipColumn("orders", "Order", "AlbumOrder")),
Expand Down
4 changes: 2 additions & 2 deletions wren-base/src/test/java/io/wren/base/sqlrewrite/TestMetric.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,8 @@ public TestMetric()
Column.column("clerk", WrenTypes.VARCHAR, null, true),
Column.column("shippriority", WrenTypes.INTEGER, null, true),
Column.column("comment", WrenTypes.VARCHAR, null, true),
Column.caluclatedColumn("customer_name", WrenTypes.VARCHAR, "customer.name"),
Column.caluclatedColumn("cumstomer_address", WrenTypes.VARCHAR, "customer.address"),
Column.calculatedColumn("customer_name", WrenTypes.VARCHAR, "customer.name"),
Column.calculatedColumn("cumstomer_address", WrenTypes.VARCHAR, "customer.address"),
Column.column("customer", "Customer", "OrdersCustomer", true),
Column.column("lineitem", "Lineitem", "OrdersLineitem", true)),
"orderkey"),
Expand Down
20 changes: 10 additions & 10 deletions wren-base/src/test/java/io/wren/base/sqlrewrite/TestWrenDataLineage.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,20 +108,20 @@ public void testAnalyze()
Model newCustomer = addColumnsToModel(
customer,
Column.column("orders", "Orders", "OrdersCustomer", true),
Column.caluclatedColumn("total_price", WrenTypes.BIGINT, "sum(orders.totalprice)"),
Column.caluclatedColumn("discount_extended_price", WrenTypes.BIGINT, "sum(orders.lineitem.discount + orders.extended_price)"),
Column.caluclatedColumn("lineitem_price", WrenTypes.BIGINT, "sum(orders.lineitem.discount * orders.lineitem.extendedprice)"));
Column.calculatedColumn("total_price", WrenTypes.BIGINT, "sum(orders.totalprice)"),
Column.calculatedColumn("discount_extended_price", WrenTypes.BIGINT, "sum(orders.lineitem.discount + orders.extended_price)"),
Column.calculatedColumn("lineitem_price", WrenTypes.BIGINT, "sum(orders.lineitem.discount * orders.lineitem.extendedprice)"));
Model newOrders = addColumnsToModel(
orders,
Column.column("customer", "Customer", "OrdersCustomer", true),
Column.column("lineitem", "Lineitem", "OrdersLineitem", true),
Column.caluclatedColumn("customer_name", WrenTypes.BIGINT, "customer.name"),
Column.caluclatedColumn("extended_price", WrenTypes.BIGINT, "sum(lineitem.extendedprice)"),
Column.caluclatedColumn("extended_price_2", WrenTypes.BIGINT, "sum(lineitem.extendedprice + totalprice)"));
Column.calculatedColumn("customer_name", WrenTypes.BIGINT, "customer.name"),
Column.calculatedColumn("extended_price", WrenTypes.BIGINT, "sum(lineitem.extendedprice)"),
Column.calculatedColumn("extended_price_2", WrenTypes.BIGINT, "sum(lineitem.extendedprice + totalprice)"));
Model newLineitem = addColumnsToModel(
lineitem,
Column.column("orders", "Orders", "OrdersLineitem", true),
Column.caluclatedColumn("test_column", WrenTypes.BIGINT, "orders.customer.total_price + extendedprice"));
Column.calculatedColumn("test_column", WrenTypes.BIGINT, "orders.customer.total_price + extendedprice"));
Manifest manifest = withDefaultCatalogSchema()
.setModels(List.of(newCustomer, newOrders, newLineitem))
.setRelationships(List.of(ordersCustomer, ordersLineitem))
Expand Down Expand Up @@ -203,7 +203,7 @@ public void testAnalyzeModelOnModel()
Model newCustomer = addColumnsToModel(
customer,
Column.column("orders", "Orders", "OrdersCustomer", true),
Column.caluclatedColumn("total_price", WrenTypes.BIGINT, "sum(orders.totalprice)"));
Column.calculatedColumn("total_price", WrenTypes.BIGINT, "sum(orders.totalprice)"));
Model onCustomer = Model.onBaseObject(
"OnCustomer",
"Customer",
Expand All @@ -215,7 +215,7 @@ public void testAnalyzeModelOnModel()
Model newOrders = addColumnsToModel(
orders,
Column.column("on_customer", "OnCustomer", "OrdersOnCustomer", true),
Column.caluclatedColumn("customer_name", WrenTypes.BIGINT, "on_customer.mom_name"));
Column.calculatedColumn("customer_name", WrenTypes.BIGINT, "on_customer.mom_name"));
Relationship ordersOnCustomer = Relationship.relationship("OrdersOnCustomer", List.of("Orders", "OnCustomer"), JoinType.MANY_TO_ONE, "Orders.custkey = OnCustomer.mom_custkey");
Manifest manifest = withDefaultCatalogSchema()
.setModels(List.of(newOrders, newCustomer, onCustomer))
Expand Down Expand Up @@ -516,7 +516,7 @@ public void testGetSourceColumns()
Model newCustomer = addColumnsToModel(
customer,
Column.column("orders", "Orders", "OrdersCustomer", true),
Column.caluclatedColumn("discount_extended_price", WrenTypes.BIGINT, "sum(orders.lineitem.discount + orders.lineitem.extendedprice)"));
Column.calculatedColumn("discount_extended_price", WrenTypes.BIGINT, "sum(orders.lineitem.discount + orders.lineitem.extendedprice)"));
Manifest manifest = withDefaultCatalogSchema()
.setModels(List.of(newCustomer, orders, lineitem))
.setRelationships(List.of(ordersCustomer, ordersLineitem))
Expand Down
2 changes: 1 addition & 1 deletion wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestDecisionPointAnalyzer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ public TestDecisionPointAnalyzer()
Column.column("shippriority", WrenTypes.INTEGER, null, true),
Column.column("comment", WrenTypes.VARCHAR, null, true),
Column.column("customer", "customer", "CustomerOrders", false),
Column.caluclatedColumn("customer_name", WrenTypes.VARCHAR, "customer.name"));
Column.calculatedColumn("customer_name", WrenTypes.VARCHAR, "customer.name"));
List<Column> lineitemColumns = List.of(
Column.column("orderkey", WrenTypes.INTEGER, null, true),
Column.column("partkey", WrenTypes.INTEGER, null, true),
Expand Down
Loading

0 comments on commit e23469d

Please sign in to comment.