Skip to content

fix: add missing dict_id #1068

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 8 commits into from
Jul 14, 2023
Merged

fix: add missing dict_id #1068

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
58cdb1d
fix missing dict id
jiacai2050 Jul 12, 2023
1499224
typo
jiacai2050 Jul 12, 2023
2f977e4
update ut
jiacai2050 Jul 12, 2023
5cd1e7c
add more ut
jiacai2050 Jul 12, 2023
c37c61c
remove comments
jiacai2050 Jul 13, 2023
f936250
fix CI
jiacai2050 Jul 13, 2023
a6c6196
fix comments
jiacai2050 Jul 13, 2023
b4f3074
fix CR
jiacai2050 Jul 13, 2023
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
64 changes: 59 additions & 5 deletions components/arrow_ext/src/ipc.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,10 @@

//! Utilities for `RecordBatch` serialization using Arrow IPC

use std::io::Cursor;
use std::{borrow::Cow, io::Cursor, sync::Arc};

use arrow::{
datatypes::{DataType, Field, Schema, SchemaRef},
ipc::{reader::StreamReader, writer::StreamWriter},
record_batch::RecordBatch,
};
Expand Down Expand Up @@ -116,20 +117,73 @@ impl RecordBatchesEncoder {
self.num_rows
}

// Workaround for https://github.com/apache/arrow-datafusion/issues/6784
fn convert_schema(schema: &SchemaRef) -> Cow<SchemaRef> {
let dict_field_num: usize = schema
.fields()
.iter()
.map(|f| {
if let DataType::Dictionary(_, _) = f.data_type() {
1
} else {
0
}
})
.sum();
if dict_field_num <= 1 {
return Cow::Borrowed(schema);
}

let new_fields = schema
.fields()
.iter()
.enumerate()
.map(|(i, f)| {
if let DataType::Dictionary(_, _) = f.data_type() {
let dict_id = i as i64;
Arc::new(Field::new_dict(
f.name(),
f.data_type().clone(),
f.is_nullable(),
dict_id,
f.dict_is_ordered().unwrap_or(false),
))
} else {
f.clone()
}
})
.collect::<Vec<_>>();

let schema_ref = Arc::new(Schema::new_with_metadata(
new_fields,
schema.metadata.clone(),
));

Cow::Owned(schema_ref)
}

/// Append one batch into the encoder for encoding.
pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
let stream_writer = if let Some(v) = &mut self.stream_writer {
v
} else {
// TODO: pre-allocate the buffer.
let buffer: Vec<u8> = Vec::new();
let mem_size = batch
.columns()
.iter()
.map(|col| col.get_buffer_memory_size())
.sum();
let buffer: Vec<u8> = Vec::with_capacity(mem_size);
let stream_writer =
StreamWriter::try_new(buffer, &batch.schema()).context(ArrowError)?;
StreamWriter::try_new(buffer, &Self::convert_schema(&batch.schema()))
.context(ArrowError)?;
self.stream_writer = Some(stream_writer);
self.stream_writer.as_mut().unwrap()
};

stream_writer.write(batch).context(ArrowError)?;
let schema = batch.schema();
let schema = Self::convert_schema(&schema);
let batch = RecordBatch::try_new(schema.into_owned(), batch.columns().to_vec()).unwrap();
stream_writer.write(&batch).context(ArrowError)?;
self.num_rows += batch.num_rows();
Ok(())
}
Expand Down
33 changes: 30 additions & 3 deletions integration_tests/cases/env/local/ddl/alter_table.result
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,32 @@ UInt64(0),Timestamp(1),Int32(1),String("d1"),String(""),
UInt64(0),Timestamp(2),Int32(2),String("d2"),String("2"),


ALTER TABLE `05_alter_table_t0` add COLUMN (add_dic string dictionary);

affected_rows: 0

DESCRIBE TABLE `05_alter_table_t0`;

name,type,is_primary,is_nullable,is_tag,is_dictionary,
String("tsid"),String("uint64"),Boolean(true),Boolean(false),Boolean(false),Boolean(false),
String("t"),String("timestamp"),Boolean(true),Boolean(false),Boolean(false),Boolean(false),
String("a"),String("int"),Boolean(false),Boolean(true),Boolean(false),Boolean(false),
String("dic"),String("string"),Boolean(false),Boolean(true),Boolean(false),Boolean(true),
String("b"),String("string"),Boolean(false),Boolean(true),Boolean(false),Boolean(false),
String("add_dic"),String("string"),Boolean(false),Boolean(true),Boolean(false),Boolean(true),


INSERT INTO TABLE `05_alter_table_t0`(a, b, t, dic, add_dic) values (2, '2', 2, "d2", "d3");

affected_rows: 1

SELECT * FROM `05_alter_table_t0`;

tsid,t,a,dic,b,add_dic,
UInt64(0),Timestamp(1),Int32(1),String("d1"),String(""),String(""),
UInt64(0),Timestamp(2),Int32(2),String("d2"),String("2"),String("d3"),


ALTER TABLE `05_alter_table_t0` DROP COLUMN b;

Failed to execute query, err: Server(ServerError { code: 500, msg: "Failed to create plan, query: ALTER TABLE `05_alter_table_t0` DROP COLUMN b;. Caused by: Failed to create plan, err:Unsupported SQL statement" })
Expand All @@ -57,13 +83,14 @@ String("t"),String("timestamp"),Boolean(true),Boolean(false),Boolean(false),Bool
String("a"),String("int"),Boolean(false),Boolean(true),Boolean(false),Boolean(false),
String("dic"),String("string"),Boolean(false),Boolean(true),Boolean(false),Boolean(true),
String("b"),String("string"),Boolean(false),Boolean(true),Boolean(false),Boolean(false),
String("add_dic"),String("string"),Boolean(false),Boolean(true),Boolean(false),Boolean(true),


SELECT * FROM `05_alter_table_t0`;

tsid,t,a,dic,b,
UInt64(0),Timestamp(1),Int32(1),String("d1"),String(""),
UInt64(0),Timestamp(2),Int32(2),String("d2"),String("2"),
tsid,t,a,dic,b,add_dic,
UInt64(0),Timestamp(1),Int32(1),String("d1"),String(""),String(""),
UInt64(0),Timestamp(2),Int32(2),String("d2"),String("2"),String("d3"),


DROP TABLE `05_alter_table_t0`;
Expand Down
10 changes: 5 additions & 5 deletions integration_tests/cases/env/local/ddl/alter_table.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,15 +13,15 @@ INSERT INTO TABLE `05_alter_table_t0`(a, b, t, dic) values (2, '2', 2, "d2");
SELECT * FROM `05_alter_table_t0`;

-- waiting for datafusion's bug fix
-- ALTER TABLE `05_alter_table_t0` add COLUMN (add_dic string dictionary);
-- DESCRIBE TABLE `05_alter_table_t0`;
-- INSERT INTO TABLE `05_alter_table_t0`(a, b, t, dic, add_dic) values (2, '2', 2, "d2", "d3");
-- SELECT * FROM `05_alter_table_t0`;
ALTER TABLE `05_alter_table_t0` add COLUMN (add_dic string dictionary);
DESCRIBE TABLE `05_alter_table_t0`;
INSERT INTO TABLE `05_alter_table_t0`(a, b, t, dic, add_dic) values (2, '2', 2, "d2", "d3");
SELECT * FROM `05_alter_table_t0`;


-- doesn't support drop column
ALTER TABLE `05_alter_table_t0` DROP COLUMN b;
DESCRIBE TABLE `05_alter_table_t0`;
SELECT * FROM `05_alter_table_t0`;

DROP TABLE `05_alter_table_t0`;
DROP TABLE `05_alter_table_t0`;