Skip to content

Commit a5daead

Browse files
committed
Copy from largelist
1 parent 063c596 commit a5daead

File tree

2 files changed

+60
-4
lines changed

2 files changed

+60
-4
lines changed

src/arrow_parquet/schema_parser.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -497,7 +497,9 @@ fn is_coercible(
497497

498498
true
499499
}
500-
(DataType::List(from_field), DataType::List(to_field)) => {
500+
(DataType::List(from_field), DataType::List(to_field))
501+
| (DataType::FixedSizeList(from_field, _), DataType::List(to_field))
502+
| (DataType::LargeList(from_field), DataType::List(to_field)) => {
501503
let element_oid = array_element_typoid(to_typoid);
502504
let element_typmod = to_typmod;
503505

src/pgrx_tests/copy_from_coerce.rs

Lines changed: 57 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@ mod tests {
1212
};
1313
use arrow::array::{
1414
ArrayRef, BinaryArray, BooleanArray, Date32Array, Decimal128Array, Float32Array,
15-
Float64Array, Int16Array, Int32Array, Int8Array, LargeBinaryArray, LargeStringArray,
16-
ListArray, MapArray, RecordBatch, StringArray, StructArray, Time64MicrosecondArray,
17-
TimestampMicrosecondArray, UInt16Array, UInt32Array, UInt64Array,
15+
Float64Array, Int16Array, Int32Array, Int8Array, LargeBinaryArray, LargeListArray,
16+
LargeStringArray, ListArray, MapArray, RecordBatch, StringArray, StructArray,
17+
Time64MicrosecondArray, TimestampMicrosecondArray, UInt16Array, UInt32Array, UInt64Array,
1818
};
1919
use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
2020
use arrow::datatypes::UInt16Type;
@@ -553,6 +553,60 @@ mod tests {
553553
Spi::run(drop_table).unwrap();
554554
}
555555

556+
#[pg_test]
557+
fn test_coerce_large_list() {
558+
// [UINT16] => {int[], bigint[]}
559+
let x_nullable = false;
560+
let field_x = Field::new(
561+
"x",
562+
DataType::LargeList(Field::new("item", DataType::UInt16, false).into()),
563+
x_nullable,
564+
);
565+
566+
let x = Arc::new(UInt16Array::from(vec![1, 2]));
567+
let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 2]));
568+
let x = Arc::new(LargeListArray::new(
569+
Arc::new(Field::new("item", DataType::UInt16, false)),
570+
offsets,
571+
x,
572+
None,
573+
));
574+
575+
let y_nullable = true;
576+
let field_y = Field::new(
577+
"y",
578+
DataType::LargeList(Field::new("item", DataType::UInt16, true).into()),
579+
y_nullable,
580+
);
581+
582+
let y = Arc::new(LargeListArray::from_iter_primitive::<UInt16Type, _, _>(
583+
vec![Some(vec![Some(3), Some(4)])],
584+
));
585+
586+
let schema = Arc::new(Schema::new(vec![field_x, field_y]));
587+
588+
let batch = RecordBatch::try_new(schema.clone(), vec![x, y]).unwrap();
589+
write_record_batch_to_parquet(schema, batch);
590+
591+
let create_table = "CREATE TABLE test_table (x int[], y bigint[])";
592+
Spi::run(create_table).unwrap();
593+
594+
let copy_from = format!("COPY test_table FROM '{}'", LOCAL_TEST_FILE_PATH);
595+
Spi::run(&copy_from).unwrap();
596+
597+
let value = Spi::get_two::<Vec<Option<i32>>, Vec<Option<i64>>>(
598+
"SELECT x, y FROM test_table LIMIT 1",
599+
)
600+
.unwrap();
601+
assert_eq!(
602+
value,
603+
(Some(vec![Some(1), Some(2)]), Some(vec![Some(3), Some(4)]))
604+
);
605+
606+
let drop_table = "DROP TABLE test_table";
607+
Spi::run(drop_table).unwrap();
608+
}
609+
556610
#[pg_test]
557611
fn test_coerce_struct_types() {
558612
// STRUCT {a: UINT16, b: UINT16} => test_type {a: int, b: bigint}

0 commit comments

Comments
 (0)