sunchao commented on code in PR #6784: URL: https://github.com/apache/datafusion-comet/pull/6784#discussion_r4433048103
########## native/core/src/execution/operators/text_scan.rs: ########## @@ -0,0 +1,686 @@ +// 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 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 borrow + // the fetched bytes directly 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), so in + // practice this reads each small file once; the range filter below is a correctness + // safety net for any split that still reaches here. + let buf: Cow<[u8]> = 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)?; + Cow::Owned(decoded) + } else { + Cow::Borrowed(raw.as_ref()) + }; + + let mut timer = baseline_metrics.elapsed_compute().timer(); + // Spark's line reader (Hadoop LineRecordReader.skipUtfByteOrderMark) strips a leading + // UTF-8 BOM from the first record at file offset 0, whatever the separator. It does so + // after finding the first record, so split first and then strip. Native reads unsplit + // files only. wholetext keeps the BOM. + let at_file_start = range.as_ref().map(|r| r.start == 0).unwrap_or(true); + let content: &[u8] = &buf; + let mut lines = split_lines(content, whole_text, line_sep.as_deref()); Review Comment: A file consisting of newlines allocates one 16-byte slice descriptor per input byte before the first batch is returned. An isolated 8 MiB input used 128 MiB for the slice table alone. At the default 64 MiB file cap the table reaches 1 GiB, including for `count(*)`. Could the reader split lines incrementally and retain only what the current output batch needs? ########## 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.isCompressedFile also catches the non-standard .gzip/.zstd + // extensions on Spark 4.1+ that CompressionCodecFactory alone misses. + val hasCompressedFile = selectedFiles.exists { file => Review Comment: On Spark 3.4, 3.5 and 4.0, the shim constructs a Hadoop `CompressionCodecFactory` for each path. For a directory of uncompressed files, this loop rebuilds the same codec registry for every file on the driver. Could this plan check create the factory once and reuse it across the selected paths? -- 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]
