Skip to content

fix(core): fix the behaviors of dynamic query #776

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

Merged
merged 5 commits into from
Sep 4, 2024
Merged

fix(core): fix the behaviors of dynamic query #776

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ef7a539
fix count(*) with calculated field filter
goldmedal Sep 2, 2024
719ef9a
fix count(*) from the model or join
goldmedal Sep 2, 2024
52ad361
fix selecting wildcard exclude calculated field
goldmedal Sep 2, 2024
53ff4de
fix the style
goldmedal Sep 2, 2024
b34a5af
fix ibis clickhouse test
goldmedal Sep 2, 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
6 changes: 2 additions & 4 deletions ibis-server/tests/routers/v2/connector/test_clickhouse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ def test_query(clickhouse: ClickHouseContainer):
)
assert response.status_code == 200
result = response.json()
assert len(result["columns"]) == 10
assert len(result["columns"]) == 9
assert len(result["data"]) == 1
assert result["data"][0] == [
1,
Expand All @@ -179,7 +179,6 @@ def test_query(clickhouse: ClickHouseContainer):
"2024-01-01 23:59:59.000000",
"2024-01-01 23:59:59.000000 UTC",
None,
"Customer#000000370",
]
assert result["dtypes"] == {
"orderkey": "int32",
Expand All @@ -191,7 +190,6 @@ def test_query(clickhouse: ClickHouseContainer):
"timestamp": "object",
"timestamptz": "object",
"test_null_time": "object",
"customer_name": "object",
}


Expand All @@ -207,7 +205,7 @@ def test_query_with_connection_url(clickhouse: ClickHouseContainer):
)
assert response.status_code == 200
result = response.json()
assert len(result["columns"]) == 10
assert len(result["columns"]) == 9
assert len(result["data"]) == 1
assert result["data"][0][0] == 1
assert result["dtypes"] is not None
Expand Down
4 changes: 2 additions & 2 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 @@ -103,7 +103,7 @@ def test_query(postgres: PostgresContainer):
)
assert response.status_code == 200
result = response.json()
assert len(result["columns"]) == len(manifest["models"][0]["columns"])
assert len(result["columns"]) == 9
assert len(result["data"]) == 1
assert result["data"][0] == [
1,
Expand Down Expand Up @@ -141,7 +141,7 @@ def test_query_with_connection_url(postgres: PostgresContainer):
)
assert response.status_code == 200
result = response.json()
assert len(result["columns"]) == len(manifest["models"][0]["columns"])
assert len(result["columns"]) == 9
assert len(result["data"]) == 1
assert result["data"][0][0] == 1
assert result["dtypes"] is not None
Expand Down
2 changes: 1 addition & 1 deletion wren-base/src/main/java/io/wren/base/sqlrewrite/WrenSqlRewrite.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ public Statement apply(Statement root, SessionContext sessionContext, Analysis a
if (!tableRequiredFields.containsKey(tableName)) {
Relationable relationable = wrenMDL.getRelationable(tableName)
.orElseThrow(() -> new IllegalArgumentException(format("dataset not found: %s", tableName)));
tableRequiredFields.put(tableName, relationable.getColumns().stream().map(Column::getName).collect(toImmutableSet()));
tableRequiredFields.put(tableName, relationable.getColumns().stream().filter(column -> !column.isCalculated()).map(Column::getName).collect(toImmutableSet()));
}
});

Expand Down
24 changes: 21 additions & 3 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 @@ -19,6 +19,7 @@
import io.trino.sql.tree.AliasedRelation;
import io.trino.sql.tree.AllColumns;
import io.trino.sql.tree.AstVisitor;
import io.trino.sql.tree.DefaultTraversalVisitor;
import io.trino.sql.tree.DereferenceExpression;
import io.trino.sql.tree.Expression;
import io.trino.sql.tree.FrameBound;
Expand Down Expand Up @@ -372,7 +373,12 @@ private void analyzeSelectAllColumns(AllColumns allColumns, Scope scope, Immutab
// TODO handle target.*
}
else {
analysis.addCollectedColumns(scope.getRelationType().getFields());
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()))))
Expand All @@ -387,8 +393,20 @@ private void analyzeSelectSingleColumn(SingleColumn singleColumn, Scope scope, I
ExpressionAnalysis expressionAnalysis = analyzeExpression(singleColumn.getExpression(), scope);

if (expressionAnalysis.isRequireRelation()) {
analysis.addRequiredSourceNode(scope.getRelationId().getSourceNode()
.orElseThrow(() -> new IllegalArgumentException("count(*) should have a followed source")));
Node source = scope.getRelationId().getSourceNode()
.orElseThrow(() -> new IllegalArgumentException("count(*) should have a followed source"));

// collect only the source node that is a table for generating the required column for models
DefaultTraversalVisitor<Void> visitor = new DefaultTraversalVisitor<>()
{
@Override
protected Void visitTable(Table node, Void scope)
{
analysis.addRequiredSourceNode(node);
return null;
}
};
visitor.process(source, null);
}
}

Expand Down
38 changes: 35 additions & 3 deletions wren-tests/src/test/java/io/wren/testing/duckdb/TestWrenWithDuckDB.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,6 @@ public void testQueryOnlyModelColumn()
column("custkey", "integer"),
column("orderstatus", "varchar"),
column("totalprice", "DECIMAL(15,2)"),
column("nation_name", "varchar"),
column("orderdate", "date")));
assertThat(queryResultDto.getData().size()).isEqualTo(100);
}
Expand Down Expand Up @@ -139,7 +138,6 @@ public void testView()
column("custkey", "integer"),
column("orderstatus", "varchar"),
column("totalprice", "DECIMAL(15,2)"),
column("nation_name", "varchar"),
column("orderdate", "date")));
assertThat(queryResultDto.getData().size()).isEqualTo(100);

Expand All @@ -161,7 +159,6 @@ public void testView()
column("custkey", "integer"),
column("orderstatus", "varchar"),
column("totalprice", "DECIMAL(15,2)"),
column("nation_name", "varchar"),
column("orderdate", "date")));
assertThat(queryResultDto.getData().size()).isEqualTo(100);
}
Expand Down Expand Up @@ -198,4 +195,39 @@ public void testQueryJoinAliased()
.isEqualTo(List.of(column("totalprice", "DECIMAL(15,2)")));
assertThat(queryResultDto.getData().size()).isEqualTo(100);
}

@Test
public void testCountWithCalculatedFieldFilter()
{
QueryResultDto queryResultDto = query(manifest, "select count(*) from \"Orders\" where nation_name = 'ALGERIA'");
assertThat(queryResultDto.getData().getFirst()[0]).isEqualTo(691);
}

@Test
public void testCountJoin()
{
QueryResultDto queryResultDto = query(manifest, "select count(*) from Orders a");
assertThat(queryResultDto.getData().getFirst()[0]).isEqualTo(15000);

queryResultDto = query(manifest, "select count(*) from Orders, Customer");
assertThat(queryResultDto.getData().getFirst()[0]).isEqualTo(22500000);

queryResultDto = query(manifest, "select count(*) from Orders a, Customer b");
assertThat(queryResultDto.getData().getFirst()[0]).isEqualTo(22500000);

queryResultDto = query(manifest, "select count(*) from Orders a JOIN Customer b ON a.custkey = b.custkey");
assertThat(queryResultDto.getData().getFirst()[0]).isEqualTo(15000);
}

@Test
public void testSelectAllExcludeCalculatedField()
{
QueryResultDto queryResultDto = query(manifest, "select * from Orders limit 100");
assertThat(queryResultDto.getColumns())
.isEqualTo(List.of(column("orderkey", "INTEGER"),
column("custkey", "INTEGER"),
column("orderstatus", "VARCHAR"),
column("totalprice", "DECIMAL(15,2)"),
column("orderdate", "DATE")));
}
}
Loading