Skip to content
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

feat(engine): handle target.* #933

Merged
merged 7 commits into from
Nov 21, 2024
Merged

feat(engine): handle target.* #933

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
88effd2
feat(engine): handle target.*
grieve54706 Nov 20, 2024
e61db70
chore: merge two block
grieve54706 Nov 20, 2024
8c1bd3b
test: make it complicated and with calculated column
grieve54706 Nov 21, 2024
ac0e7b7
chore: fix typo
grieve54706 Nov 21, 2024
9085a9a
test: add more e2e test
grieve54706 Nov 21, 2024
031bdf7
chore: fix style
grieve54706 Nov 21, 2024
5a26db8
chore: comment
grieve54706 Nov 21, 2024
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
17 changes: 17 additions & 0 deletions ibis-server/tests/routers/v2/connector/test_postgres.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,23 @@ 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)
response = client.post(
url=f"{base_url}/query",
params={"limit": 1},
json={
"connectionInfo": connection_info,
"manifestStr": manifest_str,
"sql": 'SELECT "o".* FROM "Orders" AS "o"',
},
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you add the end-to-end SQL test for the calculated field and join syntax? Maybe we should add tested in TestDynamicFields, too. We need to make sure the model-generating behavior is fine in WrenSqlRewrite.

assert response.status_code == 200
result = response.json()
assert len(result["columns"]) == len(manifest["models"][0]["columns"])
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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