andygrove commented on code in PR #5808:
URL: https://github.com/apache/datafusion-comet/pull/5808#discussion_r3973649180
##########
native/jni-bridge/src/errors.rs:
##########
@@ -573,7 +573,12 @@ fn throw_exception(env: &mut Env, error: &CometError,
backtrace: Option<String>)
// FAILED_READ_FILE / FileNotFound via the structured SparkError
channel. Anything else
// falls back to generic handling.
CometError::DataFusion { msg: _, source } => {
- if let Some(spark_error) =
try_classify_file_read_error(source) {
+ if let Some(json_message) = spark_error_json_in_chain(source) {
+ env.throw_new(
+
jni::jni_str!("org/apache/comet/exceptions/CometQueryExecutionException"),
Review Comment:
This makes four copies of
`throw_new("org/apache/comet/exceptions/CometQueryExecutionException",
json_message)` in this file, at 543, 549, 578 and 667. Would you pull out a
`throw_spark_error_json(env, json_message)` and have
`throw_spark_error_as_json` call it too? The `External` arm just above also
hand-rolls the same two downcasts `spark_error_json_in_chain` does, only
without the chain walk, so pointing it at the new function would remove that
duplication and stop the two arms disagreeing about how deep they look.
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -57,27 +64,99 @@ use
datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion_datasource::PartitionedFile;
use futures::future::BoxFuture;
use futures::{FutureExt, TryFutureExt};
+use object_store::path::Path;
use object_store::{ObjectStore, ObjectStoreExt};
use parquet::arrow::arrow_reader::ArrowReaderOptions;
use parquet::arrow::async_reader::AsyncFileReader;
+use parquet::arrow::parquet_to_arrow_schema;
use parquet::errors::ParquetError;
use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData};
+use std::collections::HashMap;
use std::fmt::Debug;
use std::ops::Range;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex, MutexGuard, PoisonError, Weak};
#[derive(Debug)]
pub struct EagerPageIndexReaderFactory {
store: Arc<dyn ObjectStore>,
metadata_cache: Arc<FileMetadataCache>,
+ field_id_check: Option<Arc<FieldIdCheck>>,
}
impl EagerPageIndexReaderFactory {
pub fn new(store: Arc<dyn ObjectStore>, metadata_cache:
Arc<FileMetadataCache>) -> Self {
Self {
store,
metadata_cache,
+ field_id_check: None,
+ }
+ }
+
+ /// Validate the ids `requested_schema` carries against each file's schema
as its footer
+ /// loads. Installs nothing when field id matching is off or the schema
carries no id, so
+ /// ordinary reads pay nothing.
+ pub fn with_field_id_check(
+ mut self,
+ requested_schema: SchemaRef,
+ parquet_options: &SparkParquetOptions,
+ ) -> Self {
+ if parquet_options.use_field_id &&
schema_holds_field_ids(&requested_schema) {
+ self.field_id_check = Some(Arc::new(FieldIdCheck {
+ requested_schema,
+ parquet_options: parquet_options.clone(),
+ validated: Mutex::new(HashMap::new()),
+ }));
}
+ self
+ }
+}
+
+/// Validates requested field ids for files the expression adapter never sees:
DataFusion's
+/// opener creates the adapter only when a predicate is pushed or the file
schema differs from
+/// the requested one, so a metadata-free file whose schema equals it is read
positionally
+/// (comet#5801). Resolves the mapping the adapter resolves, so both raise the
same error.
+#[derive(Debug)]
+struct FieldIdCheck {
+ requested_schema: SchemaRef,
+ parquet_options: SparkParquetOptions,
+ /// Files already validated, keyed by path to the metadata they were
checked against, so a
+ /// footer served from `FileMetadataCache` is not rechecked on every open.
+ validated: Mutex<HashMap<Path, Weak<ParquetMetaData>>>,
Review Comment:
`parking_lot::Mutex` is already a dependency of `native/core` and is what
every other mutex in the crate uses, in `fair_pool.rs`, `task_shared.rs` and
`jni_api.rs`. Switching `validated` to it would let the `lock()` helper and the
`MutexGuard` and `PoisonError` imports go away. `PoisonError::into_inner` is
currently the only occurrence in the whole native tree.
##########
native/core/src/parquet/parquet_support.rs:
##########
@@ -148,43 +150,327 @@ impl SparkParquetOptions {
/// Spark-compatible cast implementation. Defers to DataFusion's cast where
that is known
/// to be compatible, and returns an error when a not supported and not
DF-compatible cast
-/// is requested.
+/// is requested. Resolves the nested field mapping for this one value; a
per-file caller
+/// resolves once and uses [`spark_parquet_convert_with_mapping`] for every
batch.
pub fn spark_parquet_convert(
arg: ColumnarValue,
data_type: &DataType,
parquet_options: &SparkParquetOptions,
+) -> DataFusionResult<ColumnarValue> {
+ let mapping =
+ resolve_field_mapping(&arg.data_type(), data_type,
parquet_options).map_err(spark_error)?;
+ spark_parquet_convert_with_mapping(arg, data_type, &mapping,
parquet_options)
+}
+
+/// [`spark_parquet_convert`] with a mapping already resolved for the value's
type.
+pub(crate) fn spark_parquet_convert_with_mapping(
+ arg: ColumnarValue,
+ data_type: &DataType,
+ mapping: &FieldMapping,
+ parquet_options: &SparkParquetOptions,
) -> DataFusionResult<ColumnarValue> {
match arg {
- ColumnarValue::Array(array) =>
Ok(ColumnarValue::Array(parquet_convert_array(
+ ColumnarValue::Array(array) => Ok(ColumnarValue::Array(convert_array(
array,
data_type,
+ mapping,
parquet_options,
+ None,
)?)),
ColumnarValue::Scalar(scalar) => {
// Note that normally CAST(scalar) should be fold in Spark JVM
side. However, for
// some cases e.g., scalar subquery, Spark will not fold it, so we
need to handle it
// here.
let array = scalar.to_array()?;
let scalar = ScalarValue::try_from_array(
- &parquet_convert_array(array, data_type, parquet_options)?,
+ &convert_array(array, data_type, mapping, parquet_options,
None)?,
0,
)?;
Ok(ColumnarValue::Scalar(scalar))
}
}
}
-fn parquet_convert_array(
- array: ArrayRef,
+/// Wrap a [`SparkError`] the way every native operator surfaces it to the JVM.
+pub(crate) fn spark_error(error: SparkError) -> DataFusionError {
+ DataFusionError::External(Box::new(error))
+}
+
+/// Outcome of matching one requested id or name against a struct's file
fields: the last
+/// file field that matched and whether more than one did. A plain `Copy`
value, so resolving
+/// a wide struct allocates nothing per id or per name; the matched names are
only gathered
+/// when an ambiguity is reported.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) struct FieldMatch {
+ pub(crate) index: usize,
+ pub(crate) ambiguous: bool,
+}
+
+impl FieldMatch {
+ pub(crate) fn new(index: usize, ambiguous: bool) -> Self {
+ Self { index, ambiguous }
+ }
+
+ /// The first file field carrying this id or name.
+ pub(crate) fn first(index: usize) -> Self {
+ Self::new(index, false)
+ }
+
+ /// A further file field carrying the same id or name: the later index
wins, as Spark's
+ /// `toMap` does for exact names, and the entry turns ambiguous.
+ pub(crate) fn also(self, index: usize) -> Self {
+ Self::new(index, true)
+ }
+}
+
+/// Record file field `index` under `key`, keeping the entry `Copy`-sized
however many fields
+/// share the key.
+pub(crate) fn record_field_match<K: Hash + Eq>(
+ matches: &mut HashMap<K, FieldMatch>,
+ key: K,
+ index: usize,
+) {
+ matches
+ .entry(key)
+ .and_modify(|m| *m = m.also(index))
+ .or_insert_with(|| FieldMatch::first(index));
+}
+
+/// Comma-joined names of the fields carrying `id`, for the duplicate-id error
message.
+pub(crate) fn field_names_with_id(fields: &Fields, id: i32) -> String {
+ fields
+ .iter()
+ .filter(|f| field_id(f) == Some(id))
+ .map(|f| f.name().as_str())
+ .collect::<Vec<_>>()
+ .join(", ")
+}
+
+/// Which file field supplies each requested field, resolved once per file and
reused for
+/// every batch. Follows the requested type as Spark's `clipParquetSchema`
does: a struct
+/// lists one source per requested field, a list (large or not) or map carries
the mapping
+/// of its element or key and value types, and anything else is a leaf
converted by type.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub(crate) enum FieldMapping {
+ Struct(Vec<StructFieldSource>),
+ List(Box<FieldMapping>),
+ Map(Box<FieldMapping>, Box<FieldMapping>),
+ Leaf,
+}
+
+/// The file field behind one requested struct field.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub(crate) struct StructFieldSource {
+ /// Index of the file field supplying the requested field; `None`
null-fills it.
+ pub(crate) from_index: Option<usize>,
+ /// Mapping of the requested field's own type.
+ pub(crate) nested: FieldMapping,
+}
+
+impl FieldMapping {
+ /// Mapping of a list's element type. A `Leaf` list converts its elements
by type alone,
+ /// as the adapter hands one to every column whose type holds no struct.
+ fn list_element(&self) -> DataFusionResult<&FieldMapping> {
+ match self {
+ FieldMapping::List(inner) => Ok(inner),
+ FieldMapping::Leaf => Ok(&FieldMapping::Leaf),
+ other => Err(DataFusionError::Internal(format!(
+ "list column resolved to a non-list field mapping: {other:?}"
+ ))),
+ }
+ }
+
+ /// Mappings of a map's key and value types; see
[`FieldMapping::list_element`].
+ fn map_entries(&self) -> DataFusionResult<(&FieldMapping, &FieldMapping)> {
+ match self {
+ FieldMapping::Map(key, value) => Ok((key, value)),
+ FieldMapping::Leaf => Ok((&FieldMapping::Leaf,
&FieldMapping::Leaf)),
+ other => Err(DataFusionError::Internal(format!(
+ "map column resolved to a non-map field mapping: {other:?}"
+ ))),
+ }
+ }
+
+ /// True when every requested field reads the file field at its own
position, so a
+ /// metadata-only relabel of the file array already yields the requested
layout.
+ pub(crate) fn is_positional(&self) -> bool {
+ match self {
+ FieldMapping::Struct(sources) => sources
+ .iter()
+ .enumerate()
+ .all(|(i, s)| s.from_index == Some(i) &&
s.nested.is_positional()),
+ FieldMapping::List(inner) => inner.is_positional(),
+ FieldMapping::Map(key, value) => key.is_positional() &&
value.is_positional(),
+ FieldMapping::Leaf => true,
+ }
+ }
+}
+
+/// True when a field of `schema`, at any nesting depth, carries a Parquet
field id.
+pub(crate) fn schema_holds_field_ids(schema: &Schema) -> bool {
+ schema.fields().iter().any(|f| field_holds_id(f))
+}
+
+fn field_holds_id(field: &Field) -> bool {
+ field_id(field).is_some()
+ || match field.data_type() {
+ DataType::Struct(fields) => fields.iter().any(|f|
field_holds_id(f)),
+ DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _)
=> field_holds_id(f),
+ _ => false,
+ }
+}
+
+/// Resolve every requested root field against `file_schema` the way the
expression adapter
+/// does, keeping only the ambiguity Spark reports. DataFusion's opener
creates the adapter
+/// only when a predicate is pushed or the file schema differs from the
requested one, so the
Review Comment:
The reason the opener skips the adapter is spelled out here, on
`FieldIdCheck`, and again in `init_datasource_exec`. Keeping it once on
`FieldIdCheck` and having the other two point at it would mean it only has to
be corrected in one place if DataFusion changes that condition.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
