Skip to content

Commit d2b57e5

Browse files
coderfendercoderfender
authored
Merge branch 'main' into bitmap_instead_hll_smaller_int_types
2 parents c16592d + 1fbbba5 commit d2b57e5

23 files changed

Lines changed: 665 additions & 225 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎datafusion-cli/src/print_format.rs‎

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ fn keep_only_maxrows(s: &str, maxrows: usize) -> String {
9797
let last_line = &lines[lines.len() - 1]; // bottom border line
9898

9999
let spaces = last_line.len().saturating_sub(4);
100-
let dotted_line = format!("| .{:<spaces$}|", "", spaces = spaces);
100+
let dotted_line = format!("| .{}|", " ".repeat(spaces));
101101

102102
let mut result = lines[0..(maxrows + 3)].to_vec(); // Keep top border and `maxrows` lines
103103
result.extend(vec![dotted_line; 3]); // Append ... lines
@@ -632,6 +632,41 @@ mod tests {
632632
.unwrap()
633633
}
634634

635+
#[test]
636+
fn print_maxrows_limited_wide_table() {
637+
let output = PrintBatchesTest::new()
638+
.with_format(PrintFormat::Table)
639+
.with_batches(vec![wide_column_batch()])
640+
.with_maxrows(MaxRows::Limited(1))
641+
.run();
642+
assert_snapshot!(output, @r"
643+
+----+----+----+----+----+----+----+----+----+----+
644+
| c0 | c1 | c2 | c3 | c4 | c5 | c6 | c7 | c8 | c9 |
645+
+----+----+----+----+----+----+----+----+----+----+
646+
| 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 |
647+
| . |
648+
| . |
649+
| . |
650+
+----+----+----+----+----+----+----+----+----+----+
651+
");
652+
}
653+
654+
/// return a schema with many columns (to exercise wide table formatting)
655+
fn wide_column_schema() -> SchemaRef {
656+
let fields: Vec<Field> = (0..10)
657+
.map(|i| Field::new(format!("c{i}"), DataType::Int32, false))
658+
.collect();
659+
Arc::new(Schema::new(fields))
660+
}
661+
662+
/// return a batch with many columns and three rows
663+
fn wide_column_batch() -> RecordBatch {
664+
let arrays: Vec<Arc<dyn arrow::array::Array>> = (0..10)
665+
.map(|_| Arc::new(Int32Array::from(vec![0, 1, 2])) as _)
666+
.collect();
667+
RecordBatch::try_new(wide_column_schema(), arrays).unwrap()
668+
}
669+
635670
/// Slice the record batch into 2 batches
636671
fn split_batch(batch: &RecordBatch) -> Vec<RecordBatch> {
637672
assert!(batch.num_rows() > 1);

‎datafusion-cli/tests/cli_integration.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -287,6 +287,23 @@ fn test_cli_with_unbounded_memory_pool() {
287287
assert_cmd_snapshot!(cmd);
288288
}
289289

290+
#[test]
291+
fn test_cli_wide_result_set_no_crash() {
292+
let mut settings = make_settings();
293+
294+
settings.set_snapshot_suffix("wide_result_set");
295+
296+
let _bound = settings.bind_to_scope();
297+
298+
let mut cmd = cli();
299+
let sql = "SELECT v1 as c0, v1+1 as c1, v1+2 as c2, v1+3 as c3, v1+4 as c4, \
300+
v1+5 as c5, v1+6 as c6, v1+7 as c7, v1+8 as c8, v1+9 as c9 \
301+
FROM generate_series(1, 100) as t1(v1);";
302+
cmd.args(["--maxrows", "5", "--command", sql]);
303+
304+
assert_cmd_snapshot!(cmd);
305+
}
306+
290307
#[tokio::test]
291308
async fn test_cli() {
292309
if env::var("TEST_STORAGE_INTEGRATION").is_err() {

‎datafusion-cli/tests/snapshots/cli_wide_result_set_no_crash@wide_result_set.snap‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
---
2+
source: datafusion-cli/tests/cli_integration.rs
3+
assertion_line: 307
4+
info:
5+
program: datafusion-cli
6+
args:
7+
- "--maxrows"
8+
- "5"
9+
- "--command"
10+
- "SELECT v1 as c0, v1+1 as c1, v1+2 as c2, v1+3 as c3, v1+4 as c4, v1+5 as c5, v1+6 as c6, v1+7 as c7, v1+8 as c8, v1+9 as c9 FROM generate_series(1, 100) as t1(v1);"
11+
---
12+
success: true
13+
exit_code: 0
14+
----- stdout -----
15+
[CLI_VERSION]
16+
+----+----+----+----+----+----+----+----+----+----+
17+
| c0 | c1 | c2 | c3 | c4 | c5 | c6 | c7 | c8 | c9 |
18+
+----+----+----+----+----+----+----+----+----+----+
19+
| 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 |
20+
| 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 |
21+
| 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 |
22+
| 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 |
23+
| 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 |
24+
| . |
25+
| . |
26+
| . |
27+
+----+----+----+----+----+----+----+----+----+----+
28+
100 row(s) fetched. (First 5 displayed. Use --maxrows to adjust)
29+
[ELAPSED]
30+
31+
32+
----- stderr -----

‎datafusion/core/tests/dataframe/dataframe_functions.rs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -402,7 +402,7 @@ async fn test_fn_approx_median() -> Result<()> {
402402
+-----------------------+
403403
| approx_median(test.b) |
404404
+-----------------------+
405-
| 10 |
405+
| 10.0 |
406406
+-----------------------+
407407
");
408408

@@ -422,7 +422,7 @@ async fn test_fn_approx_percentile_cont() -> Result<()> {
422422
+---------------------------------------------------------------------------+
423423
| approx_percentile_cont(Float64(0.5)) WITHIN GROUP [test.b ASC NULLS LAST] |
424424
+---------------------------------------------------------------------------+
425-
| 10 |
425+
| 10.0 |
426426
+---------------------------------------------------------------------------+
427427
");
428428

@@ -437,7 +437,7 @@ async fn test_fn_approx_percentile_cont() -> Result<()> {
437437
+----------------------------------------------------------------------------+
438438
| approx_percentile_cont(Float64(0.1)) WITHIN GROUP [test.b DESC NULLS LAST] |
439439
+----------------------------------------------------------------------------+
440-
| 100 |
440+
| 100.0 |
441441
+----------------------------------------------------------------------------+
442442
");
443443

@@ -457,7 +457,7 @@ async fn test_fn_approx_percentile_cont() -> Result<()> {
457457
+--------------------------------------------------------------------+
458458
| approx_percentile_cont(arg_2) WITHIN GROUP [test.b ASC NULLS LAST] |
459459
+--------------------------------------------------------------------+
460-
| 10 |
460+
| 10.0 |
461461
+--------------------------------------------------------------------+
462462
"
463463
);
@@ -477,7 +477,7 @@ async fn test_fn_approx_percentile_cont() -> Result<()> {
477477
+---------------------------------------------------------------------+
478478
| approx_percentile_cont(arg_2) WITHIN GROUP [test.b DESC NULLS LAST] |
479479
+---------------------------------------------------------------------+
480-
| 100 |
480+
| 100.0 |
481481
+---------------------------------------------------------------------+
482482
"
483483
);
@@ -494,7 +494,7 @@ async fn test_fn_approx_percentile_cont() -> Result<()> {
494494
+------------------------------------------------------------------------------------+
495495
| approx_percentile_cont(Float64(0.5),Int32(2)) WITHIN GROUP [test.b ASC NULLS LAST] |
496496
+------------------------------------------------------------------------------------+
497-
| 30 |
497+
| 30.25 |
498498
+------------------------------------------------------------------------------------+
499499
");
500500

@@ -510,7 +510,7 @@ async fn test_fn_approx_percentile_cont() -> Result<()> {
510510
+-------------------------------------------------------------------------------------+
511511
| approx_percentile_cont(Float64(0.1),Int32(2)) WITHIN GROUP [test.b DESC NULLS LAST] |
512512
+-------------------------------------------------------------------------------------+
513-
| 69 |
513+
| 69.85 |
514514
+-------------------------------------------------------------------------------------+
515515
");
516516

‎datafusion/core/tests/dataframe/mod.rs‎

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1204,26 +1204,26 @@ async fn window_using_aggregates() -> Result<()> {
12041204
| first_value | last_val | approx_distinct | approx_median | median | max | min | c2 | c3 |
12051205
+-------------+----------+-----------------+---------------+--------+-----+------+----+------+
12061206
| | | | | | | | 1 | -85 |
1207-
| -85 | -101 | 14 | -12 | -12 | 83 | -101 | 4 | -54 |
1208-
| -85 | -101 | 17 | -25 | -25 | 83 | -101 | 5 | -31 |
1209-
| -85 | -12 | 10 | -32 | -34 | 83 | -85 | 3 | 13 |
1210-
| -85 | -25 | 3 | -56 | -56 | -25 | -85 | 1 | -5 |
1211-
| -85 | -31 | 18 | -29 | -28 | 83 | -101 | 5 | 36 |
1212-
| -85 | -38 | 16 | -25 | -25 | 83 | -101 | 4 | 65 |
1213-
| -85 | -43 | 7 | -43 | -43 | 83 | -85 | 2 | 45 |
1214-
| -85 | -48 | 6 | -35 | -36 | 83 | -85 | 2 | -43 |
1215-
| -85 | -5 | 4 | -37 | -40 | -5 | -85 | 1 | 83 |
1216-
| -85 | -54 | 15 | -17 | -18 | 83 | -101 | 4 | -38 |
1217-
| -85 | -56 | 2 | -70 | -70 | -56 | -85 | 1 | -25 |
1218-
| -85 | -72 | 9 | -43 | -43 | 83 | -85 | 3 | -12 |
1219-
| -85 | -85 | 1 | -85 | -85 | -85 | -85 | 1 | -56 |
1220-
| -85 | 13 | 11 | -17 | -18 | 83 | -85 | 3 | 14 |
1221-
| -85 | 13 | 11 | -25 | -25 | 83 | -85 | 3 | 13 |
1222-
| -85 | 14 | 12 | -12 | -12 | 83 | -85 | 3 | 17 |
1223-
| -85 | 17 | 13 | -11 | -8 | 83 | -85 | 4 | -101 |
1224-
| -85 | 45 | 8 | -34 | -34 | 83 | -85 | 3 | -72 |
1225-
| -85 | 65 | 17 | -17 | -18 | 83 | -101 | 5 | -101 |
1226-
| -85 | 83 | 5 | -25 | -25 | 83 | -85 | 2 | -48 |
1207+
| -85 | -101 | 14 | -12.0 | -12 | 83 | -101 | 4 | -54 |
1208+
| -85 | -101 | 17 | -25.0 | -25 | 83 | -101 | 5 | -31 |
1209+
| -85 | -12 | 10 | -32.75 | -34 | 83 | -85 | 3 | 13 |
1210+
| -85 | -25 | 3 | -56.0 | -56 | -25 | -85 | 1 | -5 |
1211+
| -85 | -31 | 18 | -29.75 | -28 | 83 | -101 | 5 | 36 |
1212+
| -85 | -38 | 16 | -25.0 | -25 | 83 | -101 | 4 | 65 |
1213+
| -85 | -43 | 7 | -43.0 | -43 | 83 | -85 | 2 | 45 |
1214+
| -85 | -48 | 6 | -35.75 | -36 | 83 | -85 | 2 | -43 |
1215+
| -85 | -5 | 4 | -37.75 | -40 | -5 | -85 | 1 | 83 |
1216+
| -85 | -54 | 15 | -17.0 | -18 | 83 | -101 | 4 | -38 |
1217+
| -85 | -56 | 2 | -70.5 | -70 | -56 | -85 | 1 | -25 |
1218+
| -85 | -72 | 9 | -43.0 | -43 | 83 | -85 | 3 | -12 |
1219+
| -85 | -85 | 1 | -85.0 | -85 | -85 | -85 | 1 | -56 |
1220+
| -85 | 13 | 11 | -17.0 | -18 | 83 | -85 | 3 | 14 |
1221+
| -85 | 13 | 11 | -25.0 | -25 | 83 | -85 | 3 | 13 |
1222+
| -85 | 14 | 12 | -12.0 | -12 | 83 | -85 | 3 | 17 |
1223+
| -85 | 17 | 13 | -11.25 | -8 | 83 | -85 | 4 | -101 |
1224+
| -85 | 45 | 8 | -34.5 | -34 | 83 | -85 | 3 | -72 |
1225+
| -85 | 65 | 17 | -17.0 | -18 | 83 | -101 | 5 | -101 |
1226+
| -85 | 83 | 5 | -25.0 | -25 | 83 | -85 | 2 | -48 |
12271227
+-------------+----------+-----------------+---------------+--------+-----+------+----+------+
12281228
"
12291229
);

‎datafusion/expr-common/src/signature.rs‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ use std::fmt::Display;
2121
use std::hash::Hash;
2222
use std::sync::Arc;
2323

24-
use crate::type_coercion::aggregates::NUMERICS;
2524
use arrow::datatypes::{
2625
DECIMAL32_MAX_PRECISION, DECIMAL64_MAX_PRECISION, DECIMAL128_MAX_PRECISION, DataType,
2726
Decimal128Type, DecimalType, Field, IntervalUnit, TimeUnit,
@@ -596,6 +595,20 @@ impl Display for ArrayFunctionArgument {
596595
}
597596
}
598597

598+
static NUMERICS: &[DataType] = &[
599+
DataType::Int8,
600+
DataType::Int16,
601+
DataType::Int32,
602+
DataType::Int64,
603+
DataType::UInt8,
604+
DataType::UInt16,
605+
DataType::UInt32,
606+
DataType::UInt64,
607+
DataType::Float16,
608+
DataType::Float32,
609+
DataType::Float64,
610+
];
611+
599612
impl TypeSignature {
600613
pub fn to_string_repr(&self) -> Vec<String> {
601614
match self {

‎datafusion/expr-common/src/type_coercion/aggregates.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,7 @@ use arrow::datatypes::{DataType, FieldRef};
2020

2121
use datafusion_common::{Result, internal_err, plan_err};
2222

23-
// TODO: remove usage of these (INTEGERS and NUMERICS) in favour of signatures
24-
// see https://github.com/apache/datafusion/issues/18092
23+
#[deprecated(since = "54.0.0", note = "Use functions signatures")]
2524
pub static INTEGERS: &[DataType] = &[
2625
DataType::Int8,
2726
DataType::Int16,
@@ -33,6 +32,7 @@ pub static INTEGERS: &[DataType] = &[
3332
DataType::UInt64,
3433
];
3534

35+
#[deprecated(since = "54.0.0", note = "Use functions signatures")]
3636
pub static NUMERICS: &[DataType] = &[
3737
DataType::Int8,
3838
DataType::Int16,

‎datafusion/expr/Cargo.toml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,3 +66,6 @@ sqlparser = { workspace = true, optional = true }
6666
ctor = { workspace = true }
6767
env_logger = { workspace = true }
6868
insta = { workspace = true }
69+
# Makes sure `test_display_pg_json` behaves in a consistent way regardless of
70+
# feature unification with dependencies
71+
serde_json = { workspace = true, features = ["preserve_order"] }

‎datafusion/expr/src/test/function_stub.rs‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,14 @@ use datafusion_common::plan_err;
2929
use datafusion_common::{Result, exec_err, not_impl_err, utils::take_function_args};
3030

3131
use crate::Volatility::Immutable;
32-
use crate::type_coercion::aggregates::NUMERICS;
3332
use crate::{
34-
Accumulator, AggregateUDFImpl, Expr, GroupsAccumulator, ReversedUDAF, Signature,
33+
Accumulator, AggregateUDFImpl, Coercion, Expr, GroupsAccumulator, ReversedUDAF,
34+
Signature, TypeSignature, TypeSignatureClass,
3535
expr::AggregateFunction,
3636
function::{AccumulatorArgs, StateFieldsArgs},
3737
utils::AggregateOrderSensitivity,
3838
};
39+
use datafusion_common::types::{NativeType, logical_float64};
3940

4041
macro_rules! create_func {
4142
($UDAF:ty, $AGGREGATE_UDF_FN:ident) => {
@@ -444,9 +445,22 @@ pub struct Avg {
444445

445446
impl Avg {
446447
pub fn new() -> Self {
448+
let signature = Signature::one_of(
449+
vec![
450+
TypeSignature::Coercible(vec![Coercion::new_exact(
451+
TypeSignatureClass::Decimal,
452+
)]),
453+
TypeSignature::Coercible(vec![Coercion::new_implicit(
454+
TypeSignatureClass::Native(logical_float64()),
455+
vec![TypeSignatureClass::Integer, TypeSignatureClass::Float],
456+
NativeType::Float64,
457+
)]),
458+
],
459+
Immutable,
460+
);
447461
Self {
448462
aliases: vec![String::from("mean")],
449-
signature: Signature::uniform(1, NUMERICS.to_vec(), Immutable),
463+
signature,
450464
}
451465
}
452466
}

0 commit comments

Comments
 (0)