parthchandra commented on code in PR #6784:
URL: https://github.com/apache/datafusion-comet/pull/6784#discussion_r4461295849
##########
native/core/src/execution/planner.rs:
##########
@@ -1922,6 +1923,55 @@ impl PhysicalPlanner {
Arc::new(SparkPlan::new(spark_plan.plan_id, scan, vec![])),
))
}
+ OpStruct::TextScan(scan) => {
+ let data_schema =
convert_spark_types_to_arrow_schema(scan.data_schema.as_slice());
+ let partition_schema =
+
convert_spark_types_to_arrow_schema(scan.partition_schema.as_slice());
+ let projection_vector: Vec<usize> =
+ scan.projection_vector.iter().map(|i| *i as
usize).collect();
+ let object_store_options: HashMap<String, String> =
+ scan.object_store_options.clone();
+ let partition = usize::try_from(self.partition)
+ .ok()
+ .and_then(|p| scan.file_partitions.get(p))
+ .ok_or_else(|| {
+ GeneralError(format!(
+ "TextScan has no file partition {} ({}
partitions)",
+ self.partition,
+ scan.file_partitions.len()
+ ))
+ })?;
+ // Build the store from this partition's files; other
partitions may read
+ // from a different one.
+ let object_store_url = match
partition.partitioned_file.first() {
+ Some(file) => {
+ prepare_object_store_with_configs(
+ self.session_ctx.runtime_env(),
+ file.file_path.clone(),
+ &object_store_options,
+ )?
+ .0
Review Comment:
Done. The Text arm now builds its store through scan_store, so a plan made
without data builds no real store. The partition lookup, store and file group
are one helper, scan_task_files, used by both the CSV and Text arms so they
can't drift again. I also removed the unused partition_schema from TextScan and
reserved tag 2.
##########
native/core/src/execution/operators/text_scan.rs:
##########
@@ -0,0 +1,796 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Native Text file scan, mirroring Spark's `text` data source.
+//!
+//! Spark's text reader produces a single `value: string` column. With the
default options each
+//! line of the file is one row (split on `\n`, `\r\n`, or `\r`, with the
terminator stripped).
+//! With `wholetext=true` the whole file is a single row. A custom `lineSep`
splits on exactly that
+//! byte sequence instead of the universal-newline set. See the
`TextOptions`/`HadoopFileLinesReader`
+//! logic in Spark's `org.apache.spark.sql.execution.datasources.text` package.
+
+use crate::execution::operators::ExecutionError;
+use arrow::array::StringBuilder;
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::{RecordBatch, RecordBatchOptions};
+use bytes::Bytes;
+use datafusion::common::Result;
+use datafusion::datasource::object_store::ObjectStoreUrl;
+use datafusion::physical_expr::projection::ProjectionExprs;
+use datafusion::physical_plan::metrics::{BaselineMetrics,
ExecutionPlanMetricsSet, MetricBuilder};
+use datafusion::physical_plan::DisplayFormatType;
+use datafusion_comet_common::decode_utf8_spark_lossy;
+use datafusion_comet_proto::spark_operator::TextOptions;
+use datafusion_datasource::file::FileSource;
+use datafusion_datasource::file_compression_type::FileCompressionType;
+use datafusion_datasource::file_groups::FileGroup;
+use datafusion_datasource::file_scan_config::{FileScanConfig,
FileScanConfigBuilder};
+use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
+use datafusion_datasource::projection::{ProjectionOpener, SplitProjection};
+use datafusion_datasource::source::DataSourceExec;
+use datafusion_datasource::{as_file_source, PartitionedFile, TableSchema};
+use futures::StreamExt;
+use object_store::{ObjectStore, ObjectStoreExt};
+use std::borrow::Cow;
+use std::fmt;
+use std::io::Read;
+use std::sync::Arc;
+
+/// Batch size used when the execution framework does not set one before
opening the file.
+const DEFAULT_TEXT_BATCH_SIZE: usize = 8192;
+
+/// UTF-8 byte-order mark. Spark's line reader strips a leading one; wholetext
keeps it.
+const UTF8_BOM: &[u8] = &[0xEF, 0xBB, 0xBF];
+
+/// Cap on a batch's decoded `value` buffer bytes, kept well below Arrow's 2GB
32-bit offset limit
+/// so a batch of many rows (or a few long lines) cannot overflow. Measured
against the decoded
+/// string length, since a lossy UTF-8 decode can expand ill-formed bytes up
to 3x.
+const MAX_BATCH_VALUE_BYTES: usize = 1 << 30;
+
+pub fn init_text_datasource_exec(
+ object_store_url: ObjectStoreUrl,
+ file_groups: Vec<Vec<PartitionedFile>>,
+ data_schema: SchemaRef,
+ _partition_schema: SchemaRef,
+ projection_vector: Vec<usize>,
+ text_options: &TextOptions,
+) -> Result<Arc<DataSourceExec>, ExecutionError> {
+ let text_source = TextSource::new(data_schema, text_options);
+
+ let file_groups = file_groups.into_iter().map(FileGroup::new).collect();
+
+ let file_scan_config = FileScanConfigBuilder::new(object_store_url,
Arc::new(text_source))
+ .with_file_groups(file_groups)
+ .with_projection_indices(Some(projection_vector))?
+ .build();
+
+ Ok(DataSourceExec::from_data_source(file_scan_config))
+}
+
+/// A [`FileSource`] for Spark's text format. The file schema is always the
single `value` column;
+/// projection (and partition columns, if any) is handled by
[`ProjectionOpener`], exactly as the
+/// CSV source does.
+#[derive(Debug, Clone)]
+struct TextSource {
+ table_schema: TableSchema,
+ projection: SplitProjection,
+ batch_size: Option<usize>,
+ metrics: ExecutionPlanMetricsSet,
+ whole_text: bool,
+ // The raw separator bytes when a custom `lineSep` is set; None means
universal newline
+ // (`\n`, `\r\n`, `\r`).
+ line_sep: Option<Vec<u8>>,
+}
+
+impl TextSource {
+ fn new(file_schema: SchemaRef, options: &TextOptions) -> Self {
+ let table_schema: TableSchema = file_schema.into();
+ // `line_sep` arrives as the exact separator bytes Spark computed
(`lineSeparatorInRead`,
+ // already encoded with the file's charset), so we split on them
directly.
+ let line_sep = options.line_sep.clone().filter(|sep| !sep.is_empty());
+ Self {
+ projection: SplitProjection::unprojected(&table_schema),
+ table_schema,
+ batch_size: None,
+ metrics: ExecutionPlanMetricsSet::new(),
+ whole_text: options.whole_text,
+ line_sep,
+ }
+ }
+}
+
+impl From<TextSource> for Arc<dyn FileSource> {
+ fn from(source: TextSource) -> Self {
+ as_file_source(source)
+ }
+}
+
+impl FileSource for TextSource {
+ fn create_file_opener(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ base_config: &FileScanConfig,
+ partition_index: usize,
+ ) -> Result<Arc<dyn FileOpener>> {
+ // The inner opener must emit batches whose schema is the file schema
projected by
+ // `file_indices` (that is what ProjectionOpener feeds its projector).
For text that is
+ // either the single `value` column or, for a projection-less count,
an empty schema.
+ let file_schema = self.table_schema.file_schema();
+ let projected_file_schema =
Arc::new(file_schema.project(&self.projection.file_indices)?);
+ let opener = Arc::new(TextOpener {
+ projected_file_schema,
+ file_compression_type: base_config.file_compression_type,
+ object_store,
+ partition_index,
+ batch_size: self.batch_size.unwrap_or(DEFAULT_TEXT_BATCH_SIZE),
+ metrics: self.metrics.clone(),
+ whole_text: self.whole_text,
+ line_sep: self.line_sep.clone(),
+ }) as Arc<dyn FileOpener>;
+ ProjectionOpener::try_new(self.projection.clone(), opener, file_schema)
+ }
+
+ fn table_schema(&self) -> &TableSchema {
+ &self.table_schema
+ }
+
+ fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
+ let mut conf = self.clone();
+ conf.batch_size = Some(batch_size);
+ Arc::new(conf)
+ }
+
+ fn try_pushdown_projection(
+ &self,
+ projection: &ProjectionExprs,
+ ) -> Result<Option<Arc<dyn FileSource>>> {
+ let mut source = self.clone();
+ let new_projection = self.projection.source.try_merge(projection)?;
+ source.projection =
SplitProjection::new(self.table_schema.file_schema(), &new_projection);
+ Ok(Some(Arc::new(source)))
+ }
+
+ fn projection(&self) -> Option<&ProjectionExprs> {
+ Some(&self.projection.source)
+ }
+
+ fn metrics(&self) -> &ExecutionPlanMetricsSet {
+ &self.metrics
+ }
+
+ fn file_type(&self) -> &str {
+ "text"
+ }
+
+ // A custom line separator can straddle a byte-range boundary, and we read
whole files rather
+ // than boundary-aligned ranges, so this source must not let DataFusion
repartition files by
+ // byte range.
+ fn supports_repartitioning(&self) -> bool {
+ false
+ }
+
+ fn fmt_extra(&self, t: DisplayFormatType, f: &mut fmt::Formatter) ->
fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, ", whole_text={}", self.whole_text)
+ }
+ DisplayFormatType::TreeRender => Ok(()),
+ }
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(
+ &Arc<dyn datafusion::physical_plan::PhysicalExpr>,
+ ) -> Result<datafusion::common::tree_node::TreeNodeRecursion>,
+ ) -> Result<datafusion::common::tree_node::TreeNodeRecursion> {
+
datafusion::physical_plan::apply_expression_roots(self.projection.source.iter(),
f)
+ }
+}
+
+/// A [`FileOpener`] that reads a whole text file and yields the projected
file schema in batches.
+struct TextOpener {
+ // The file schema projected by `file_indices`: the single `value` column,
or an empty schema
+ // for a projection-less count.
+ projected_file_schema: SchemaRef,
+ file_compression_type: FileCompressionType,
+ object_store: Arc<dyn ObjectStore>,
+ partition_index: usize,
+ batch_size: usize,
+ metrics: ExecutionPlanMetricsSet,
+ whole_text: bool,
+ line_sep: Option<Vec<u8>>,
+}
+
+impl FileOpener for TextOpener {
+ fn open(&self, partitioned_file: PartitionedFile) ->
Result<FileOpenFuture> {
+ let store = Arc::clone(&self.object_store);
+ let projected_file_schema = Arc::clone(&self.projected_file_schema);
+ let file_compression_type = self.file_compression_type;
+ let batch_size = self.batch_size;
+ let whole_text = self.whole_text;
+ let line_sep = self.line_sep.clone();
+ let baseline_metrics = BaselineMetrics::new(&self.metrics,
self.partition_index);
+ // Feeds the task's input bytes in the Spark UI (see
CometMetricNode.reportScanInputMetrics).
+ let bytes_scanned =
+ MetricBuilder::new(&self.metrics).counter("bytes_scanned",
self.partition_index);
+
+ Ok(Box::pin(async move {
+ let range = partitioned_file.range.clone();
+ let location = partitioned_file.object_meta.location;
+ let raw = store.get(&location).await?.bytes().await?;
+ bytes_scanned.add(raw.len());
+
+ // Decompress only when the file is actually compressed. On
Comet's plan config the file
+ // is always UNCOMPRESSED (CometScanRule falls back for compressed
files), so we keep
+ // the fetched bytes as they are and avoid a full-file copy; the
decompress arm is a
+ // defensive handler should compression ever be plumbed through. A
byte range never
+ // coexists with decompression (compressed files are
non-splittable).
+ //
+ // NOTE: the whole object is read into memory. CometScanRule
restricts native text to
+ // unsplit files (it falls back when Spark splits a file into byte
ranges) and caps the
+ // file size, so in practice this reads each small file once.
Lines are found and
+ // batches built on demand, so only the bytes and the current
batch are held.
+ let buf: Bytes = if file_compression_type.is_compressed() {
+ let mut decoder =
file_compression_type.convert_read(std::io::Cursor::new(raw))?;
+ let mut decoded = Vec::new();
+ decoder.read_to_end(&mut decoded)?;
+ Bytes::from(decoded)
+ } else {
+ raw
+ };
+
+ // Spark's line reader (Hadoop
LineRecordReader.skipUtfByteOrderMark) strips a leading
+ // UTF-8 BOM from the first record at file offset 0, whatever the
separator. wholetext
+ // keeps it. Native reads unsplit files only.
+ let at_file_start = range.as_ref().map(|r| r.start ==
0).unwrap_or(true);
+ // Spark may split an uncompressed file into byte ranges
(TextScan.isSplitable). A line
+ // belongs to the split whose [start, end) contains the line's
start offset -- so each
+ // line is emitted by exactly one split, with no duplication or
gaps across splits. See
+ // DataFusion's CsvOpener for the same rule. wholetext files are
never split.
+ let byte_range = if whole_text {
+ None
+ } else {
+ range.map(|r| (r.start.max(0) as usize, r.end.max(0) as usize))
+ };
+ let reader = TextBatchReader::new(
+ buf,
+ projected_file_schema,
+ batch_size,
+ TextLayout {
+ whole_text,
+ line_sep,
+ strip_bom: !whole_text && at_file_start,
+ byte_range,
+ },
+ );
+
+ Ok(futures::stream::unfold(
+ (reader, baseline_metrics),
+ |(mut reader, metrics)| async move {
+ // The timer covers only building this batch, not the wait
for the next poll.
+ let next = {
+ let _timer = metrics.elapsed_compute().timer();
+ reader.next_batch()
+ };
+ match next {
+ Ok(Some(batch)) => Some((Ok(batch), (reader,
metrics))),
+ Ok(None) => None,
+ Err(e) => {
+ // Report the error once, then end the stream.
+ reader.finished = true;
+ Some((Err(e), (reader, metrics)))
+ }
+ }
+ },
+ )
+ .boxed())
+ }))
+ }
+}
+
+/// How the bytes of one file split into rows.
+struct TextLayout {
+ /// One row for the whole file, separators included.
+ whole_text: bool,
+ /// A custom line separator's exact bytes, or `None` for the universal
newline set.
+ line_sep: Option<Vec<u8>>,
+ /// Strip a leading UTF-8 BOM from the first record (line mode, split at
file offset 0).
+ strip_bom: bool,
+ /// Keep only the lines whose start offset is in `[start, end)`.
+ byte_range: Option<(usize, usize)>,
+}
+
+/// Reads a file's bytes as batches of `value` rows, one batch per call.
+///
+/// Lines are found by scanning forward from `pos`, so nothing is kept per
line: the memory beyond
+/// the file bytes is the batch being built. A file of nothing but newlines
would otherwise need a
+/// slice descriptor per byte.
+struct TextBatchReader {
+ buf: Bytes,
+ schema: SchemaRef,
+ batch_rows: usize,
+ layout: TextLayout,
+ /// Offset of the next unread byte in line mode.
+ pos: usize,
+ /// wholetext: the single row has been returned.
+ whole_taken: bool,
+ /// The first-record BOM check has not run yet.
+ first_record: bool,
+ /// A row read but held back because it would have overflowed the current
batch.
+ pending: Option<(usize, usize)>,
+ finished: bool,
+}
+
+impl TextBatchReader {
+ fn new(buf: Bytes, schema: SchemaRef, batch_size: usize, layout:
TextLayout) -> Self {
+ let first_record = layout.strip_bom;
+ Self {
+ buf,
+ schema,
+ batch_rows: batch_size.max(1),
+ layout,
+ pos: 0,
+ whole_taken: false,
+ first_record,
+ pending: None,
+ finished: false,
+ }
+ }
+
+ /// Offsets `(start, end)` of the next line, terminator excluded, or
`None` at the end.
+ ///
+ /// Default separators are `\n`, `\r\n` and `\r`. A custom separator is an
exact byte match. A
+ /// trailing terminator does not start another, empty line (Hadoop
LineRecordReader).
+ fn scan_line(&mut self) -> Option<(usize, usize)> {
+ let buf = &self.buf[..];
+ if self.layout.whole_text {
+ if self.whole_taken {
+ return None;
+ }
+ self.whole_taken = true;
+ return Some((0, buf.len()));
+ }
+ if self.pos >= buf.len() {
+ return None;
+ }
+ let start = self.pos;
+ match &self.layout.line_sep {
+ None => match buf[start..].iter().position(|&b| b == b'\n' || b ==
b'\r') {
+ None => {
+ self.pos = buf.len();
+ Some((start, buf.len()))
+ }
+ Some(i) => {
+ let end = start + i;
+ self.pos = if buf[end] == b'\r' && buf.get(end + 1) ==
Some(&b'\n') {
+ end + 2
+ } else {
+ end + 1
+ };
+ Some((start, end))
+ }
+ },
+ Some(sep) => match find_subslice(&buf[start..], sep) {
+ None => {
+ self.pos = buf.len();
+ Some((start, buf.len()))
+ }
+ Some(i) => {
+ let end = start + i;
+ self.pos = end + sep.len();
+ Some((start, end))
+ }
+ },
+ }
+ }
+
+ /// Offsets of the next row to emit, after the first-record BOM and the
byte range apply.
+ fn next_row(&mut self) -> Option<(usize, usize)> {
+ if let Some(row) = self.pending.take() {
+ return Some(row);
+ }
+ loop {
+ let (start, end) = self.scan_line()?;
+ let mut row_start = start;
+ // Hadoop reads the first record and then drops the BOM from its
text. When that
+ // leaves nothing and the record was not terminated (the file is
only the BOM), it
+ // counts as end of input, so the file has no rows. A BOM followed
by a separator
+ // leaves an empty row. A custom `lineSep` equal to the BOM never
reaches the strip,
+ // because the BOM is then the separator and the first record does
not start with it.
+ if self.first_record {
+ self.first_record = false;
+ if self.buf[start..end].starts_with(UTF8_BOM) {
+ if self.buf.len() == UTF8_BOM.len() {
+ return None;
+ }
+ row_start = start + UTF8_BOM.len();
+ }
+ }
+ if let Some((range_start, range_end)) = self.layout.byte_range {
+ if start >= range_end {
+ // Lines come in increasing offset order, so the rest are
outside the range.
+ self.pos = self.buf.len();
+ return None;
+ }
+ if start < range_start {
+ continue;
+ }
+ }
+ return Some((row_start, end));
+ }
+ }
+
+ /// The next batch of up to `batch_size` rows, or `None` when the file is
exhausted.
+ fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
+ if self.finished {
+ return Ok(None);
+ }
+ if self.schema.fields().is_empty() {
+ // The count(*) case: rows with no columns. Count lines without
decoding them.
+ let mut rows = 0;
+ while rows < self.batch_rows && self.next_row().is_some() {
+ rows += 1;
+ }
+ if rows == 0 {
+ self.finished = true;
+ return Ok(None);
+ }
+ let options = RecordBatchOptions::new().with_row_count(Some(rows));
+ return Ok(Some(RecordBatch::try_new_with_options(
+ Arc::clone(&self.schema),
+ vec![],
+ &options,
+ )?));
+ }
+
+ // Decode each line once (decode_utf8_spark_lossy borrows for valid
UTF-8, allocates only
+ // for ill-formed input) and accumulate against the DECODED size,
since a lossy decode can
+ // expand ill-formed bytes up to 3x. Start a new batch before the
value buffer would reach
+ // Arrow's 32-bit offset limit. Always take at least one line so
progress is guaranteed;
+ // if a single decoded value alone would overflow, fail cleanly rather
than panic in Arrow.
+ // A handle to the bytes lets the decoded values borrow them while the
reader advances.
+ let buf = self.buf.clone();
+ let mut decoded: Vec<Cow<str>> = Vec::new();
+ let mut data_bytes = 0usize;
+ while decoded.len() < self.batch_rows {
+ let (start, end) = match self.next_row() {
+ Some(row) => row,
+ None => break,
+ };
+ let value = decode_utf8_spark_lossy(&buf[start..end]);
+ let len = value.len();
+ if len > i32::MAX as usize {
+ return
Err(datafusion::common::DataFusionError::Execution(format!(
+ "Text value of {len} bytes exceeds Arrow's 2GB
string-offset limit"
+ )));
+ }
+ if !decoded.is_empty() && data_bytes + len > MAX_BATCH_VALUE_BYTES
{
+ self.pending = Some((start, end));
+ break;
+ }
+ data_bytes += len;
+ decoded.push(value);
+ }
+ if decoded.is_empty() {
+ self.finished = true;
+ return Ok(None);
+ }
+ let mut builder = StringBuilder::with_capacity(decoded.len(),
data_bytes);
+ for value in &decoded {
+ builder.append_value(value);
+ }
+ Ok(Some(RecordBatch::try_new(
+ Arc::clone(&self.schema),
+ vec![Arc::new(builder.finish())],
+ )?))
Review Comment:
Done. next_batch now appends each decoded value straight into the
StringBuilder and tracks the row count and decoded bytes itself, so each batch
is walked once and there is no intermediate Vec. The overflow check uses the
length of the next decoded value. The builder is sized from the batch size,
capped by the file size.
##########
native/core/src/execution/operators/text_scan.rs:
##########
@@ -0,0 +1,796 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Native Text file scan, mirroring Spark's `text` data source.
+//!
+//! Spark's text reader produces a single `value: string` column. With the
default options each
+//! line of the file is one row (split on `\n`, `\r\n`, or `\r`, with the
terminator stripped).
+//! With `wholetext=true` the whole file is a single row. A custom `lineSep`
splits on exactly that
+//! byte sequence instead of the universal-newline set. See the
`TextOptions`/`HadoopFileLinesReader`
+//! logic in Spark's `org.apache.spark.sql.execution.datasources.text` package.
+
+use crate::execution::operators::ExecutionError;
+use arrow::array::StringBuilder;
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::{RecordBatch, RecordBatchOptions};
+use bytes::Bytes;
+use datafusion::common::Result;
+use datafusion::datasource::object_store::ObjectStoreUrl;
+use datafusion::physical_expr::projection::ProjectionExprs;
+use datafusion::physical_plan::metrics::{BaselineMetrics,
ExecutionPlanMetricsSet, MetricBuilder};
+use datafusion::physical_plan::DisplayFormatType;
+use datafusion_comet_common::decode_utf8_spark_lossy;
+use datafusion_comet_proto::spark_operator::TextOptions;
+use datafusion_datasource::file::FileSource;
+use datafusion_datasource::file_compression_type::FileCompressionType;
+use datafusion_datasource::file_groups::FileGroup;
+use datafusion_datasource::file_scan_config::{FileScanConfig,
FileScanConfigBuilder};
+use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
+use datafusion_datasource::projection::{ProjectionOpener, SplitProjection};
+use datafusion_datasource::source::DataSourceExec;
+use datafusion_datasource::{as_file_source, PartitionedFile, TableSchema};
+use futures::StreamExt;
+use object_store::{ObjectStore, ObjectStoreExt};
+use std::borrow::Cow;
+use std::fmt;
+use std::io::Read;
+use std::sync::Arc;
+
+/// Batch size used when the execution framework does not set one before
opening the file.
+const DEFAULT_TEXT_BATCH_SIZE: usize = 8192;
+
+/// UTF-8 byte-order mark. Spark's line reader strips a leading one; wholetext
keeps it.
+const UTF8_BOM: &[u8] = &[0xEF, 0xBB, 0xBF];
+
+/// Cap on a batch's decoded `value` buffer bytes, kept well below Arrow's 2GB
32-bit offset limit
+/// so a batch of many rows (or a few long lines) cannot overflow. Measured
against the decoded
+/// string length, since a lossy UTF-8 decode can expand ill-formed bytes up
to 3x.
+const MAX_BATCH_VALUE_BYTES: usize = 1 << 30;
+
+pub fn init_text_datasource_exec(
+ object_store_url: ObjectStoreUrl,
+ file_groups: Vec<Vec<PartitionedFile>>,
+ data_schema: SchemaRef,
+ _partition_schema: SchemaRef,
+ projection_vector: Vec<usize>,
+ text_options: &TextOptions,
+) -> Result<Arc<DataSourceExec>, ExecutionError> {
+ let text_source = TextSource::new(data_schema, text_options);
+
+ let file_groups = file_groups.into_iter().map(FileGroup::new).collect();
+
+ let file_scan_config = FileScanConfigBuilder::new(object_store_url,
Arc::new(text_source))
+ .with_file_groups(file_groups)
+ .with_projection_indices(Some(projection_vector))?
+ .build();
+
+ Ok(DataSourceExec::from_data_source(file_scan_config))
+}
+
+/// A [`FileSource`] for Spark's text format. The file schema is always the
single `value` column;
+/// projection (and partition columns, if any) is handled by
[`ProjectionOpener`], exactly as the
+/// CSV source does.
+#[derive(Debug, Clone)]
+struct TextSource {
+ table_schema: TableSchema,
+ projection: SplitProjection,
+ batch_size: Option<usize>,
+ metrics: ExecutionPlanMetricsSet,
+ whole_text: bool,
+ // The raw separator bytes when a custom `lineSep` is set; None means
universal newline
+ // (`\n`, `\r\n`, `\r`).
+ line_sep: Option<Vec<u8>>,
+}
+
+impl TextSource {
+ fn new(file_schema: SchemaRef, options: &TextOptions) -> Self {
+ let table_schema: TableSchema = file_schema.into();
+ // `line_sep` arrives as the exact separator bytes Spark computed
(`lineSeparatorInRead`,
+ // already encoded with the file's charset), so we split on them
directly.
+ let line_sep = options.line_sep.clone().filter(|sep| !sep.is_empty());
+ Self {
+ projection: SplitProjection::unprojected(&table_schema),
+ table_schema,
+ batch_size: None,
+ metrics: ExecutionPlanMetricsSet::new(),
+ whole_text: options.whole_text,
+ line_sep,
+ }
+ }
+}
+
+impl From<TextSource> for Arc<dyn FileSource> {
+ fn from(source: TextSource) -> Self {
+ as_file_source(source)
+ }
+}
+
+impl FileSource for TextSource {
+ fn create_file_opener(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ base_config: &FileScanConfig,
+ partition_index: usize,
+ ) -> Result<Arc<dyn FileOpener>> {
+ // The inner opener must emit batches whose schema is the file schema
projected by
+ // `file_indices` (that is what ProjectionOpener feeds its projector).
For text that is
+ // either the single `value` column or, for a projection-less count,
an empty schema.
+ let file_schema = self.table_schema.file_schema();
+ let projected_file_schema =
Arc::new(file_schema.project(&self.projection.file_indices)?);
+ let opener = Arc::new(TextOpener {
+ projected_file_schema,
+ file_compression_type: base_config.file_compression_type,
+ object_store,
+ partition_index,
+ batch_size: self.batch_size.unwrap_or(DEFAULT_TEXT_BATCH_SIZE),
+ metrics: self.metrics.clone(),
+ whole_text: self.whole_text,
+ line_sep: self.line_sep.clone(),
+ }) as Arc<dyn FileOpener>;
+ ProjectionOpener::try_new(self.projection.clone(), opener, file_schema)
+ }
+
+ fn table_schema(&self) -> &TableSchema {
+ &self.table_schema
+ }
+
+ fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
+ let mut conf = self.clone();
+ conf.batch_size = Some(batch_size);
+ Arc::new(conf)
+ }
+
+ fn try_pushdown_projection(
+ &self,
+ projection: &ProjectionExprs,
+ ) -> Result<Option<Arc<dyn FileSource>>> {
+ let mut source = self.clone();
+ let new_projection = self.projection.source.try_merge(projection)?;
+ source.projection =
SplitProjection::new(self.table_schema.file_schema(), &new_projection);
+ Ok(Some(Arc::new(source)))
+ }
+
+ fn projection(&self) -> Option<&ProjectionExprs> {
+ Some(&self.projection.source)
+ }
+
+ fn metrics(&self) -> &ExecutionPlanMetricsSet {
+ &self.metrics
+ }
+
+ fn file_type(&self) -> &str {
+ "text"
+ }
+
+ // A custom line separator can straddle a byte-range boundary, and we read
whole files rather
+ // than boundary-aligned ranges, so this source must not let DataFusion
repartition files by
+ // byte range.
+ fn supports_repartitioning(&self) -> bool {
+ false
+ }
+
+ fn fmt_extra(&self, t: DisplayFormatType, f: &mut fmt::Formatter) ->
fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, ", whole_text={}", self.whole_text)
+ }
+ DisplayFormatType::TreeRender => Ok(()),
+ }
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(
+ &Arc<dyn datafusion::physical_plan::PhysicalExpr>,
+ ) -> Result<datafusion::common::tree_node::TreeNodeRecursion>,
+ ) -> Result<datafusion::common::tree_node::TreeNodeRecursion> {
+
datafusion::physical_plan::apply_expression_roots(self.projection.source.iter(),
f)
+ }
+}
+
+/// A [`FileOpener`] that reads a whole text file and yields the projected
file schema in batches.
+struct TextOpener {
+ // The file schema projected by `file_indices`: the single `value` column,
or an empty schema
+ // for a projection-less count.
+ projected_file_schema: SchemaRef,
+ file_compression_type: FileCompressionType,
+ object_store: Arc<dyn ObjectStore>,
+ partition_index: usize,
+ batch_size: usize,
+ metrics: ExecutionPlanMetricsSet,
+ whole_text: bool,
+ line_sep: Option<Vec<u8>>,
+}
+
+impl FileOpener for TextOpener {
+ fn open(&self, partitioned_file: PartitionedFile) ->
Result<FileOpenFuture> {
+ let store = Arc::clone(&self.object_store);
+ let projected_file_schema = Arc::clone(&self.projected_file_schema);
+ let file_compression_type = self.file_compression_type;
+ let batch_size = self.batch_size;
+ let whole_text = self.whole_text;
+ let line_sep = self.line_sep.clone();
+ let baseline_metrics = BaselineMetrics::new(&self.metrics,
self.partition_index);
+ // Feeds the task's input bytes in the Spark UI (see
CometMetricNode.reportScanInputMetrics).
+ let bytes_scanned =
+ MetricBuilder::new(&self.metrics).counter("bytes_scanned",
self.partition_index);
+
+ Ok(Box::pin(async move {
+ let range = partitioned_file.range.clone();
+ let location = partitioned_file.object_meta.location;
+ let raw = store.get(&location).await?.bytes().await?;
+ bytes_scanned.add(raw.len());
+
+ // Decompress only when the file is actually compressed. On
Comet's plan config the file
+ // is always UNCOMPRESSED (CometScanRule falls back for compressed
files), so we keep
+ // the fetched bytes as they are and avoid a full-file copy; the
decompress arm is a
+ // defensive handler should compression ever be plumbed through. A
byte range never
+ // coexists with decompression (compressed files are
non-splittable).
+ //
+ // NOTE: the whole object is read into memory. CometScanRule
restricts native text to
+ // unsplit files (it falls back when Spark splits a file into byte
ranges) and caps the
+ // file size, so in practice this reads each small file once.
Lines are found and
+ // batches built on demand, so only the bytes and the current
batch are held.
+ let buf: Bytes = if file_compression_type.is_compressed() {
+ let mut decoder =
file_compression_type.convert_read(std::io::Cursor::new(raw))?;
+ let mut decoded = Vec::new();
+ decoder.read_to_end(&mut decoded)?;
+ Bytes::from(decoded)
+ } else {
+ raw
+ };
+
+ // Spark's line reader (Hadoop
LineRecordReader.skipUtfByteOrderMark) strips a leading
+ // UTF-8 BOM from the first record at file offset 0, whatever the
separator. wholetext
+ // keeps it. Native reads unsplit files only.
+ let at_file_start = range.as_ref().map(|r| r.start ==
0).unwrap_or(true);
+ // Spark may split an uncompressed file into byte ranges
(TextScan.isSplitable). A line
+ // belongs to the split whose [start, end) contains the line's
start offset -- so each
+ // line is emitted by exactly one split, with no duplication or
gaps across splits. See
+ // DataFusion's CsvOpener for the same rule. wholetext files are
never split.
+ let byte_range = if whole_text {
+ None
+ } else {
+ range.map(|r| (r.start.max(0) as usize, r.end.max(0) as usize))
+ };
Review Comment:
Done. Removed the decompress arm, file_compression_type, byte_range, the
offset-0 flag, the std::io::Read import and the two range_split_* tests. The
Scala gate declines compressed and split files, so none of it could run. Good
catch on the boundary rule: Hadoop gives a line that starts exactly at a split
boundary to the earlier split, so the filter would have been wrong anyway.
##########
spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala:
##########
@@ -477,6 +419,186 @@ case class CometScanRule(session: SparkSession)
withFallbackReasons(scanExec, fallbackReasons.toSet)
}
+ case scan: TextScan if COMET_TEXT_V2_NATIVE_ENABLED.get() =>
+ if (scanExec.output.exists(_.isMetadataCol)) {
+ return withFallbackReason(
+ scanExec,
+ "Metadata columns are not supported for Text V2 scans")
+ }
+ // A native scan is useless without native execution: with exec
disabled, CometExecRule
+ // leaves this CometBatchScanExec unconverted and it fails at runtime.
Require native exec.
+ if (!COMET_EXEC_ENABLED.get()) {
+ return withFallbackReason(
+ scanExec,
+ s"Native Text scan requires ${COMET_EXEC_ENABLED.key} to be
enabled")
+ }
+ // input_file_name/_block_start/_block_length are already handled by
the blanket V2 gate
+ // near the top of transformV2Scan, so no Text-specific check is
needed here.
+ val fallbackReasons = new ListBuffer[String]()
+ val schemaSupported =
+ CometBatchScanExec.isSchemaSupported(scan.readDataSchema,
fallbackReasons)
+ if (!schemaSupported) {
+ fallbackReasons += s"Schema ${scan.readDataSchema} is not supported"
+ }
+ // Spark's text source allows only a single data column
(TextScan.verifyReadSchema). A
+ // user-forced multi-column read schema is invalid; fall back so Spark
raises its structured
+ // AnalysisException rather than letting the native reader fail with
an opaque error.
+ if (scan.readDataSchema.size > 1) {
+ fallbackReasons += "Text data source supports only a single column"
+ }
+ // Partition columns are not implemented for native Text scans (unlike
the native CSV path,
+ // which forwards them). Fall back to Spark for partitioned reads
rather than drop them.
+ if (scan.readPartitionSchema.nonEmpty) {
+ fallbackReasons += "Comet does not support partition columns in
native Text scans"
+ }
+ // The native reader loads each file whole. Spark splits a large
uncompressed text file
+ // into byte-range partitions (TextScan.isSplitable), and each split
would then re-read and
+ // buffer the whole file with no memory-pool reservation. Restrict
native execution to
+ // unsplit files (the small build-side lookups this targets); Spark
handles large files.
+ val hasSplitFile = scanExec.inputPartitions.exists {
+ case fp: FilePartition => fp.files.exists(_.start != 0)
+ case _ => false
+ }
+ if (hasSplitFile) {
+ fallbackReasons += "Comet native Text scan does not support files
split into byte " +
+ "ranges (large files)"
+ }
+ // The native reader loads each file whole and does not reserve the
buffer against the
+ // memory pool, so bound the per-file size; larger files fall back to
Spark. This keeps the
+ // feature to the small lookup tables it targets regardless of
maxPartitionBytes.
+ val maxTextFileSize = COMET_SCAN_TEXT_MAX_FILE_SIZE.get()
+ val hasOversizedFile = scanExec.inputPartitions.exists {
+ case fp: FilePartition => fp.files.exists(_.length > maxTextFileSize)
+ case _ => false
+ }
+ if (hasOversizedFile) {
+ fallbackReasons += "Comet native Text scan does not read files
larger than " +
+ s"${COMET_SCAN_TEXT_MAX_FILE_SIZE.key} ($maxTextFileSize bytes)"
+ }
Review Comment:
Done. selectedFiles is built once at the top of the Text branch, with a
plain FilePartition cast like the serde uses, and the split, size, compression,
checksum and path checks all use it.
##########
spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala:
##########
@@ -477,6 +419,186 @@ case class CometScanRule(session: SparkSession)
withFallbackReasons(scanExec, fallbackReasons.toSet)
}
+ case scan: TextScan if COMET_TEXT_V2_NATIVE_ENABLED.get() =>
+ if (scanExec.output.exists(_.isMetadataCol)) {
+ return withFallbackReason(
+ scanExec,
+ "Metadata columns are not supported for Text V2 scans")
+ }
+ // A native scan is useless without native execution: with exec
disabled, CometExecRule
+ // leaves this CometBatchScanExec unconverted and it fails at runtime.
Require native exec.
+ if (!COMET_EXEC_ENABLED.get()) {
+ return withFallbackReason(
+ scanExec,
+ s"Native Text scan requires ${COMET_EXEC_ENABLED.key} to be
enabled")
+ }
+ // input_file_name/_block_start/_block_length are already handled by
the blanket V2 gate
+ // near the top of transformV2Scan, so no Text-specific check is
needed here.
+ val fallbackReasons = new ListBuffer[String]()
+ val schemaSupported =
+ CometBatchScanExec.isSchemaSupported(scan.readDataSchema,
fallbackReasons)
+ if (!schemaSupported) {
+ fallbackReasons += s"Schema ${scan.readDataSchema} is not supported"
+ }
+ // Spark's text source allows only a single data column
(TextScan.verifyReadSchema). A
+ // user-forced multi-column read schema is invalid; fall back so Spark
raises its structured
+ // AnalysisException rather than letting the native reader fail with
an opaque error.
+ if (scan.readDataSchema.size > 1) {
+ fallbackReasons += "Text data source supports only a single column"
+ }
+ // Partition columns are not implemented for native Text scans (unlike
the native CSV path,
+ // which forwards them). Fall back to Spark for partitioned reads
rather than drop them.
+ if (scan.readPartitionSchema.nonEmpty) {
+ fallbackReasons += "Comet does not support partition columns in
native Text scans"
+ }
+ // The native reader loads each file whole. Spark splits a large
uncompressed text file
+ // into byte-range partitions (TextScan.isSplitable), and each split
would then re-read and
+ // buffer the whole file with no memory-pool reservation. Restrict
native execution to
+ // unsplit files (the small build-side lookups this targets); Spark
handles large files.
+ val hasSplitFile = scanExec.inputPartitions.exists {
+ case fp: FilePartition => fp.files.exists(_.start != 0)
+ case _ => false
+ }
+ if (hasSplitFile) {
+ fallbackReasons += "Comet native Text scan does not support files
split into byte " +
+ "ranges (large files)"
+ }
+ // The native reader loads each file whole and does not reserve the
buffer against the
+ // memory pool, so bound the per-file size; larger files fall back to
Spark. This keeps the
+ // feature to the small lookup tables it targets regardless of
maxPartitionBytes.
+ val maxTextFileSize = COMET_SCAN_TEXT_MAX_FILE_SIZE.get()
+ val hasOversizedFile = scanExec.inputPartitions.exists {
+ case fp: FilePartition => fp.files.exists(_.length > maxTextFileSize)
+ case _ => false
+ }
+ if (hasOversizedFile) {
+ fallbackReasons += "Comet native Text scan does not read files
larger than " +
+ s"${COMET_SCAN_TEXT_MAX_FILE_SIZE.key} ($maxTextFileSize bytes)"
+ }
+ // Build the Hadoop conf from the scan's read options with their case
preserved, like
+ // Spark's TextScan.createReaderFactory (Hadoop keys are case
sensitive, but the scan's
+ // option map lowercases them).
+ val hadoopConf = session.sessionState.newHadoopConfWithOptions(
+ scan.options.asCaseSensitiveMap.asScala.toMap)
+ val textOptions =
+ new org.apache.spark.sql.execution.datasources.text.TextOptions(
+ scan.options.asScala.toMap)
+ // The files this scan reads. `fileIndex.inputFiles` also lists files
that partition
+ // pruning dropped, so the per-file checks below use the selected
partitions instead.
+ val selectedFiles = scanExec.inputPartitions.flatMap {
+ case fp: FilePartition => fp.files.toSeq
+ case _ => Seq.empty
+ }
+ // Comet's native plan does not carry a compression codec, so the
native reader would feed
+ // raw compressed bytes to the line splitter. Spark auto-detects a
codec by file extension,
+ // from the decoded path (`Path(URI)`), so a name like `a b.gz` keeps
its suffix. The
+ // version-aware ShimFileFormat.compressedFileCheck also catches the
non-standard
+ // .gzip/.zstd extensions on Spark 4.1+ that CompressionCodecFactory
alone misses.
+ val isCompressed = ShimFileFormat.compressedFileCheck(hadoopConf)
+ val hasCompressedFile =
+ selectedFiles.exists(file => isCompressed(new
org.apache.hadoop.fs.Path(file.pathUri)))
+ if (hasCompressedFile) {
+ fallbackReasons += "Comet does not support compressed text files"
+ }
+ // Hadoop's local filesystem verifies a `.<name>.crc` checksum file
when one exists and
+ // raises a ChecksumException if the data no longer matches it. The
native reader reads the
+ // file directly and would return the altered rows, so fall back for
such files.
+ if (selectedFiles.exists(file =>
+ CometScanRule.hasLocalChecksumFile(hadoopConf, file.pathUri))) {
+ fallbackReasons += "Comet native Text scan does not support local
files that have a " +
+ "Hadoop checksum (.crc) file"
+ }
+ // Spark skips unreadable/missing files when these are set; the native
reader's DataFusion
+ // FileStream fails the whole scan instead (OnError::Fail). Fall back
to match Spark.
+ if (SQLConf.get.ignoreCorruptFiles || scan.options.getBoolean(
+ "ignoreCorruptFiles",
+ false)) {
+ fallbackReasons += "Comet native Text scan does not support
ignoreCorruptFiles"
+ }
+ if (SQLConf.get.ignoreMissingFiles || scan.options.getBoolean(
+ "ignoreMissingFiles",
+ false)) {
+ fallbackReasons += "Comet native Text scan does not support
ignoreMissingFiles"
+ }
+ // Only Spark's line reader consults the settings below. The wholetext
reader does not,
+ // so a value it never reads must not fail or fall back a wholetext
query.
+ if (!textOptions.wholeText) {
+ // Spark's line reader skips any line at or over
+ // mapreduce.input.linerecordreader.line.maxlength (default
unlimited); the key can come
+ // from spark.hadoop.*, core-site.xml, or a read option. The native
reader does not
+ // truncate, so fall back when a finite limit is set. A value that
is not an integer
+ // fails Spark's reader, so fall back and let Spark report it.
+ val maxLineLength = scala.util.Try(
+
hadoopConf.getInt("mapreduce.input.linerecordreader.line.maxlength",
Int.MaxValue))
+ if (maxLineLength.map(_ != Int.MaxValue).getOrElse(true)) {
+ fallbackReasons += "Comet native Text scan does not support " +
+ "mapreduce.input.linerecordreader.line.maxlength"
+ }
+ // Spark passes `mapreduce.job.input.file.must.*` to the
filesystem's open call as
+ // mandatory options, and the read fails if the filesystem cannot
honor one. The native
+ // reader opens files itself and cannot enforce them, so fall back
and let Spark decide.
+ if (hadoopConf
+ .iterator()
+ .asScala
+ .exists(_.getKey.startsWith("mapreduce.job.input.file.must."))) {
+ fallbackReasons += "Comet native Text scan does not support
mandatory file-open " +
+ "options (mapreduce.job.input.file.must.*)"
+ }
+ // Spark reads a custom line separator with Hadoop's
LineReader.readCustomLine, which on
+ // a mismatch backtracks only one byte instead of following the
separator's own prefix
+ // overlap. For a separator that overlaps itself (e.g. "aab", whose
"aa" prefix can
+ // restart a match) that reader can split differently from the
native reader's exact
+ // substring scan, and the difference depends on where the separator
falls across
+ // Hadoop's input buffer boundaries. Fall back so the two always
agree.
+ textOptions.lineSeparatorInRead.foreach { sep =>
+ if (CometScanRule.separatorSelfOverlaps(sep)) {
+ fallbackReasons += "Comet native Text scan does not support a
self-overlapping " +
+ "line separator (Spark's line reader can split it differently)"
+ }
+ }
+ }
+ // A wholetext file is never split (TextScan.isSplitable is false for
wholeText), so the
+ // split-file guard above does not bound it. Reading a very large file
whole is unbounded
+ // memory, and the whole file becomes one Arrow string value. Because
a lossy UTF-8 decode
+ // can expand ill-formed bytes up to 3x, cap the raw file size at
Int.MaxValue/3 so the
+ // decoded value stays under Arrow's 2GB 32-bit string-offset limit;
fall back above that.
+ // (The native reader also fails cleanly rather than panicking if a
value still overflows.)
+ if (textOptions.wholeText) {
+ val maxWholeTextBytes =
+ math.min(session.sessionState.conf.filesMaxPartitionBytes,
Int.MaxValue.toLong / 3)
+ val hasLargeWholeTextFile = scanExec.inputPartitions.exists {
+ case fp: FilePartition => fp.files.exists(_.length >
maxWholeTextBytes)
+ case _ => false
+ }
+ if (hasLargeWholeTextFile) {
+ fallbackReasons += "Comet native Text scan does not support large
wholetext files " +
+ s"(over ${SQLConf.FILES_MAX_PARTITION_BYTES.key}, capped below
2GB)"
+ }
+ }
Review Comment:
Done. The wholetext cap is now one check,
min(spark.comet.scan.text.maxFileSize, Int.MaxValue/3), so lowering
maxPartitionBytes no longer turns off native wholetext. I replaced the old test
with two: one that wholetext falls back above the size ceiling, and one that
maxPartitionBytes=16 leaves a 100-byte wholetext file native.
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometTextNativeScanExec.scala:
##########
@@ -0,0 +1,196 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder}
+import org.apache.spark.sql.catalyst.plans.QueryPlan
+import org.apache.spark.sql.catalyst.plans.physical.{Partitioning,
UnknownPartitioning}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.text.TextOptions
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.execution.datasources.v2.text.TextScan
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+
+import com.google.common.base.Objects
+
+import org.apache.comet.{CometConf, ConfigEntry}
+import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
+import org.apache.comet.objectstore.NativeConfig
+import org.apache.comet.serde.{CometOperatorSerde, OperatorOuterClass}
+import org.apache.comet.serde.OperatorOuterClass.Operator
+import org.apache.comet.serde.operator.{partition2Proto, schema2Proto}
+
+/*
+ * Native Text scan operator that delegates file reading to datafusion.
Mirrors the native CSV
+ * scan path; produces a single `value: string` column.
+ */
+case class CometTextNativeScanExec(
+ override val nativeOp: Operator,
+ override val output: Seq[Attribute],
+ @transient override val originalPlan: BatchScanExec,
+ override val serializedPlanOpt: SerializedPlan)
+ extends CometLeafExec {
+ override val supportsColumnar: Boolean = true
+
+ override val nodeName: String = "CometTextNativeScan"
+
+ // Show only the semantic output in EXPLAIN / the SQL tab / the event log.
The default would print
+ // the whole `nativeOp` protobuf, which carries every input file path (large
for a lookup
+ // directory of many files). Mirrors CometNativeScanExec.
+ override def stringArgs: Iterator[Any] = Iterator(output)
+
+ // The serialized partitions, which split any Spark partition that mixes
object stores.
+ override def outputPartitioning: Partitioning = UnknownPartitioning(
+ nativeOp.getTextScan.getFilePartitionsCount)
+
+ override def outputOrdering: Seq[SortOrder] = Nil
+
+ // The scan reports the bytes it fetched as `bytes_scanned`, which
CometMetricNode turns into the
+ // task's input bytes and records, so they show on the stage UI like any
other scan.
+ override lazy val metrics: Map[String, SQLMetric] =
+ CometMetricNode.baselineMetrics(sparkContext) +
+ ("bytes_scanned" -> SQLMetrics.createSizeMetric(sparkContext, "number of
bytes scanned"))
+
+ // Two scans of the same files must compare equal so Spark reuses their
exchanges and stages.
+ // That needs normalized output attribute ids and the wrapped scan's
canonical form, and no
+ // serialization state, which differs between otherwise identical copies.
+ override protected def doCanonicalize(): SparkPlan = {
+ CometTextNativeScanExec(
+ nativeOp,
+ output.map(QueryPlan.normalizeExpressions(_, output)),
+ if (originalPlan == null) null else originalPlan.doCanonicalize(),
+ SerializedPlan(None))
+ }
+
+ override def equals(obj: Any): Boolean = {
+ obj match {
+ case other: CometTextNativeScanExec =>
+ this.output == other.output &&
+ this.serializedPlanOpt == other.serializedPlanOpt &&
+ this.originalPlan == other.originalPlan
+ case _ =>
+ false
+ }
+ }
+
+ override def hashCode(): Int = {
+ Objects.hashCode(output, serializedPlanOpt, originalPlan)
+ }
+}
+
+object CometTextNativeScanExec extends CometOperatorSerde[CometBatchScanExec] {
+
+ override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
+ CometConf.COMET_TEXT_V2_NATIVE_ENABLED)
+
+ override def convert(
+ op: CometBatchScanExec,
+ builder: Operator.Builder,
+ childOp: Operator*): Option[Operator] = {
+ val textScanBuilder = OperatorOuterClass.TextScan.newBuilder()
+ val textScan = op.wrapped.scan.asInstanceOf[TextScan]
+ val sessionState = op.session.sessionState
+ val options = new TextOptions(textScan.options.asScala.toMap)
+ // Build the Hadoop conf from the scan's read options (case preserved),
matching Spark's
+ // TextScan.createReaderFactory, so per-read settings like fs.s3a.endpoint
reach native.
+ val hadoopConf =
+
sessionState.newHadoopConfWithOptions(textScan.options.asCaseSensitiveMap.asScala.toMap)
+ val s3CompliantSchemes = NativeConfig.resolveS3CompliantSchemes(hadoopConf)
+ // The scan gate already declined schemes that only the read options route
through libhdfs.
+ val libhdfsSchemes = NativeConfig
+ .resolveScanLibhdfsSchemes(
+ NativeConfig.resolveLibhdfsSchemes(sessionState.newHadoopConf()),
+ hadoopConf)
+ .libhdfs
+ val inputPartitions = op.inputPartitions.map(_.asInstanceOf[FilePartition])
+ val multiStoreReason = CometScanUtils.multiStoreFallbackReason(
+ "Native Text scan",
+ inputPartitions.view.flatMap(_.files.view.map(_.pathUri)),
+ s3CompliantSchemes,
+ libhdfsSchemes,
+ isBucketedScan = false)
+ if (multiStoreReason.nonEmpty) {
+ // CometExecRule falls back to the wrapped scan, so tag it too for the
explain output.
+ withFallbackReason(op, multiStoreReason.get)
+ withFallbackReason(op.wrapped, multiStoreReason.get)
+ return None
+ }
+ // The options of every scheme the scan reads, or a decline for the ABFS
layouts one
+ // resolution cannot serve. Checked before the split below so that its
reason is reported.
+ val fileUris = inputPartitions.view.flatMap(_.files).map(_.pathUri)
+ val objectStoreOptions =
+ NativeConfig.extractScanObjectStoreOptions(hadoopConf, fileUris) match {
+ case Right(options) => options
+ case Left(reason) =>
+ withFallbackReason(op, reason)
+ withFallbackReason(op.wrapped, reason)
+ return None
+ }
+ // Native planning reads each partition through one object store.
+ val filePartitions = CometScanUtils.splitPartitionsByStore(
+ inputPartitions,
+ file => NativeConfig.objectStoreKey(file.pathUri, s3CompliantSchemes,
libhdfsSchemes))
+ // Spark 3.4 reports a one-partition scan as SinglePartition and plans no
exchange above it,
+ // so the operators above rely on the scan's partition count.
+ if (filePartitions.length != inputPartitions.length &&
+ !op.wrapped.outputPartitioning.isInstanceOf[UnknownPartitioning]) {
+ val reason = "Native Text scan would split a partition that mixes object
stores, but " +
+ "Spark planned the operators above it for the scan's own partitioning"
+ withFallbackReason(op, reason)
+ withFallbackReason(op.wrapped, reason)
+ return None
+ }
+ val textOptionsProto = textOptions2Proto(options)
+ val dataSchemaProto = schema2Proto(textScan.dataSchema)
+ val readSchemaFieldNames = textScan.readDataSchema.fieldNames
+ val projectionVector = textScan.dataSchema.zipWithIndex
+ .filter { case (field, _) =>
+ readSchemaFieldNames.contains(field.name)
+ }
+ .map(_._2.asInstanceOf[Integer])
+ val partitionSchemaProto = schema2Proto(textScan.readPartitionSchema)
+ val partitionsProto = filePartitions.map(partition2Proto(_,
textScan.readPartitionSchema))
+
+ textScanBuilder.putAllObjectStoreOptions(objectStoreOptions.asJava)
+ textScanBuilder.setTextOptions(textOptionsProto)
+ textScanBuilder.addAllFilePartitions(partitionsProto.asJava)
+ textScanBuilder.addAllDataSchema(dataSchemaProto.asJava)
+ textScanBuilder.addAllProjectionVector(projectionVector.asJava)
+ textScanBuilder.addAllPartitionSchema(partitionSchemaProto.asJava)
+ Some(builder.setTextScan(textScanBuilder).build())
Review Comment:
Done. The format-independent part of the two convert methods (Hadoop conf
from the read options, libhdfs, multi-store and ABFS checks, per-store split,
Spark 3.4 partition-count check, projection vector) is now
CometScanUtils.prepareNativeFileScan and projectionVector. Both nodes decline
reader-only libhdfs schemes the same way now, which fixes the drift you pointed
out. The CSV node also has the canonical form the Text node has, with a test
that two CSV scans of the same file are sameResult. I left equals, hashCode and
stringArgs in each case class.
--
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]
