[GitHub] [arrow-datafusion] izveigor commented on a diff in pull request #6936: feat: reading array elements

["izveigor (via GitHub)" <gi...@apache.org>]

izveigor commented on code in PR #6936:
URL: https://github.com/apache/arrow-datafusion/pull/6936#discussion_r1278016584


##########
datafusion/physical-expr/src/expressions/get_indexed_field.rs:
##########
@@ -74,70 +90,63 @@ impl PhysicalExpr for GetIndexedFieldExpr {
     }
 
     fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
-        let data_type = self.arg.data_type(input_schema)?;
-        get_data_type_field(&data_type, &self.key).map(|f| f.data_type().clone())
+        let arg_dt = self.arg.data_type(input_schema)?;
+        let key_dt = self.key.data_type(input_schema)?;
+        let extra_key_dt = if let Some(extra_key) = &self.extra_key {
+            Some(extra_key.data_type(input_schema)?)
+        } else {
+            None
+        };
+        get_data_type_field(&arg_dt, &key_dt, &extra_key_dt)
+            .map(|f| f.data_type().clone())
     }
 
     fn nullable(&self, input_schema: &Schema) -> Result<bool> {
-        let data_type = self.arg.data_type(input_schema)?;
-        get_data_type_field(&data_type, &self.key).map(|f| f.is_nullable())
+        let arg_dt = self.arg.data_type(input_schema)?;
+        let key_dt = self.key.data_type(input_schema)?;
+        let extra_key_dt = if let Some(extra_key) = &self.extra_key {
+            Some(extra_key.data_type(input_schema)?)
+        } else {
+            None
+        };
+        get_data_type_field(&arg_dt, &key_dt, &extra_key_dt).map(|f| f.is_nullable())
     }
 
     fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
         let array = self.arg.evaluate(batch)?.into_array(1);
-        match (array.data_type(), &self.key) {
-            (DataType::List(_) | DataType::Struct(_), _) if self.key.is_null() => {
-                let scalar_null: ScalarValue = array.data_type().try_into()?;
-                Ok(ColumnarValue::Scalar(scalar_null))
+        let key = self.key.evaluate(batch)?.into_array(1);
+        if let Some(extra_key) = &self.extra_key {
+            let extra_key = extra_key.evaluate(batch)?.into_array(1);
+            match (array.data_type(), key.data_type(), extra_key.data_type()) {
+                (DataType::List(_), DataType::Int64, DataType::Int64)
+                | (DataType::List(_), DataType::Int64, DataType::Null)
+                | (DataType::List(_), DataType::Null, DataType::Int64)
+                | (DataType::List(_), DataType::Null, DataType::Null) => Ok(ColumnarValue::Array(array_slice(&[

Review Comment:
   It does not work with columns, but all functions support column format.
   @alamb if you have free time can you give me some advices how to solve the problem?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org