Skip to content

Commit 2574b16

Browse files
TangruilinTangruilin
committed
[task #8213]Port tests in select.rs to sqllogictest
Signed-off-by: tangruilin <[email protected]>
1 parent 827668a commit 2574b16

File tree

2 files changed

+179
-373
lines changed

2 files changed

+179
-373
lines changed

‎datafusion/core/tests/sql/select.rs

Lines changed: 0 additions & 373 deletions
Original file line numberDiff line numberDiff line change
@@ -19,379 +19,6 @@ use super::*;
1919
use datafusion_common::ScalarValue;
2020
use tempfile::TempDir;
2121

22-
#[tokio::test]
23-
async fn query_get_indexed_field() -> Result<()> {
24-
let ctx = SessionContext::new();
25-
let schema = Arc::new(Schema::new(vec![Field::new_list(
26-
"some_list",
27-
Field::new("item", DataType::Int64, true),
28-
false,
29-
)]));
30-
let builder = PrimitiveBuilder::<Int64Type>::with_capacity(3);
31-
let mut lb = ListBuilder::new(builder);
32-
for int_vec in [[0, 1, 2], [4, 5, 6], [7, 8, 9]] {
33-
let builder = lb.values();
34-
for int in int_vec {
35-
builder.append_value(int);
36-
}
37-
lb.append(true);
38-
}
39-
40-
let data = RecordBatch::try_new(schema.clone(), vec![Arc::new(lb.finish())])?;
41-
42-
ctx.register_batch("ints", data)?;
43-
44-
// Original column is micros, convert to millis and check timestamp
45-
let sql = "SELECT some_list[1] as i0 FROM ints LIMIT 3";
46-
let actual = execute_to_batches(&ctx, sql).await;
47-
#[rustfmt::skip]
48-
let expected = ["+----+",
49-
"| i0 |",
50-
"+----+",
51-
"| 0 |",
52-
"| 4 |",
53-
"| 7 |",
54-
"+----+"];
55-
assert_batches_eq!(expected, &actual);
56-
Ok(())
57-
}
58-
59-
#[tokio::test]
60-
async fn query_nested_get_indexed_field() -> Result<()> {
61-
let ctx = SessionContext::new();
62-
let nested_dt = DataType::List(Arc::new(Field::new("item", DataType::Int64, true)));
63-
// Nested schema of { "some_list": [[i64]] }
64-
let schema = Arc::new(Schema::new(vec![Field::new(
65-
"some_list",
66-
DataType::List(Arc::new(Field::new("item", nested_dt.clone(), true))),
67-
false,
68-
)]));
69-
70-
let builder = PrimitiveBuilder::<Int64Type>::with_capacity(3);
71-
let nested_lb = ListBuilder::new(builder);
72-
let mut lb = ListBuilder::new(nested_lb);
73-
for int_vec_vec in [
74-
[[0, 1], [2, 3], [3, 4]],
75-
[[5, 6], [7, 8], [9, 10]],
76-
[[11, 12], [13, 14], [15, 16]],
77-
] {
78-
let nested_builder = lb.values();
79-
for int_vec in int_vec_vec {
80-
let builder = nested_builder.values();
81-
for int in int_vec {
82-
builder.append_value(int);
83-
}
84-
nested_builder.append(true);
85-
}
86-
lb.append(true);
87-
}
88-
89-
let data = RecordBatch::try_new(schema.clone(), vec![Arc::new(lb.finish())])?;
90-
91-
ctx.register_batch("ints", data)?;
92-
93-
// Original column is micros, convert to millis and check timestamp
94-
let sql = "SELECT some_list[1] as i0 FROM ints LIMIT 3";
95-
let actual = execute_to_batches(&ctx, sql).await;
96-
let expected = [
97-
"+----------+",
98-
"| i0 |",
99-
"+----------+",
100-
"| [0, 1] |",
101-
"| [5, 6] |",
102-
"| [11, 12] |",
103-
"+----------+",
104-
];
105-
assert_batches_eq!(expected, &actual);
106-
let sql = "SELECT some_list[1][1] as i0 FROM ints LIMIT 3";
107-
let actual = execute_to_batches(&ctx, sql).await;
108-
#[rustfmt::skip]
109-
let expected = ["+----+",
110-
"| i0 |",
111-
"+----+",
112-
"| 0 |",
113-
"| 5 |",
114-
"| 11 |",
115-
"+----+"];
116-
assert_batches_eq!(expected, &actual);
117-
Ok(())
118-
}
119-
120-
#[tokio::test]
121-
async fn query_nested_get_indexed_field_on_struct() -> Result<()> {
122-
let ctx = SessionContext::new();
123-
let nested_dt = DataType::List(Arc::new(Field::new("item", DataType::Int64, true)));
124-
// Nested schema of { "some_struct": { "bar": [i64] } }
125-
let struct_fields = vec![Field::new("bar", nested_dt.clone(), true)];
126-
let schema = Arc::new(Schema::new(vec![Field::new(
127-
"some_struct",
128-
DataType::Struct(struct_fields.clone().into()),
129-
false,
130-
)]));
131-
132-
let builder = PrimitiveBuilder::<Int64Type>::with_capacity(3);
133-
let nested_lb = ListBuilder::new(builder);
134-
let mut sb = StructBuilder::new(struct_fields, vec![Box::new(nested_lb)]);
135-
for int_vec in [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]] {
136-
let lb = sb.field_builder::<ListBuilder<Int64Builder>>(0).unwrap();
137-
for int in int_vec {
138-
lb.values().append_value(int);
139-
}
140-
lb.append(true);
141-
sb.append(true);
142-
}
143-
let s = sb.finish();
144-
let data = RecordBatch::try_new(schema.clone(), vec![Arc::new(s)])?;
145-
146-
ctx.register_batch("structs", data)?;
147-
148-
// Original column is micros, convert to millis and check timestamp
149-
let sql = "SELECT some_struct['bar'] as l0 FROM structs LIMIT 3";
150-
let actual = execute_to_batches(&ctx, sql).await;
151-
let expected = [
152-
"+----------------+",
153-
"| l0 |",
154-
"+----------------+",
155-
"| [0, 1, 2, 3] |",
156-
"| [4, 5, 6, 7] |",
157-
"| [8, 9, 10, 11] |",
158-
"+----------------+",
159-
];
160-
assert_batches_eq!(expected, &actual);
161-
162-
// Access to field of struct by CompoundIdentifier
163-
let sql = "SELECT some_struct.bar as l0 FROM structs LIMIT 3";
164-
let actual = execute_to_batches(&ctx, sql).await;
165-
let expected = [
166-
"+----------------+",
167-
"| l0 |",
168-
"+----------------+",
169-
"| [0, 1, 2, 3] |",
170-
"| [4, 5, 6, 7] |",
171-
"| [8, 9, 10, 11] |",
172-
"+----------------+",
173-
];
174-
assert_batches_eq!(expected, &actual);
175-
176-
let sql = "SELECT some_struct['bar'][1] as i0 FROM structs LIMIT 3";
177-
let actual = execute_to_batches(&ctx, sql).await;
178-
#[rustfmt::skip]
179-
let expected = ["+----+",
180-
"| i0 |",
181-
"+----+",
182-
"| 0 |",
183-
"| 4 |",
184-
"| 8 |",
185-
"+----+"];
186-
assert_batches_eq!(expected, &actual);
187-
Ok(())
188-
}
189-
190-
#[tokio::test]
191-
async fn query_on_string_dictionary() -> Result<()> {
192-
// Test to ensure DataFusion can operate on dictionary types
193-
// Use StringDictionary (32 bit indexes = keys)
194-
let d1: DictionaryArray<Int32Type> =
195-
vec![Some("one"), None, Some("three")].into_iter().collect();
196-
197-
let d2: DictionaryArray<Int32Type> = vec![Some("blarg"), None, Some("three")]
198-
.into_iter()
199-
.collect();
200-
201-
let d3: StringArray = vec![Some("XYZ"), None, Some("three")].into_iter().collect();
202-
203-
let batch = RecordBatch::try_from_iter(vec![
204-
("d1", Arc::new(d1) as ArrayRef),
205-
("d2", Arc::new(d2) as ArrayRef),
206-
("d3", Arc::new(d3) as ArrayRef),
207-
])
208-
.unwrap();
209-
210-
let ctx = SessionContext::new();
211-
ctx.register_batch("test", batch)?;
212-
213-
// Basic SELECT
214-
let sql = "SELECT d1 FROM test";
215-
let actual = execute_to_batches(&ctx, sql).await;
216-
let expected = [
217-
"+-------+",
218-
"| d1 |",
219-
"+-------+",
220-
"| one |",
221-
"| |",
222-
"| three |",
223-
"+-------+",
224-
];
225-
assert_batches_eq!(expected, &actual);
226-
227-
// basic filtering
228-
let sql = "SELECT d1 FROM test WHERE d1 IS NOT NULL";
229-
let actual = execute_to_batches(&ctx, sql).await;
230-
let expected = [
231-
"+-------+",
232-
"| d1 |",
233-
"+-------+",
234-
"| one |",
235-
"| three |",
236-
"+-------+",
237-
];
238-
assert_batches_eq!(expected, &actual);
239-
240-
// comparison with constant
241-
let sql = "SELECT d1 FROM test WHERE d1 = 'three'";
242-
let actual = execute_to_batches(&ctx, sql).await;
243-
let expected = [
244-
"+-------+",
245-
"| d1 |",
246-
"+-------+",
247-
"| three |",
248-
"+-------+",
249-
];
250-
assert_batches_eq!(expected, &actual);
251-
252-
// comparison with another dictionary column
253-
let sql = "SELECT d1 FROM test WHERE d1 = d2";
254-
let actual = execute_to_batches(&ctx, sql).await;
255-
let expected = [
256-
"+-------+",
257-
"| d1 |",
258-
"+-------+",
259-
"| three |",
260-
"+-------+",
261-
];
262-
assert_batches_eq!(expected, &actual);
263-
264-
// order comparison with another dictionary column
265-
let sql = "SELECT d1 FROM test WHERE d1 <= d2";
266-
let actual = execute_to_batches(&ctx, sql).await;
267-
let expected = [
268-
"+-------+",
269-
"| d1 |",
270-
"+-------+",
271-
"| three |",
272-
"+-------+",
273-
];
274-
assert_batches_eq!(expected, &actual);
275-
276-
// comparison with a non dictionary column
277-
let sql = "SELECT d1 FROM test WHERE d1 = d3";
278-
let actual = execute_to_batches(&ctx, sql).await;
279-
let expected = [
280-
"+-------+",
281-
"| d1 |",
282-
"+-------+",
283-
"| three |",
284-
"+-------+",
285-
];
286-
assert_batches_eq!(expected, &actual);
287-
288-
// filtering with constant
289-
let sql = "SELECT d1 FROM test WHERE d1 = 'three'";
290-
let actual = execute_to_batches(&ctx, sql).await;
291-
let expected = [
292-
"+-------+",
293-
"| d1 |",
294-
"+-------+",
295-
"| three |",
296-
"+-------+",
297-
];
298-
assert_batches_eq!(expected, &actual);
299-
300-
// Expression evaluation
301-
let sql = "SELECT concat(d1, '-foo') FROM test";
302-
let actual = execute_to_batches(&ctx, sql).await;
303-
let expected = [
304-
"+------------------------------+",
305-
"| concat(test.d1,Utf8(\"-foo\")) |",
306-
"+------------------------------+",
307-
"| one-foo |",
308-
"| -foo |",
309-
"| three-foo |",
310-
"+------------------------------+",
311-
];
312-
assert_batches_eq!(expected, &actual);
313-
314-
// Expression evaluation with two dictionaries
315-
let sql = "SELECT concat(d1, d2) FROM test";
316-
let actual = execute_to_batches(&ctx, sql).await;
317-
let expected = [
318-
"+-------------------------+",
319-
"| concat(test.d1,test.d2) |",
320-
"+-------------------------+",
321-
"| oneblarg |",
322-
"| |",
323-
"| threethree |",
324-
"+-------------------------+",
325-
];
326-
assert_batches_eq!(expected, &actual);
327-
328-
// aggregation
329-
let sql = "SELECT COUNT(d1) FROM test";
330-
let actual = execute_to_batches(&ctx, sql).await;
331-
let expected = [
332-
"+----------------+",
333-
"| COUNT(test.d1) |",
334-
"+----------------+",
335-
"| 2 |",
336-
"+----------------+",
337-
];
338-
assert_batches_eq!(expected, &actual);
339-
340-
// aggregation min
341-
let sql = "SELECT MIN(d1) FROM test";
342-
let actual = execute_to_batches(&ctx, sql).await;
343-
let expected = [
344-
"+--------------+",
345-
"| MIN(test.d1) |",
346-
"+--------------+",
347-
"| one |",
348-
"+--------------+",
349-
];
350-
assert_batches_eq!(expected, &actual);
351-
352-
// aggregation max
353-
let sql = "SELECT MAX(d1) FROM test";
354-
let actual = execute_to_batches(&ctx, sql).await;
355-
let expected = [
356-
"+--------------+",
357-
"| MAX(test.d1) |",
358-
"+--------------+",
359-
"| three |",
360-
"+--------------+",
361-
];
362-
assert_batches_eq!(expected, &actual);
363-
364-
// grouping
365-
let sql = "SELECT d1, COUNT(*) FROM test group by d1";
366-
let actual = execute_to_batches(&ctx, sql).await;
367-
let expected = [
368-
"+-------+----------+",
369-
"| d1 | COUNT(*) |",
370-
"+-------+----------+",
371-
"| | 1 |",
372-
"| one | 1 |",
373-
"| three | 1 |",
374-
"+-------+----------+",
375-
];
376-
assert_batches_sorted_eq!(expected, &actual);
377-
378-
// window functions
379-
let sql = "SELECT d1, row_number() OVER (partition by d1) as rn1 FROM test";
380-
let actual = execute_to_batches(&ctx, sql).await;
381-
let expected = [
382-
"+-------+-----+",
383-
"| d1 | rn1 |",
384-
"+-------+-----+",
385-
"| | 1 |",
386-
"| one | 1 |",
387-
"| three | 1 |",
388-
"+-------+-----+",
389-
];
390-
assert_batches_sorted_eq!(expected, &actual);
391-
392-
Ok(())
393-
}
394-
39522
#[tokio::test]
39623
async fn sort_on_window_null_string() -> Result<()> {
39724
let d1: DictionaryArray<Int32Type> =

0 commit comments

Comments
 (0)