sunchao commented on code in PR #6784: URL: https://github.com/apache/datafusion-comet/pull/6784#discussion_r4253018828
########## native/core/src/execution/operators/text_scan.rs: ########## @@ -0,0 +1,642 @@ +// 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}; +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); + + 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?; + + // 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 strips a leading UTF-8 BOM at the start of the file, but only with + // the default separator. With a custom lineSep the BOM is ordinary data (Spark leaves it + // for the splitter, so e.g. a file `BOM a BOM b` read with lineSep=BOM yields + // ["", "a", "b"]). Only strip in line mode, at file offset 0 (native reads unsplit files + // only), and when no custom lineSep is set. wholetext keeps the BOM. + let at_file_start = range.as_ref().map(|r| r.start == 0).unwrap_or(true); + let content = strip_leading_bom(&buf, whole_text, at_file_start, line_sep.is_some()); + let mut lines = split_lines(content, whole_text, line_sep.as_deref()); + // 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. + if !whole_text { + if let Some(range) = range { + lines = filter_lines_by_range(content, lines, range.start, range.end); + } + } + let batches = build_batches(&projected_file_schema, lines, batch_size)?; + timer.stop(); + + Ok(futures::stream::iter(batches.into_iter().map(Ok)).boxed()) + })) + } +} + +/// Strip a single leading UTF-8 BOM the way Spark's default line reader does: only in line mode +/// (not wholetext), only for the split that starts at file offset 0, and only with the default +/// separator. With a custom lineSep the BOM is left in place as ordinary data. +fn strip_leading_bom( + buf: &[u8], + whole_text: bool, + at_file_start: bool, + custom_line_sep: bool, +) -> &[u8] { + if !whole_text && at_file_start && !custom_line_sep && buf.starts_with(UTF8_BOM) { + &buf[UTF8_BOM.len()..] + } else { + buf + } +} + +/// Split a decompressed file buffer into text lines following Spark's semantics. +fn split_lines<'a>(buf: &'a [u8], whole_text: bool, line_sep: Option<&[u8]>) -> Vec<&'a [u8]> { + if whole_text { + // One row per file; the whole content (including internal separators) is a single value. + return vec![buf]; + } + match line_sep { + Some(sep) => split_on_separator(buf, sep), + None => split_universal(buf), + } +} + +/// Universal-newline splitting: a line ends at `\n`, `\r\n`, or `\r`, and the terminator is +/// stripped. A trailing terminator does not produce a final empty line (Hadoop LineRecordReader). +fn split_universal(buf: &[u8]) -> Vec<&[u8]> { + let mut lines = Vec::new(); + let mut start = 0; + let mut i = 0; + while i < buf.len() { + match buf[i] { + b'\n' => { + lines.push(&buf[start..i]); + i += 1; + start = i; + } + b'\r' => { + lines.push(&buf[start..i]); + i += 1; + if i < buf.len() && buf[i] == b'\n' { + i += 1; + } + start = i; + } + _ => i += 1, + } + } + if start < buf.len() { + lines.push(&buf[start..]); + } + lines +} + +/// Split on an exact separator byte sequence; the trailing separator does not produce an empty +/// final line. +fn split_on_separator<'a>(buf: &'a [u8], sep: &[u8]) -> Vec<&'a [u8]> { + let mut lines = Vec::new(); + let mut start = 0; + let mut i = 0; + while i + sep.len() <= buf.len() { + if &buf[i..i + sep.len()] == sep { Review Comment: [P2] Match Hadoop’s custom-separator behavior at buffer boundaries. With native V2 text enabled, read an uncompressed file containing 65,534 `x` bytes followed by `aaabz`, using `lineSep="aab"`. With its default 65,536-byte input buffer, Spark returns the entire file as one row. This substring matcher returns two rows: 65,534 `x` characters followed by `a`, and `z`. The file passes the size and split guards, so enabling native text silently changes counts, filters, and joins. Hadoop’s partial-delimiter handling across buffer boundaries differs from searching the complete file. Could the reader reproduce that behavior, including the configured buffer size, or fall back for affected repeated-prefix separators? Evidence: Compiled `split_lines` and `split_on_separator` unchanged from this head. The 65,539-byte case produced two native-helper rows of lengths 65,535 and 1. Spark V2 SQL 3.5.9 and 4.1.3 both reported `io.file.buffer.size=65536` and returned one row of length 65,539. Moving the prefix one byte earlier made both implementations return two rows. A 19-byte reproduction using 14 `x` bytes plus `aaabz` and `.option("io.file.buffer.size", "16")` showed the same mismatch. Hadoop 3.4.2 `LineReader.readCustomLine` confirms the boundary-dependent partial-match behavior. ########## spark/src/test/scala/org/apache/comet/text/CometTextNativeReadSuite.scala: ########## @@ -0,0 +1,424 @@ +/* + * 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.comet.text + +import java.io.File +import java.nio.charset.StandardCharsets +import java.nio.file.Files + +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.comet.CometTextNativeScanExec +import org.apache.spark.sql.functions.input_file_name +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus + +class CometTextNativeReadSuite extends CometTestBase { + private val TEST_TEXT_PATH = "src/test/resources/test-data/text-test-1.txt" + + private def withTextConf(f: => Unit): Unit = { + withSQLConf( + CometConf.COMET_TEXT_V2_NATIVE_ENABLED.key -> "true", + // Route only `text` through the V2 path; keep every other format on V1 so a Parquet probe + // side (used by the broadcast-join test) is not pushed through the unsupported V2 scan. + SQLConf.USE_V1_SOURCE_LIST.key -> "avro,csv,json,kafka,orc,parquet")(f) + } + + test("native text read - lines to value column") { + withTextConf { + val df = spark.read.text(TEST_TEXT_PATH) + checkSparkAnswerAndOperator(df) + } + } + + test("native text read - count") { + withTextConf { + val df = spark.read.text(TEST_TEXT_PATH).selectExpr("count(*) as c") + checkSparkAnswerAndOperator(df) + } + } + + test("native text read - filter and projection") { + withTextConf { + val df = spark.read.text(TEST_TEXT_PATH).where("value like '%bar%'") + checkSparkAnswerAndOperator(df) + } + } + + test("native text read - preserves line order within a file") { + withTempDir { dir => + val file = new File(dir, "ordered.txt") + val lines = (1 to 500).map(i => s"row-$i") + Files.write(file.toPath, lines.mkString("\n").getBytes(StandardCharsets.UTF_8)) + withTextConf { + val df = spark.read.text(file.getAbsolutePath) + // Assert the scan actually ran natively (a Spark fallback would also preserve order). + checkSparkAnswerAndOperator(df, Seq(classOf[CometTextNativeScanExec])) + // checkSparkAnswer sorts both sides, so assert order explicitly here. + val got = df.collect().map(_.getString(0)).toSeq + assert(got == lines, "native text scan must preserve within-file line order") + } + } + } + + test("native text read - fallback for input_file_name") { + withTextConf { + val df = spark.read.text(TEST_TEXT_PATH).withColumn("f", input_file_name()) + checkSparkAnswerAndFallbackReason( + df, + "Native Text scan is not compatible with input_file_name") Review Comment: [P2] Update this assertion to the shared V2 fallback reason. `transformV2Scan` handles `input_file_name` before reaching the text-specific branch, so the actual reason starts with `Native V2 scan is not compatible with input_file_name`. The expected `Native Text scan...` string cannot match. This new test fails on all five Spark profiles and blocks `Required Checks`, although the fallback itself works. Assert the shared V2 message while retaining the answer and fallback checks. Evidence: Exact-head run 38305916651 reports this as the sole failed test in every Spark 3.4, 3.5, 4.0, 4.1, and 4.2 scan job. The Spark 4.1 job, https://github.com/apache/datafusion-comet/actions/runs/38305916651/job/115002925730, reports 737 successful tests and one failure: expected `Native Text scan is not compatible with input_file_name`, actual `Native V2 scan is not compatible with input_file_name, input_file_block_start, or input_file_block_length`. -- 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]
