Revert dictionary lookup for indexed fields

Author: Igosuki

datafusion/Cargo.toml

@@ -45,7 +45,7 @@ unicode_expressions = ["unicode-segmentation"]
 # Used for testing ONLY: causes all values to hash to the same value (test for collisions)
 force_hash_collisions = []
 # Used to enable the avro format
-avro = ["avro-rs"]
+avro = ["avro-rs", "num-traits"]
 
 [dependencies]
 ahash = "0.7"
@@ -72,7 +72,7 @@ lazy_static = { version = "^1.4.0", optional = true }
 smallvec = { version = "1.6", features = ["union"] }
 rand = "0.8"
 avro-rs = { version = "0.13", features = ["snappy"], optional = true }
-num-traits = { version = "0.2" }
+num-traits = { version = "0.2", optional = true }
 
 [dev-dependencies]
 criterion = "0.3"

datafusion/src/field_util.rs

@@ -29,21 +29,6 @@ use crate::scalar::ScalarValue;
 /// * there is no field key is not of the required index type
 pub fn get_indexed_field(data_type: &DataType, key: &ScalarValue) -> Result<Field> {
     match (data_type, key) {
-        (DataType::Dictionary(ref kt, ref _vt), ScalarValue::Utf8(Some(k))) => {
-            match kt.as_ref() {
-                DataType::Int8 | DataType::Int16 |DataType::Int32 |DataType::Int64 |DataType::UInt8
-                | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => {
-                    Ok(Field::new(k, *kt.clone(), true))
-                },
-                _ => Err(DataFusionError::Plan(format!("The key for a dictionary has to be a primitive type, was : \"{}\"", key))),
-            }
-        },
-        (DataType::Dictionary(_, _), _) => {
-            Err(DataFusionError::Plan(
-                "Only utf8 types are valid for dictionary lookup"
-                    .to_string(),
-            ))
-        }
         (DataType::List(lt), ScalarValue::Int64(Some(i))) => {
             Ok(Field::new(&i.to_string(), lt.data_type().clone(), false))
         }

datafusion/src/physical_plan/expressions/get_indexed_field.rs

@@ -26,19 +26,14 @@ use arrow::{
 
 use crate::arrow::array::Array;
 use crate::arrow::compute::concat;
-use crate::arrow::datatypes::{
-    Int16Type, Int32Type, Int64Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
-};
 use crate::scalar::ScalarValue;
 use crate::{
     error::DataFusionError,
     error::Result,
     field_util::get_indexed_field as get_data_type_field,
     physical_plan::{ColumnarValue, PhysicalExpr},
 };
-use arrow::array::{ArrayRef, DictionaryArray, ListArray};
-use arrow::datatypes::{ArrowPrimitiveType, Int8Type};
-use num_traits::ToPrimitive;
+use arrow::array::ListArray;
 use std::fmt::Debug;
 
 /// expression to get a field of a struct array.
@@ -96,25 +91,8 @@ impl PhysicalExpr for GetIndexedFieldExpr {
                     let iter = concat(vec.as_slice()).unwrap();
                     Ok(ColumnarValue::Array(iter))
                 }
-                (DataType::Dictionary(ref kt, _), ScalarValue::Utf8(Some(s))) => {
-                    match **kt {
-                        DataType::Int8 => dict_lookup::<Int8Type>(array, s),
-                        DataType::Int16 => dict_lookup::<Int16Type>(array, s),
-                        DataType::Int32 => dict_lookup::<Int32Type>(array, s),
-                        DataType::Int64 => dict_lookup::<Int64Type>(array, s),
-                        DataType::UInt8 => dict_lookup::<UInt8Type>(array, s),
-                        DataType::UInt16 => dict_lookup::<UInt16Type>(array, s),
-                        DataType::UInt32 => dict_lookup::<UInt32Type>(array, s),
-                        DataType::UInt64 => dict_lookup::<UInt64Type>(array, s),
-                        _ => Err(DataFusionError::NotImplemented(
-                            "dictionary lookup only available for numeric keys"
-                                .to_string(),
-                        )),
-                    }
-                }
                 _ => Err(DataFusionError::NotImplemented(
-                    "get indexed field is only possible on dictionary and list"
-                        .to_string(),
+                    "get indexed field is only possible on lists".to_string(),
                 )),
             },
             ColumnarValue::Scalar(_) => Err(DataFusionError::NotImplemented(
@@ -131,25 +109,3 @@ pub fn get_indexed_field(
 ) -> Result<Arc<dyn PhysicalExpr>> {
     Ok(Arc::new(GetIndexedFieldExpr::new(arg, key)))
 }
-
-fn dict_lookup<T: ArrowPrimitiveType>(
-    array: ArrayRef,
-    lookup: &str,
-) -> Result<ColumnarValue>
-where
-    T::Native: num_traits::cast::ToPrimitive,
-{
-    let as_dict_array = array.as_any().downcast_ref::<DictionaryArray<T>>().unwrap();
-    if let Some(index) = as_dict_array.lookup_key(lookup) {
-        Ok(ColumnarValue::Array(
-            as_dict_array
-                .keys()
-                .slice(ToPrimitive::to_usize(&index).unwrap(), 1),
-        ))
-    } else {
-        Err(DataFusionError::NotImplemented(format!(
-            "key not found in dictionary for : {}",
-            lookup
-        )))
-    }
-}

datafusion/tests/sql.rs

@@ -5168,18 +5168,12 @@ async fn query_nested_get_indexed_field() -> Result<()> {
     let mut ctx = ExecutionContext::new();
     let nested_dt = DataType::List(Box::new(Field::new("item", DataType::Int64, true)));
     // Nested schema of { "some_list": [[i64]] }
-    let schema = Arc::new(Schema::new(vec![
-        Field::new(
-            "some_list",
-            DataType::List(Box::new(Field::new("item", nested_dt.clone(), true))),
-            false,
-        ),
-        Field::new(
-            "some_dict",
-            DataType::Dictionary(Box::new(DataType::Int64), Box::new(DataType::Utf8)),
-            false,
-        ),
-    ]));
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "some_list",
+        DataType::List(Box::new(Field::new("item", nested_dt.clone(), true))),
+        false,
+    )]));
+
     let builder = PrimitiveBuilder::<Int64Type>::new(3);
     let nested_lb = ListBuilder::new(builder);
     let mut lb = ListBuilder::new(nested_lb);
@@ -5199,20 +5193,7 @@ async fn query_nested_get_indexed_field() -> Result<()> {
         lb.append(true).unwrap();
     }
 
-    let dictionary_values = StringArray::from(vec![Some("a"), Some("b"), Some("c")]);
-    let mut sb = StringDictionaryBuilder::new_with_dictionary(
-        PrimitiveBuilder::<Int64Type>::new(3),
-        &dictionary_values,
-    )
-    .unwrap();
-    for s in &["b", "a", "c"] {
-        sb.append(s).unwrap();
-    }
-
-    let data = RecordBatch::try_new(
-        schema.clone(),
-        vec![Arc::new(lb.finish()), Arc::new(sb.finish())],
-    )?;
+    let data = RecordBatch::try_new(schema.clone(), vec![Arc::new(lb.finish())])?;
     let table = MemTable::try_new(schema, vec![vec![data]])?;
     let table_a = Arc::new(table);
 
@@ -5227,9 +5208,4 @@ async fn query_nested_get_indexed_field() -> Result<()> {
     let actual = execute(&mut ctx, sql).await;
     let expected = vec![vec!["0"], vec!["5"], vec!["11"]];
     assert_eq!(expected, actual);
-    let sql = r#"SELECT some_dict["b"], some_dict["a"] FROM ints LIMIT 3"#;
-    let actual = execute(&mut ctx, sql).await;
-    let expected = vec![vec!["0", "1"]];
-    assert_eq!(expected, actual);
-    Ok(())
 }