Skip to content

Commit

Permalink
fix: missing dict attribute
Browse files Browse the repository at this point in the history
  • Loading branch information
tanruixiang committed Jun 28, 2023
1 parent acb5d97 commit 64727de
Show file tree
Hide file tree
Showing 2 changed files with 360 additions and 14 deletions.
66 changes: 64 additions & 2 deletions datafusion/common/src/dfschema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -586,27 +586,53 @@ pub trait ExprSchema: std::fmt::Debug {

/// What is the datatype of this column?
fn data_type(&self, col: &Column) -> Result<&DataType>;

/// Is this column reference dict_is_ordered?
fn dict_is_ordered(&self, col: &Column) -> Result<bool>;

/// What is the dict_id of this column?
fn dict_id(&self, col: &Column) -> Result<i64>;
}

// Implement `ExprSchema` for `Arc<DFSchema>`
impl<P: AsRef<DFSchema> + std::fmt::Debug> ExprSchema for P {
fn nullable(&self, col: &Column) -> Result<bool> {
self.as_ref().nullable(col)
}

fn data_type(&self, col: &Column) -> Result<&DataType> {
self.as_ref().data_type(col)
}

fn dict_is_ordered(&self, col: &Column) -> Result<bool> {
self.as_ref().dict_is_ordered(col)
}

fn dict_id(&self, col: &Column) -> Result<i64> {
self.as_ref().dict_id(col)
}
}

impl ExprSchema for DFSchema {
fn nullable(&self, col: &Column) -> Result<bool> {
Ok(self.field_from_column(col)?.is_nullable())
}

fn data_type(&self, col: &Column) -> Result<&DataType> {
Ok(self.field_from_column(col)?.data_type())
}

fn dict_is_ordered(&self, col: &Column) -> Result<bool> {
match self.field_from_column(col)?.field().dict_is_ordered() {
Some(dict_id_ordered) => Ok(dict_id_ordered),
_ => Ok(false),
}
}

fn dict_id(&self, col: &Column) -> Result<i64> {
match self.field_from_column(col)?.field().dict_id() {
Some(dict_id_ordered) => Ok(dict_id_ordered),
_ => Ok(0),
}
}
}

/// DFField wraps an Arrow field and adds an optional qualifier
Expand Down Expand Up @@ -635,6 +661,42 @@ impl DFField {
/// Convenience method for creating new `DFField` without a qualifier
pub fn new_unqualified(name: &str, data_type: DataType, nullable: bool) -> Self {
DFField {
None,
field: Arc::new(Field::new(name, data_type, nullable)),
}
}
/// Creates a new `DFField` with dict
pub fn new_dict(
qualifier: Option<&str>,
name: &str,
data_type: DataType,
nullable: bool,
dict_id: i64,
dict_is_ordered: bool,
) -> Self {
DFField {
qualifier: qualifier.map(|s| s.into()),
field: Field::new_dict(name, data_type, nullable, dict_id, dict_is_ordered),
}
}

/// Convenience method for creating new `DFField` without a qualifier
pub fn new_unqualified_dict(
name: &str,
data_type: DataType,
nullable: bool,
dict_id: i64,
dict_is_ordered: bool,
) -> Self {
DFField {
qualifier: None,
field: Field::new_dict(name, data_type, nullable, dict_id, dict_is_ordered),
}
}

/// Create an unqualified field from an existing Arrow field
pub fn from(field: Field) -> Self {
Self {
qualifier: None,
field: Arc::new(Field::new(name, data_type, nullable)),
}
Expand Down
Loading

0 comments on commit 64727de

Please sign in to comment.