rich7420 commented on code in PR #6246:
URL: https://github.com/apache/datafusion-comet/pull/6246#discussion_r4471223745
##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -63,6 +68,187 @@ pub(crate) struct CometParquetScanOptions {
pub prefetch_next_file_metadata: bool,
}
+/// Keep top-level Spark string fields as binary while arrow-rs decodes
Parquet. arrow-rs only
+/// validates UTF-8 for columns annotated as strings, so decoding an
unannotated BINARY column that
+/// `spark.sql.parquet.binaryAsString=true` types as a string straight into
Arrow Utf8 would admit
+/// invalid UTF-8. A projection added at the scan boundary converts the binary
arrays to valid Utf8
+/// with Spark-compatible replacement semantics.
+fn binary_as_string_read_schema(schema: &SchemaRef) -> SchemaRef {
+ Arc::new(Schema::new_with_metadata(
+ schema
+ .fields()
+ .iter()
+ .map(binary_as_string_read_field)
+ .collect::<Vec<_>>(),
+ schema.metadata().clone(),
+ ))
+}
+
+/// The field `binary_as_string_read_schema` reads a top-level `field` as.
Spark's StringType is
+/// Arrow Utf8.
+fn binary_as_string_read_field(field: &FieldRef) -> FieldRef {
+ if is_binary_as_string_read(field.data_type()) {
+ Arc::new(field.as_ref().clone().with_data_type(DataType::Binary))
+ } else {
+ Arc::clone(field)
+ }
+}
+
+/// Whether `binary_as_string_read_schema` reads a top-level field of
`data_type` as Binary.
+fn is_binary_as_string_read(data_type: &DataType) -> bool {
+ data_type == &DataType::Utf8
+}
+
+/// Whether `expr` references a column whose folded name is in `folded_names`.
+fn references_any_column(
+ expr: &Arc<dyn PhysicalExpr>,
+ folded_names: &HashSet<String>,
+ case_sensitive: bool,
+) -> Result<bool, ExecutionError> {
+ for column in collect_columns(expr) {
+ if folded_names.contains(&fold_name(column.name(), case_sensitive)?) {
+ return Ok(true);
+ }
+ }
+ Ok(false)
+}
+
+/// Rewrites a data filter for a scan that reads string columns as Binary (see
+/// `binary_as_string_read_schema`), or returns `None` when it can't go into
the scan.
+///
+/// The scan decodes each such column with replacement: a valid value decodes
to itself, and an
+/// invalid one to a string containing U+FFFD. So for a literal without
U+FFFD, `s = 'lit'` on the
+/// decoded string holds exactly when `s = X'<bytes of lit>'` holds on the raw
bytes, and the same
+/// goes for `IN`. `IS [NOT] NULL` holds on both, because decoding keeps
nulls. Other filters on
+/// these columns stay above the scan. A range can disagree because decoding
can reorder values:
+/// `"a\x80"` sorts below `"aé"` as bytes, but `"a\u{FFFD}"` sorts above it.
+///
+/// Spark's Filter above the scan re-evaluates every data filter, so a pushed
filter only has to
+/// keep every row the original keeps. That's why an `AND` keeps whichever
side can be pushed.
+fn binary_as_string_pushdown_filter(
+ filter: &Arc<dyn PhysicalExpr>,
+ binary_read_columns: &HashSet<String>,
+ case_sensitive: bool,
+ filter_read_schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>, ExecutionError> {
+ if !references_any_column(filter, binary_read_columns, case_sensitive)? {
+ return Ok(Some(Arc::clone(filter)));
+ }
+ let reads_binary = |expr: &Arc<dyn PhysicalExpr>| -> Result<bool,
ExecutionError> {
+ match expr.downcast_ref::<Column>() {
+ Some(column) => {
+ Ok(binary_read_columns.contains(&fold_name(column.name(),
case_sensitive)?))
+ }
+ None => Ok(false),
+ }
+ };
+ let pushdown = |expr: &Arc<dyn PhysicalExpr>| {
+ binary_as_string_pushdown_filter(
+ expr,
+ binary_read_columns,
+ case_sensitive,
+ filter_read_schema,
+ )
+ };
+
+ if let Some(binary) = filter.downcast_ref::<BinaryExpr>() {
+ let (left, right) = (binary.left(), binary.right());
+ match binary.op() {
+ Operator::And => {
+ return Ok(match (pushdown(left)?, pushdown(right)?) {
+ (Some(left), Some(right)) => {
+ Some(Arc::new(BinaryExpr::new(left, Operator::And,
right)))
+ }
+ (side, None) | (None, side) => side,
+ });
+ }
+ Operator::Or => {
+ return Ok(match (pushdown(left)?, pushdown(right)?) {
+ (Some(left), Some(right)) => {
+ Some(Arc::new(BinaryExpr::new(left, Operator::Or,
right)))
+ }
+ _ => None,
+ });
+ }
+ Operator::Eq => {
+ if reads_binary(left)? {
+ if let Some(literal) = binary_literal(right) {
+ let eq = BinaryExpr::new(Arc::clone(left),
Operator::Eq, literal);
+ return Ok(Some(Arc::new(eq)));
+ }
+ }
+ if reads_binary(right)? {
+ if let Some(literal) = binary_literal(left) {
+ let eq = BinaryExpr::new(literal, Operator::Eq,
Arc::clone(right));
+ return Ok(Some(Arc::new(eq)));
+ }
+ }
+ }
+ _ => {}
+ }
+ } else if let Some(list) = filter.downcast_ref::<InListExpr>() {
+ if !list.negated() && reads_binary(list.expr())? {
+ if let Some(values) =
list.list().iter().map(binary_literal).collect() {
+ let expr = Arc::clone(list.expr());
+ return Ok(Some(in_list(expr, values, &false,
filter_read_schema)?));
+ }
+ }
+ } else if let Some(is_null) = filter.downcast_ref::<IsNullExpr>() {
+ if reads_binary(is_null.arg())? {
+ return Ok(Some(Arc::clone(filter)));
+ }
+ } else if let Some(is_not_null) = filter.downcast_ref::<IsNotNullExpr>() {
+ if reads_binary(is_not_null.arg())? {
+ return Ok(Some(Arc::clone(filter)));
+ }
+ }
+ Ok(None)
+}
+
+/// `expr` as a Binary literal holding the same bytes, if it is a string
literal without U+FFFD.
+fn binary_literal(expr: &Arc<dyn PhysicalExpr>) -> Option<Arc<dyn
PhysicalExpr>> {
+ match expr.downcast_ref::<Literal>()?.value() {
+ ScalarValue::Utf8(Some(value)) if
!value.contains(char::REPLACEMENT_CHARACTER) => {
+ let bytes = value.as_bytes().to_vec();
+ Some(Arc::new(Literal::new(ScalarValue::Binary(Some(bytes)))))
+ }
+ _ => None,
+ }
+}
+
+fn binary_as_string_projection(
Review Comment:
Please retain range pruning for annotated strings: `s < 'two'`.
--
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]
