Skip to content

Commit ea8e467

Browse files
hareshkhhareshkh
committed
Adds a substrait test
1 parent a0820e8 commit ea8e467

2 files changed

Lines changed: 137 additions & 0 deletions

File tree

‎datafusion/substrait/tests/cases/consumer_integration.rs‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -762,4 +762,39 @@ mod tests {
762762

763763
Ok(())
764764
}
765+
766+
/// Regression test: a Substrait join expression containing both `equal` and
767+
/// `is_not_distinct_from` (as Spark can produce) must preserve the
768+
/// null-safe semantics of `IS NOT DISTINCT FROM` by demoting it to the
769+
/// join filter when mixed with regular equality keys.
770+
///
771+
/// The plan is loaded from a JSON-encoded Substrait protobuf to exercise
772+
/// the full consumer path (`from_substrait_plan` → `from_join_rel`).
773+
#[tokio::test]
774+
async fn test_mixed_join_equal_and_indistinct_from_substrait_plan() -> Result<()> {
775+
let path = "tests/testdata/test_plans/mixed_join_equal_and_indistinct.json";
776+
let proto = serde_json::from_reader::<_, Plan>(BufReader::new(
777+
File::open(path).expect("file not found"),
778+
))
779+
.expect("failed to parse json");
780+
781+
let ctx = SessionContext::new();
782+
let plan = from_substrait_plan(&ctx.state(), &proto).await?;
783+
784+
// Execute and count rows.
785+
// Both tables have 6 identical rows; rows 3 and 4 have val=NULL.
786+
// With correct handling, IS NOT DISTINCT FROM is demoted to the join
787+
// filter, so NULL=NULL matches and all 6 rows appear in the output.
788+
let df = ctx.execute_logical_plan(plan).await?;
789+
let results = df.collect().await?;
790+
let total_rows: usize = results.iter().map(|b| b.num_rows()).sum();
791+
792+
assert_eq!(
793+
total_rows, 6,
794+
"Expected 6 rows (including NULL=NULL matches via IS NOT DISTINCT FROM), \
795+
got {total_rows}. Mixed equal/is_not_distinct_from lost null-safe semantics."
796+
);
797+
798+
Ok(())
799+
}
765800
}

‎datafusion/substrait/tests/testdata/test_plans/mixed_join_equal_and_indistinct.json‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
{
2+
"extensions": [
3+
{ "extensionFunction": { "functionAnchor": 0, "name": "is_not_distinct_from" } },
4+
{ "extensionFunction": { "functionAnchor": 1, "name": "equal" } },
5+
{ "extensionFunction": { "functionAnchor": 2, "name": "and" } }
6+
],
7+
"relations": [{
8+
"root": {
9+
"input": {
10+
"join": {
11+
"common": { "direct": {} },
12+
"left": {
13+
"read": {
14+
"common": { "direct": {} },
15+
"baseSchema": {
16+
"names": ["id", "val", "comment"],
17+
"struct": {
18+
"types": [
19+
{ "string": { "nullability": "NULLABILITY_REQUIRED" } },
20+
{ "string": { "nullability": "NULLABILITY_NULLABLE" } },
21+
{ "string": { "nullability": "NULLABILITY_REQUIRED" } }
22+
],
23+
"nullability": "NULLABILITY_REQUIRED"
24+
}
25+
},
26+
"virtualTable": {
27+
"values": [
28+
{ "fields": [{ "string": "1", "nullable": false }, { "string": "a", "nullable": true }, { "string": "c1", "nullable": false }] },
29+
{ "fields": [{ "string": "2", "nullable": false }, { "string": "b", "nullable": true }, { "string": "c2", "nullable": false }] },
30+
{ "fields": [{ "string": "3", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c3", "nullable": false }] },
31+
{ "fields": [{ "string": "4", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c4", "nullable": false }] },
32+
{ "fields": [{ "string": "5", "nullable": false }, { "string": "e", "nullable": true }, { "string": "c5", "nullable": false }] },
33+
{ "fields": [{ "string": "6", "nullable": false }, { "string": "f", "nullable": true }, { "string": "c6", "nullable": false }] }
34+
]
35+
}
36+
}
37+
},
38+
"right": {
39+
"read": {
40+
"common": { "direct": {} },
41+
"baseSchema": {
42+
"names": ["id", "val", "comment"],
43+
"struct": {
44+
"types": [
45+
{ "string": { "nullability": "NULLABILITY_REQUIRED" } },
46+
{ "string": { "nullability": "NULLABILITY_NULLABLE" } },
47+
{ "string": { "nullability": "NULLABILITY_REQUIRED" } }
48+
],
49+
"nullability": "NULLABILITY_REQUIRED"
50+
}
51+
},
52+
"virtualTable": {
53+
"values": [
54+
{ "fields": [{ "string": "1", "nullable": false }, { "string": "a", "nullable": true }, { "string": "c1", "nullable": false }] },
55+
{ "fields": [{ "string": "2", "nullable": false }, { "string": "b", "nullable": true }, { "string": "c2", "nullable": false }] },
56+
{ "fields": [{ "string": "3", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c3", "nullable": false }] },
57+
{ "fields": [{ "string": "4", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c4", "nullable": false }] },
58+
{ "fields": [{ "string": "5", "nullable": false }, { "string": "e", "nullable": true }, { "string": "c5", "nullable": false }] },
59+
{ "fields": [{ "string": "6", "nullable": false }, { "string": "f", "nullable": true }, { "string": "c6", "nullable": false }] }
60+
]
61+
}
62+
}
63+
},
64+
"expression": {
65+
"scalarFunction": {
66+
"functionReference": 2,
67+
"outputType": { "bool": { "nullability": "NULLABILITY_NULLABLE" } },
68+
"arguments": [
69+
{
70+
"value": {
71+
"scalarFunction": {
72+
"functionReference": 0,
73+
"outputType": { "bool": { "nullability": "NULLABILITY_NULLABLE" } },
74+
"arguments": [
75+
{ "value": { "selection": { "directReference": { "structField": { "field": 1 } }, "rootReference": {} } } },
76+
{ "value": { "selection": { "directReference": { "structField": { "field": 4 } }, "rootReference": {} } } }
77+
]
78+
}
79+
}
80+
},
81+
{
82+
"value": {
83+
"scalarFunction": {
84+
"functionReference": 1,
85+
"outputType": { "bool": { "nullability": "NULLABILITY_NULLABLE" } },
86+
"arguments": [
87+
{ "value": { "selection": { "directReference": { "structField": { "field": 0 } }, "rootReference": {} } } },
88+
{ "value": { "selection": { "directReference": { "structField": { "field": 3 } }, "rootReference": {} } } }
89+
]
90+
}
91+
}
92+
}
93+
]
94+
}
95+
},
96+
"type": "JOIN_TYPE_INNER"
97+
}
98+
},
99+
"names": ["id", "val", "comment", "id0", "val0", "comment0"]
100+
}
101+
}]
102+
}

0 commit comments

Comments
 (0)