diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 7f70f81122c..cfca6f2fc7e 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -549,6 +549,7 @@ jobs: org.apache.comet.iceberg.IcebergReflectionSuite org.apache.comet.serde.operator.IcebergWriteProtoTranslationSuite org.apache.comet.csv.CometCsvNativeReadSuite + org.apache.comet.text.CometTextNativeReadSuite org.apache.comet.CometFuzzTestSuite org.apache.comet.CometFuzzIcebergSuite org.apache.comet.DataGeneratorSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 44b70c9e2ca..73802326edb 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -184,6 +184,7 @@ jobs: org.apache.comet.iceberg.IcebergReflectionSuite org.apache.comet.serde.operator.IcebergWriteProtoTranslationSuite org.apache.comet.csv.CometCsvNativeReadSuite + org.apache.comet.text.CometTextNativeReadSuite org.apache.comet.CometFuzzTestSuite org.apache.comet.CometFuzzIcebergSuite org.apache.comet.DataGeneratorSuite diff --git a/docs/source/user-guide/latest/datasources.md b/docs/source/user-guide/latest/datasources.md index a24cc12a8ce..dd3afccf2aa 100644 --- a/docs/source/user-guide/latest/datasources.md +++ b/docs/source/user-guide/latest/datasources.md @@ -51,11 +51,27 @@ converted into Arrow format, allowing the Comet pipeline to take over after that ### Text -Comet does not provide a Rust-based Text scan, but when `spark.comet.convert.text.enabled` is enabled, data from Spark's -Text reader is immediately converted into Arrow format, allowing the Comet pipeline to take over after that. Spark's Text -reader keeps the bytes of each line as they are, so a file in an encoding other than UTF-8, such as Latin-1, holds strings -that are not valid UTF-8. Arrow strings must be valid UTF-8, so the conversion replaces those bytes with `U+FFFD`, and -results that depend on the original bytes can differ from Spark. Enable it only for Text files that are UTF-8. See +Comet provides experimental Rust-based text scan support (a single `value: string` column). When +`spark.comet.scan.text.v2.enabled` is enabled, text files are read in Rust. This feature is experimental and +performance benefits are workload-dependent. Only Spark's DataSource V2 text scan is accelerated, and Spark reads +text through the V1 API by default, so also remove `text` from `spark.sql.sources.useV1SourceList`. The native scan +falls back to Spark for cases it does not handle -- including compressed files, files larger than +`spark.comet.scan.text.maxFileSize` (64MB by default, since the reader loads each file whole), files Spark splits +into byte ranges, unsupported filesystem schemes, and `input_file_name()` or +`ignoreCorruptFiles`/`ignoreMissingFiles` usage. It also falls back for local files that have a Hadoop checksum +(`.crc`) file, such as data Spark itself wrote to the local filesystem, because Hadoop verifies the checksum on read and +the native reader does not. It targets small lookup tables such as a broadcast join's build side. +Reading a partition column from a partitioned directory falls back to Spark, but selecting only `value` stays native. + +Text files are a common source of ill-formed UTF-8. The native reader replaces invalid byte sequences with the +Unicode replacement character (U+FFFD), so byte-level comparisons and functions can differ from Spark on those rows; +see [Strings with non-UTF-8 bytes](compatibility/index.md#strings-with-non-utf-8-bytes) in the compatibility guide. + +Separately, `spark.comet.convert.text.enabled` converts data from Spark's own Text reader into Arrow format +immediately, allowing the Comet pipeline to take over after that. Spark's Text reader keeps the bytes of each line as +they are, so a file in an encoding other than UTF-8, such as Latin-1, holds strings that are not valid UTF-8. Arrow +strings must be valid UTF-8, so the conversion replaces those bytes with `U+FFFD`, and results that depend on the +original bytes can differ from Spark. Enable it only for Text files that are UTF-8. See [Strings with non-UTF-8 bytes](compatibility/index.md#strings-with-non-utf-8-bytes). ### Other Spark inputs @@ -130,13 +146,13 @@ Comet supports most standard storage systems, such as local file system and obje ### Scans over more than one object store -A native Parquet or CSV scan reads each of its partitions through one object store, such as one S3 +A native Parquet, CSV, or Text scan reads each of its partitions through one object store, such as one S3 bucket, one GCS bucket, or one HDFS name node. When a scan's files live in more than one store, Comet puts the files of each store into separate partitions, so such a scan can run with a different number of partitions than Spark would use. A native Parquet scan packs each store's files on their -own, so `spark.sql.files.maxPartitionNum` applies to each store's files separately. A native CSV -scan that would split a partition falls back to Spark when Spark planned the operators above it for -the scan's own partitioning, as Spark 3.4 does for a scan of one partition. +own, so `spark.sql.files.maxPartitionNum` applies to each store's files separately. A native CSV or +Text scan that would split a partition falls back to Spark when Spark planned the operators above +it for the scan's own partitioning, as Spark 3.4 does for a scan of one partition. Such a scan forwards the object store settings of every scheme it reads, so its files can mix schemes such as `s3a://`, `gs://` and `hdfs://`. It falls back to Spark when it reads one bucket diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 2e1a49d23a1..e8d301ef2ff 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -61,7 +61,7 @@ omitted from the tables below and may be reconsidered based on demand: | Operator | Status | Notes | | ----------------------- | ------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | `FileSourceScanExec` | ✅ | Parquet only. Some types and configurations fall back. See [Parquet](parquet.md#reading). | -| `BatchScanExec` | ✅ | Apache Iceberg Parquet scans run natively. Native CSV scans are experimental and disabled by default. DataSource V2 Parquet scans are not accelerated. See [Parquet](parquet.md#reading) and the [Iceberg Guide](iceberg.md). | +| `BatchScanExec` | ✅ | Apache Iceberg Parquet scans run natively. Native CSV and Text scans are experimental and disabled by default. DataSource V2 Parquet scans are not accelerated. See [Parquet](parquet.md#reading) and the [Iceberg Guide](iceberg.md). | | `LocalTableScanExec` | ⚠️ | Disabled by default. Converting it rarely speeds up the scan itself, but a native Iceberg write needs it to write the rows of `INSERT ... VALUES` or a local DataFrame natively; see [Iceberg Writes](iceberg.md#native-parquet-write-eligibility). Enable it with `spark.comet.exec.localTableScan.enabled=true` ([#4393](https://github.com/apache/datafusion-comet/pull/4393)). | | `EmptyRelationExec` | ✅ | Spark 4.0 and later. See [Empty Relations](compatibility/operators.md#empty-relations) for native-input support, including native Parquet writes. | | `RangeExec` | ⚠️ | Disabled by default. Set `spark.comet.exec.range.enabled=true` to generate the rows of `spark.range` and SQL `range()` in native code, so the operators above them run natively. It can be slower than Spark when those operators are only cheap expressions, such as a filter, which Spark compiles together with the range into one loop. Ranges whose arithmetic overflows the `Long` range fall back to Spark. | diff --git a/docs/source/user-guide/latest/parquet.md b/docs/source/user-guide/latest/parquet.md index 04132cf8e47..3759137b4b2 100644 --- a/docs/source/user-guide/latest/parquet.md +++ b/docs/source/user-guide/latest/parquet.md @@ -98,8 +98,9 @@ reads that fail where Spark succeeds are listed in query, or cast the column to `BINARY` before persisting, if you need to preserve non-UTF-8 bytes. By contrast, Comet decodes non-UTF-8 bytes at the JVM-to-native Arrow FFI boundaries using JVM-compatible replacement semantics. This covers native query input from the JVM-exported Arrow - stream (including JVM scans, shuffle reads, and `mapInArrow`), columnar-to-row conversion, and JVM - UDF results. It does not change the native Parquet reader: rejecting invalid UTF-8 there remains + stream (including JVM scans, shuffle reads, and `mapInArrow`), columnar-to-row conversion, JVM + UDF results, and the native text scan (which replaces invalid bytes in each line's `value` with + U+FFFD). It does not change the native Parquet reader: rejecting invalid UTF-8 there remains Gap A of [#4764](https://github.com/apache/datafusion-comet/issues/4764). See [Strings with non-UTF-8 bytes](compatibility/index.md#strings-with-non-utf-8-bytes), [#4121](https://github.com/apache/datafusion-comet/issues/4121), and the tracking issue above. diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index d86f8076bde..e8abdb68e9e 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -300,6 +300,7 @@ by role. Names match what is shown in the plan output. | `CometNativeScan` | Parquet scan that runs entirely in Rust via DataFusion. | | `CometIcebergNativeScan` | Iceberg Parquet scan that runs entirely in Rust via DataFusion. | | `CometCsvNativeScan` | CSV scan that runs entirely in Rust via DataFusion (experimental). | +| `CometTextNativeScan` | Text scan (single `value` string column) that runs entirely in Rust via DataFusion (experimental). | | `CometInMemoryTableScan` | JVM-side scan of a table cached in Comet's Arrow format (experimental, disabled by default). | | `CometLocalTableScan` | JVM-side scan of a `LocalTableScanExec` (disabled by default). | | `CometEmptyRelation` | Empty input to a Rust plan that replaces `EmptyRelationExec` (Spark 4.0 and later). | diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index a01789dead3..916d025e17b 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -496,6 +496,7 @@ fn op_name(op: &OpStruct) -> &'static str { OpStruct::ParquetWriter(_) => "ParquetWriter", OpStruct::Explode(_) => "Explode", OpStruct::CsvScan(_) => "CsvScan", + OpStruct::TextScan(_) => "TextScan", OpStruct::ShuffleScan(_) => "ShuffleScan", OpStruct::BroadcastNestedLoopJoin(_) => "BroadcastNestedLoopJoin", OpStruct::CrossJoin(_) => "CrossJoin", diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index 49669c58314..ac88b787426 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -49,6 +49,8 @@ pub mod projection; mod scan; mod shuffle_scan; mod sorted_input; +mod text_scan; pub use csv_scan::init_csv_datasource_exec; pub use shuffle_scan::ShuffleScanExec; pub use sorted_input::SortedInputExec; +pub use text_scan::init_text_datasource_exec; diff --git a/native/core/src/execution/operators/text_scan.rs b/native/core/src/execution/operators/text_scan.rs new file mode 100644 index 00000000000..e1163942233 --- /dev/null +++ b/native/core/src/execution/operators/text_scan.rs @@ -0,0 +1,701 @@ +// 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_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::fmt; +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>, + data_schema: SchemaRef, + projection_vector: Vec, + text_options: &TextOptions, +) -> Result, 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, + 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>, +} + +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 for Arc { + fn from(source: TextSource) -> Self { + as_file_source(source) + } +} + +impl FileSource for TextSource { + fn create_file_opener( + &self, + object_store: Arc, + _base_config: &FileScanConfig, + partition_index: usize, + ) -> Result> { + // 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, + 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; + 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 { + let mut conf = self.clone(); + conf.batch_size = Some(batch_size); + Arc::new(conf) + } + + fn try_pushdown_projection( + &self, + projection: &ProjectionExprs, + ) -> Result>> { + 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, + ) -> Result, + ) -> Result { + 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, + object_store: Arc, + partition_index: usize, + batch_size: usize, + metrics: ExecutionPlanMetricsSet, + whole_text: bool, + line_sep: Option>, +} + +impl FileOpener for TextOpener { + fn open(&self, partitioned_file: PartitionedFile) -> Result { + let store = Arc::clone(&self.object_store); + let projected_file_schema = Arc::clone(&self.projected_file_schema); + 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 location = partitioned_file.object_meta.location; + let buf: Bytes = store.get(&location).await?.bytes().await?; + bytes_scanned.add(buf.len()); + + // NOTE: the whole object is read into memory. CometScanRule only admits unsplit, + // uncompressed files under a size cap, 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. + // + // Spark's line reader (Hadoop LineRecordReader.skipUtfByteOrderMark) strips a leading + // UTF-8 BOM from the first record, whatever the separator. wholetext keeps it. + let reader = TextBatchReader::new( + buf, + projected_file_schema, + batch_size, + TextLayout { + whole_text, + line_sep, + strip_bom: !whole_text, + }, + ); + + 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>, + /// Strip a leading UTF-8 BOM from the first record (line mode). + strip_bom: bool, +} + +/// 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 strip. + fn next_row(&mut self) -> Option<(usize, usize)> { + if let Some(row) = self.pending.take() { + return Some(row); + } + 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(); + } + } + 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> { + 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 append it straight to the builder. Count the DECODED size, + // since a lossy decode can expand ill-formed bytes up to 3x, and 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 value borrow them + // while the reader advances. + let buf = self.buf.clone(); + let mut builder = StringBuilder::with_capacity(self.batch_rows.min(buf.len().max(1)), 0); + let mut rows = 0usize; + let mut data_bytes = 0usize; + while rows < 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 rows > 0 && data_bytes + len > MAX_BATCH_VALUE_BYTES { + self.pending = Some((start, end)); + break; + } + data_bytes += len; + builder.append_value(value); + rows += 1; + } + if rows == 0 { + self.finished = true; + return Ok(None); + } + Ok(Some(RecordBatch::try_new( + Arc::clone(&self.schema), + vec![Arc::new(builder.finish())], + )?)) + } +} + +/// Offset of the first occurrence of `needle` in `haystack`. +fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option { + if needle.is_empty() || needle.len() > haystack.len() { + return None; + } + haystack.windows(needle.len()).position(|w| w == needle) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + + fn value_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])) + } + + /// Read `buf` the way the opener does and return every batch. + fn read_batches( + buf: &[u8], + schema: SchemaRef, + batch_size: usize, + layout: TextLayout, + ) -> Vec { + let mut reader = + TextBatchReader::new(Bytes::copy_from_slice(buf), schema, batch_size, layout); + let mut batches = Vec::new(); + while let Some(batch) = reader.next_batch().unwrap() { + batches.push(batch); + } + batches + } + + fn layout(whole_text: bool, line_sep: Option<&[u8]>) -> TextLayout { + TextLayout { + whole_text, + line_sep: line_sep.map(|s| s.to_vec()), + strip_bom: false, + } + } + + fn values(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|batch| { + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + (0..col.len()) + .map(|i| col.value(i).to_string()) + .collect::>() + }) + .collect() + } + + fn lines_as_str(buf: &[u8], whole_text: bool, line_sep: Option<&[u8]>) -> Vec { + values(&read_batches( + buf, + value_schema(), + 8192, + layout(whole_text, line_sep), + )) + } + + #[test] + fn universal_newline_variants() { + // \n, \r\n, and \r all split; the terminator is stripped. + assert_eq!( + lines_as_str(b"a\nb\r\nc\rd", false, None), + vec!["a", "b", "c", "d"] + ); + } + + #[test] + fn trailing_terminator_no_empty_line() { + assert_eq!(lines_as_str(b"a\nb\n", false, None), vec!["a", "b"]); + assert_eq!(lines_as_str(b"a\r\n", false, None), vec!["a"]); + assert_eq!(lines_as_str(b"a\r", false, None), vec!["a"]); + } + + #[test] + fn empty_interior_lines_are_kept() { + assert_eq!(lines_as_str(b"a\n\nb", false, None), vec!["a", "", "b"]); + // A lone newline is one empty record. + assert_eq!(lines_as_str(b"\n", false, None), vec![""]); + } + + #[test] + fn empty_buffer_has_no_lines() { + assert!(lines_as_str(b"", false, None).is_empty()); + } + + #[test] + fn no_terminator_is_single_line() { + assert_eq!( + lines_as_str(b"only one line", false, None), + vec!["only one line"] + ); + } + + #[test] + fn custom_single_char_separator() { + // A comma and a bare newline are literal content when lineSep is '|'. + assert_eq!( + lines_as_str(b"a,b|c\nd|e", false, Some(b"|")), + vec!["a,b", "c\nd", "e"] + ); + // Trailing separator does not add an empty line. + assert_eq!(lines_as_str(b"a|b|", false, Some(b"|")), vec!["a", "b"]); + } + + #[test] + fn custom_multi_byte_separator() { + assert_eq!( + lines_as_str(b"a<>b<>c", false, Some(b"<>")), + vec!["a", "b", "c"] + ); + } + + #[test] + fn whole_text_is_one_row_with_separators_preserved() { + assert_eq!(lines_as_str(b"a\nb\nc", true, None), vec!["a\nb\nc"]); + } + + #[test] + fn non_ascii_is_preserved() { + assert_eq!( + lines_as_str("café\n日本語".as_bytes(), false, None), + vec!["café", "日本語"] + ); + } + + #[test] + fn batches_split_rows_by_batch_size() { + let batches = read_batches(b"x\ny\nz", value_schema(), 2, layout(false, None)); + // batch_size 2 -> two batches (2 rows, 1 row). + assert_eq!(batches.len(), 2); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 3); + assert_eq!(values(&batches), vec!["x", "y", "z"]); + assert_eq!(batches[0].num_rows(), 2); + } + + #[test] + fn empty_projection_keeps_row_count() { + // The count(*) case: an empty projected schema, rows preserved with no columns. + let batches = read_batches( + b"x\ny\nz", + Arc::new(Schema::empty()), + 8192, + layout(false, None), + ); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_columns(), 0); + assert_eq!(batches[0].num_rows(), 3); + } + + #[test] + fn newline_only_input_counts_rows_in_bounded_batches() { + // A file of only newlines has one empty row per byte. Counting must not need a table + // that grows with the input: the rows come out in batches of at most batch_size. + let buf = vec![b'\n'; 100_000]; + let batches = read_batches(&buf, Arc::new(Schema::empty()), 8192, layout(false, None)); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 100_000); + assert!(batches.iter().all(|b| b.num_rows() <= 8192)); + } + + #[test] + fn reader_returns_one_batch_per_call() { + // Batches are built on demand: after the first call the reader still holds the rest. + let mut reader = TextBatchReader::new( + Bytes::from_static(b"a\nb\nc\nd\ne"), + value_schema(), + 2, + layout(false, None), + ); + assert_eq!( + values(&[reader.next_batch().unwrap().unwrap()]), + vec!["a", "b"] + ); + assert_eq!( + values(&[reader.next_batch().unwrap().unwrap()]), + vec!["c", "d"] + ); + assert_eq!(values(&[reader.next_batch().unwrap().unwrap()]), vec!["e"]); + assert!(reader.next_batch().unwrap().is_none()); + assert!(reader.next_batch().unwrap().is_none()); + } + + /// Read with the first-record BOM strip, as the opener sets it (line mode only). + fn lines_after_bom(buf: &[u8], whole_text: bool, line_sep: Option<&[u8]>) -> Vec { + let mut layout = layout(whole_text, line_sep); + layout.strip_bom = !whole_text; + values(&read_batches(buf, value_schema(), 8192, layout)) + } + + #[test] + fn strips_first_record_bom_in_line_mode() { + let bom = [0xEFu8, 0xBB, 0xBF]; + let with_hello = [&bom[..], b"hello"].concat(); + // Default separator at file start: the BOM is stripped. + assert_eq!(lines_after_bom(&with_hello, false, None), vec!["hello"]); + // BOM then newline -> ["", "hello"]. + let bom_nl = [&bom[..], b"\nhello"].concat(); + assert_eq!(lines_after_bom(&bom_nl, false, None), vec!["", "hello"]); + // BOM-only file, default separator -> no rows. + assert!(lines_after_bom(&bom, false, None).is_empty()); + // BOM only plus its terminator -> one empty row. + let bom_only_nl = [&bom[..], b"\n"].concat(); + assert_eq!(lines_after_bom(&bom_only_nl, false, None), vec![""]); + // Custom separator: the BOM is still stripped from the first record. + let bom_a_b = [&bom[..], b"a|b"].concat(); + assert_eq!(lines_after_bom(&bom_a_b, false, Some(b"|")), vec!["a", "b"]); + // BOM-only file with a custom separator -> no rows. + assert!(lines_after_bom(&bom, false, Some(b"|")).is_empty()); + // BOM then separator -> an empty row, then the rest. + let bom_sep_x = [&bom[..], b"|x"].concat(); + assert_eq!( + lines_after_bom(&bom_sep_x, false, Some(b"|")), + vec!["", "x"] + ); + // lineSep equal to the BOM: the BOM is the separator, so the first record does not start + // with it and nothing is stripped. `BOM a BOM b` -> ["", "a", "b"]. + let bom_a_bom_b = [&bom[..], b"a", &bom[..], b"b"].concat(); + assert_eq!( + lines_after_bom(&bom_a_bom_b, false, Some(&bom)), + vec!["", "a", "b"] + ); + // BOM-only file with lineSep=BOM -> one empty record (trailing separator dropped). + assert_eq!(lines_after_bom(&bom, false, Some(&bom)), vec![""]); + // Not stripped: wholetext keeps the BOM, and a file with no BOM is unchanged. + assert_eq!( + lines_after_bom(&with_hello, true, None).concat().as_bytes(), + &with_hello[..] + ); + assert_eq!(lines_after_bom(b"hello", false, None), vec!["hello"]); + } + + #[test] + fn invalid_utf8_is_replaced() { + // Ill-formed bytes become U+FFFD (matches Spark's new String(bytes, UTF_8) on collect). + let bad: &[u8] = &[b'a', 0xff, b'b']; + assert_eq!(lines_as_str(bad, false, None), vec!["a\u{FFFD}b"]); + } +} diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 603bba872bc..de44acc8967 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -46,6 +46,7 @@ mod delta_spark_scan; mod lance_scan; use crate::execution::operators::init_csv_datasource_exec; +use crate::execution::operators::init_text_datasource_exec; use crate::execution::operators::IcebergScanExec; use crate::execution::operators::SortedInputExec; use crate::execution::operators::TopKReaderFilterExec; @@ -385,6 +386,32 @@ impl PhysicalPlanner { Ok((url, backend)) } + /// The object store and file group one CSV or Text scan task reads: this task's own partition, + /// with the store taken from its first file, because partitions can differ in store. + fn scan_task_files( + &self, + scan_name: &str, + file_partitions: &[SparkFilePartition], + object_store_options: &HashMap, + ) -> Result<(ObjectStoreUrl, Vec>), ExecutionError> { + let partition = usize::try_from(self.partition) + .ok() + .and_then(|p| file_partitions.get(p)) + .ok_or_else(|| { + GeneralError(format!( + "{scan_name} has no file partition {} ({} partitions)", + self.partition, + file_partitions.len() + )) + })?; + let object_store_url = match partition.partitioned_file.first() { + Some(file) => self.scan_store(&file.file_path, object_store_options)?.0, + None => ObjectStoreUrl::local_filesystem(), + }; + let files = self.get_partitioned_files(partition, object_store_options)?; + Ok((object_store_url, vec![files])) + } + /// Plans `spark_plan` as the next native input when it is an input leaf and the planner has /// native inputs. `None` otherwise. fn native_input_leaf( @@ -1890,24 +1917,11 @@ impl PhysicalPlanner { .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(); - let partition = usize::try_from(self.partition) - .ok() - .and_then(|p| scan.file_partitions.get(p)) - .ok_or_else(|| { - GeneralError(format!( - "CsvScan 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) => self.scan_store(&file.file_path, &object_store_options)?.0, - None => ObjectStoreUrl::local_filesystem(), - }; - let files = self.get_partitioned_files(partition, &object_store_options)?; - let file_groups: Vec> = vec![files]; + let (object_store_url, file_groups) = self.scan_task_files( + "CsvScan", + &scan.file_partitions, + &object_store_options, + )?; let scan = init_csv_datasource_exec( object_store_url, file_groups, @@ -1922,6 +1936,32 @@ 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 projection_vector: Vec = + scan.projection_vector.iter().map(|i| *i as usize).collect(); + let object_store_options: HashMap = + scan.object_store_options.clone(); + let (object_store_url, file_groups) = self.scan_task_files( + "TextScan", + &scan.file_partitions, + &object_store_options, + )?; + let scan = init_text_datasource_exec( + object_store_url, + file_groups, + data_schema, + projection_vector, + scan.text_options + .as_ref() + .ok_or_else(|| GeneralError("TextScan missing text_options".to_string()))?, + )?; + Ok(( + vec![], + vec![], + Arc::new(SparkPlan::new(spark_plan.plan_id, scan, vec![])), + )) + } OpStruct::Scan(scan) => { let data_types = scan.fields.iter().map(to_arrow_datatype).collect_vec(); diff --git a/native/core/src/execution/planner/operator_registry.rs b/native/core/src/execution/planner/operator_registry.rs index ccc7f4afbf9..40fff69d42d 100644 --- a/native/core/src/execution/planner/operator_registry.rs +++ b/native/core/src/execution/planner/operator_registry.rs @@ -60,6 +60,7 @@ pub enum OperatorType { HashJoin, Window, CsvScan, + TextScan, MergeRows, Explode, BroadcastNestedLoopJoin, @@ -239,6 +240,7 @@ fn get_operator_type(spark_operator: &Operator) -> Option { OpStruct::Window(_) => Some(OperatorType::Window), OpStruct::Explode(_) => Some(OperatorType::Explode), OpStruct::CsvScan(_) => Some(OperatorType::CsvScan), + OpStruct::TextScan(_) => Some(OperatorType::TextScan), OpStruct::ShuffleScan(_) => Some(OperatorType::ShuffleScan), OpStruct::BroadcastNestedLoopJoin(_) => Some(OperatorType::BroadcastNestedLoopJoin), OpStruct::CrossJoin(_) => Some(OperatorType::CrossJoin), diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index fdd2c9e730d..8da5f726b97 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -72,6 +72,7 @@ message Operator { MergeRows merge_rows = 122; ArrowPythonUdf arrow_python_udf = 123; CrossJoin cross_join = 124; + TextScan text_scan = 125; // Extension point for optional, out-of-tree contrib scans (Delta, Lance, ...). The concrete // scan message (e.g. `DeltaScan`) is packed into this envelope on the JVM side and dispatched // by `type_url` on the native side. Using a single permanent field -- rather than a new oneof @@ -292,6 +293,26 @@ message CsvOptions { bool truncated_rows = 8; } +// Native Text scan, mirroring Spark's `text` data source (a single `value: string` column). +message TextScan { + // Partition columns are not supported by the native Text scan, so there is no partition schema. + reserved 2; + reserved "partition_schema"; + repeated SparkStructField data_schema = 1; + repeated int32 projection_vector = 3; + repeated SparkFilePartition file_partitions = 4; + map object_store_options = 5; + TextOptions text_options = 6; +} + +message TextOptions { + // When true, each file is read as a single row instead of being split into lines. + bool whole_text = 1; + // Custom line separator, as the exact bytes Spark computed (lineSeparatorInRead, already encoded + // with the file's charset). When unset, lines are split on `\n`, `\r\n`, or `\r`. + optional bytes line_sep = 2; +} + // A typed Iceberg primitive value, shared by partition values and predicate literals. The oneof // case is the type tag (iceberg-rust's Datum is self-typed and Predicate::bind does not coerce, so // the type must be known at construction). Grows by adding a case as Iceberg adds primitives diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index a80cdd76f1b..65b7aea1730 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -190,6 +190,27 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_TEXT_V2_NATIVE_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.scan.text.v2.enabled") + .category(CATEGORY_TESTING) + .doc( + "Whether to use the native Comet V2 reader for the text file format (a single " + + "'value: string' column). Requires the V2 text source (remove 'text' from " + + "spark.sql.sources.useV1SourceList). Default: false. " + + "Experimental: Performance benefits are workload-dependent.") + .booleanConf + .createWithDefault(false) + + val COMET_SCAN_TEXT_MAX_FILE_SIZE: ConfigEntry[Long] = + conf("spark.comet.scan.text.maxFileSize") + .category(CATEGORY_TESTING) + .doc("Maximum size of a single text file that the native Text scan will read. The native " + + "reader loads each file whole and does not reserve the buffer against the memory pool, " + + "so this bounds its peak memory; files larger than this fall back to Spark. The native " + + "Text scan targets small lookup tables such as a broadcast join's build side.") + .bytesConf(ByteUnit.BYTE) + .createWithDefault(64L * 1024 * 1024) + val COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.parquet.rowFilterPushdown.enabled") .category(CATEGORY_PARQUET) diff --git a/spark/src/main/scala/org/apache/comet/objectstore/NativeConfig.scala b/spark/src/main/scala/org/apache/comet/objectstore/NativeConfig.scala index b6f63ac77ef..13622c894ab 100644 --- a/spark/src/main/scala/org/apache/comet/objectstore/NativeConfig.scala +++ b/spark/src/main/scala/org/apache/comet/objectstore/NativeConfig.scala @@ -74,7 +74,7 @@ object NativeConfig { s3Properties.map { case (_, s3aSuffix, icebergKey) => s3aSuffix -> icebergKey }.toMap /** A URI's scheme, lowercased; None when it has none. */ - private[comet] def lowerScheme(uri: URI): Option[String] = + def lowerScheme(uri: URI): Option[String] = Option(uri.getScheme).map(_.toLowerCase(Locale.ROOT)) // Comma-separated scheme list -> trimmed, lowercased set (case-insensitive). Shared with @@ -146,7 +146,7 @@ object NativeConfig { * which native reads through libhdfs. `scanOnly` holds those listed only through the scan's * reader options, which the scan gates decline. */ - private[comet] case class ScanLibhdfsSchemes(libhdfs: Set[String], scanOnly: Set[String]) + case class ScanLibhdfsSchemes(libhdfs: Set[String], scanOnly: Set[String]) /** * Splits the `fs.comet.libhdfs.schemes` lists of `sessionConf` (the session's Hadoop conf) and @@ -154,7 +154,7 @@ object NativeConfig { * the scan's list, but libhdfs builds its configuration from the executor classpath * (`core-site.xml`), not the scan's options, so a reader option cannot add a libhdfs scheme. */ - private[comet] def resolveScanLibhdfsSchemes( + def resolveScanLibhdfsSchemes( sessionConf: Configuration, scanConf: Configuration): ScanLibhdfsSchemes = resolveScanLibhdfsSchemes(resolveLibhdfsSchemes(sessionConf), scanConf) @@ -163,7 +163,7 @@ object NativeConfig { * As above, with the session's list already resolved by [[resolveLibhdfsSchemes]], so a caller * planning many scans copies the session's Hadoop conf once. */ - private[comet] def resolveScanLibhdfsSchemes( + def resolveScanLibhdfsSchemes( sessionSchemes: Set[String], scanConf: Configuration): ScanLibhdfsSchemes = { val scan = resolveLibhdfsSchemes(scanConf) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index a73a9b55cbd..ca8936d408c 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -474,6 +474,10 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) case scan: CometBatchScanExec if scan.wrapped.scan.isInstanceOf[CSVScan] => convertToComet(scan, CometCsvNativeScanExec).getOrElse(scan.wrapped) + // A declined native Text scan falls back to Spark's scan, which produces rows. + case scan: CometBatchScanExec if scan.wrapped.scan.isInstanceOf[TextScan] => + convertToComet(scan, CometTextNativeScanExec).getOrElse(scan.wrapped) + // Comet JVM + native scan for V1 and V2 case op if isCometScan(op) => convertToComet(op, CometScanWrapper).getOrElse(op) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index 955db40ce42..7eeed2ff854 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -38,9 +38,10 @@ import org.apache.spark.sql.catalyst.util.{sideBySide, ArrayBasedMapData, DateTi import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns.getExistenceDefaultValues import org.apache.spark.sql.comet.{CometBatchScanExec, CometScanExec, CometScanUtils} import org.apache.spark.sql.execution.{FileSourceScanExec, InSubqueryExec, SparkPlan, SubqueryAdaptiveBroadcastExec} -import org.apache.spark.sql.execution.datasources.HadoopFsRelation +import org.apache.spark.sql.execution.datasources.{FilePartition, HadoopFsRelation} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.execution.datasources.v2.csv.CSVScan +import org.apache.spark.sql.execution.datasources.v2.text.TextScan import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -258,74 +259,15 @@ case class CometScanRule(session: SparkSession) return None } // Comet's native readers go through object_store, which only understands a fixed set of URL - // schemes. A custom Hadoop FileSystem (e.g. registered via spark.hadoop.fs..impl) would - // surface at execution time as `Generic URL error: Unable to recognise URL "..."`. Decline here - // so Spark's reader -- which goes through the Hadoop FS API and can resolve custom schemes -- - // handles the scan. Whether object_store recognizes a scheme is answered by the native layer - // itself (`NativeBase.isObjectStoreSchemeSupported`) rather than a hardcoded list, so the - // planner can't drift from object_store's actual support. - // - // EXCEPT schemes native routes through libhdfs (`hdfs` by default). A scheme the reader - // options add to that list is declined; see NativeConfig.resolveScanLibhdfsSchemes. - val libhdfs = - NativeConfig.resolveScanLibhdfsSchemes(sessionLibhdfsSchemes(r.sparkSession), hadoopConf) - val libhdfsSchemes = libhdfs.libhdfs - // Opt-in S3-compliant alias schemes (e.g. `blob`) from `fs.comet.s3Compliant.schemes`. Read - // from the Hadoop config, like the libhdfs list above, so `core-site.xml` is honored. The - // native object_store gate no longer claims aliases, so admit them here where the Hadoop - // config is available. - val s3CompliantSchemes = NativeConfig.resolveS3CompliantSchemes(hadoopConf) - - // Classify each root path once; see RootPathInfo. - val roots = CometScanRule.classifyRootPaths( - r.location.rootPaths.map(_.toUri), - libhdfsSchemes, - s3CompliantSchemes) - - val scanOnlyLibhdfsSchemes = roots.flatMap(_.scheme).filter(libhdfs.scanOnly.contains).toSet - if (scanOnlyLibhdfsSchemes.nonEmpty) { - withFallbackReason( - scanExec, - "Native Parquet scan does not support filesystem scheme(s) " + - s"${scanOnlyLibhdfsSchemes.toSeq.sorted.mkString(", ")} listed in " + - s"$COMET_LIBHDFS_SCHEMES_KEY only through reader options") - return None - } - val unsupportedFsSchemes = roots.iterator.flatMap { root => - root.scheme.filter(_ => - !root.isLibhdfs && !CometScanRule.isNativelyReadableScheme(root.uri, s3CompliantSchemes)) - }.toSet - if (unsupportedFsSchemes.nonEmpty) { - withFallbackReason( - scanExec, - s"Unsupported filesystem schemes: ${unsupportedFsSchemes.mkString(", ")}") - return None - } - // An early check on the root paths, with both scheme lists from the Hadoop conf as native - // reads them (not libhdfsSchemes above). CometNativeScan.convert checks the listed files. - val multiStoreReason = CometScanUtils.multiStoreFallbackReason( + // schemes; see rootPathFallbackReason for the scheme, multi-store, and path checks. + val rootPathReason = CometScanRule.rootPathFallbackReason( "Native Parquet scan", - roots.map(_.uri), - s3CompliantSchemes, - NativeConfig.resolveLibhdfsSchemes(hadoopConf), + r.location.rootPaths.map(_.toUri), + sessionLibhdfsSchemes(r.sparkSession), + hadoopConf, scanExec.bucketedScan) - if (multiStoreReason.nonEmpty) { - withFallbackReason(scanExec, multiStoreReason.get) - return None - } - // A scheme object_store recognizes can still carry a path it rejects (e.g. a directory whose - // name contains a newline -> `%0A` in the URI), which native execution would hard-fail on. - // Only object_store-native schemes are probed: aliases (normalized to s3:// natively) and - // libhdfs schemes route elsewhere. Root paths only -- a rejected character deeper in the tree - // still fails at execution. - val rejectedPath = roots.find(root => - root.scheme.isDefined && !root.isLibhdfs && !root.isAlias && - !CometScanRule.objectStoreAcceptsPath(root.uri)) - if (rejectedPath.nonEmpty) { - withFallbackReason( - scanExec, - s"Native Parquet scan cannot open path '${rejectedPath.get.uri}': object_store " + - "rejects it (e.g. an unsupported character in the path)") + if (rootPathReason.nonEmpty) { + withFallbackReason(scanExec, rootPathReason.get) return None } // Disabling the vectorized reader opts into parquet-mr's permissive behavior @@ -477,6 +419,165 @@ 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" + } + // 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(_.asInstanceOf[FilePartition].files.toSeq) + // 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. + if (selectedFiles.exists(_.start != 0)) { + 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, whatever maxPartitionBytes says. A + // wholetext file becomes one Arrow string value, and a lossy UTF-8 decode can expand + // ill-formed bytes up to 3x, so cap it at Int.MaxValue/3 to keep the decoded value under + // Arrow's 2GB 32-bit string-offset limit (the native reader also fails cleanly if a value + // still overflows). + val maxTextFileSize = COMET_SCAN_TEXT_MAX_FILE_SIZE.get() + val maxFileBytes = + if (textOptions.wholeText) math.min(maxTextFileSize, Int.MaxValue.toLong / 3) + else maxTextFileSize + if (selectedFiles.exists(_.length > maxFileBytes)) { + fallbackReasons += "Comet native Text scan does not read files larger than " + + s"${COMET_SCAN_TEXT_MAX_FILE_SIZE.key} ($maxFileBytes bytes)" + } + // 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 `..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)" + } + } + } + // Comet's native reader opens files through object_store, which only understands a fixed + // set of URL schemes. A text file on a custom Hadoop scheme (viewfs://, oss://, ...) or a + // path object_store rejects would be claimed here then hard-fail at execution, whereas + // Spark reads it via the Hadoop FS API. Use the same shared gate as the native Parquet + // path, on the selected files rather than the root paths: a legal Hadoop filename inside + // a directory (for example one with a tab in its name) can pass for its parent and still + // be rejected by object_store. A scan that spans several object stores is split per store + // in the serde. + CometScanRule + .rootPathFallbackReason( + "Native Text scan", + selectedFiles.map(_.pathUri).distinct, + NativeConfig.resolveLibhdfsSchemes(session.sessionState.newHadoopConf()), + hadoopConf, + isBucketedScan = false) + .foreach(fallbackReasons += _) + if (fallbackReasons.isEmpty) { + CometBatchScanExec( + scanExec.clone().asInstanceOf[BatchScanExec], + runtimeFilters = scanExec.runtimeFilters) + } else { + withFallbackReasons(scanExec, fallbackReasons.toSet) + } + // Iceberg scan - detected by class name. SparkStagedScan covers reads issued by // RewriteDataFiles (and similar maintenance actions) where the planner has already // staged FileScanTasks via ScanTaskSetManager. @@ -1288,6 +1389,35 @@ object CometScanRule extends Logging { case _ => false })) + /** + * True when a custom line separator overlaps itself -- some proper prefix of the separator is + * also a suffix of one of its prefixes, so a partial match can restart (e.g. "aab", or "aba"). + * Spark splits custom separators with Hadoop's LineReader.readCustomLine, which backtracks only + * one byte on a mismatch instead of following that overlap, so for such a separator it can + * split differently from the native reader's exact substring scan. Uses the KMP prefix + * function: any nonzero entry means the separator overlaps itself. + */ + private[rules] def separatorSelfOverlaps(sep: Array[Byte]): Boolean = { + val f = new Array[Int](sep.length) + var k = 0 + var overlaps = false + var i = 1 + while (i < sep.length) { + while (k > 0 && sep(i) != sep(k)) { + k = f(k - 1) + } + if (sep(i) == sep(k)) { + k += 1 + } + f(i) = k + if (k > 0) { + overlaps = true + } + i += 1 + } + overlaps + } + // Memo of `NativeBase.isObjectStoreSchemeSupported`, keyed by the probe URL rather than the // scheme: object_store's parser keys on (scheme, host-presence), so an authorityless URL would // otherwise poison the authority-bearing form of the same scheme. @@ -1340,6 +1470,97 @@ object CometScanRule extends Logging { isAlias = scheme.exists(s3CompliantSchemes.contains)) } + /** + * True when a local file has a Hadoop checksum file (`..crc`) next to it. Hadoop's local + * filesystem verifies the data against that file on every read, which the native reader does + * not. Only local paths are checked; other filesystems are not asked, because creating one is + * costly. If the filesystem cannot be asked, assume a checksum file exists. + */ + private[rules] def hasLocalChecksumFile(hadoopConf: Configuration, uri: URI): Boolean = { + val isLocal = Option(uri.getScheme).forall(_.equalsIgnoreCase("file")) + isLocal && scala.util + .Try { + val path = new org.apache.hadoop.fs.Path(uri) + path.getFileSystem(hadoopConf) match { + case fs: org.apache.hadoop.fs.ChecksumFileSystem => + fs.getRawFileSystem.exists(fs.getChecksumFile(path)) + case _ => false + } + } + .getOrElse(true) + } + + /** + * Why a native scan of `rootUris` must fall back because of the filesystem scheme or path, or + * None. A custom Hadoop FileSystem (e.g. registered via spark.hadoop.fs..impl) would + * surface at execution time as `Generic URL error: Unable to recognise URL "..."`, so those are + * declined and Spark's reader, which goes through the Hadoop FS API, handles the scan. Whether + * object_store recognizes a scheme is answered by the native layer itself + * (`NativeBase.isObjectStoreSchemeSupported`) rather than a hardcoded list, so the planner + * can't drift from object_store's actual support. Shared by the Parquet V1 and Text V2 gates; + * `scanName` starts each reason. + * + * EXCEPT schemes native routes through libhdfs (`hdfs` by default). A scheme the reader options + * add to that list is declined; see NativeConfig.resolveScanLibhdfsSchemes. + * `sessionLibhdfsSchemes` is the session's libhdfs list, and `hadoopConf` is the session's conf + * with the scan's reader options. + */ + private[rules] def rootPathFallbackReason( + scanName: String, + rootUris: Seq[URI], + sessionLibhdfsSchemes: Set[String], + hadoopConf: Configuration, + isBucketedScan: Boolean): Option[String] = { + val libhdfs = NativeConfig.resolveScanLibhdfsSchemes(sessionLibhdfsSchemes, hadoopConf) + val libhdfsSchemes = libhdfs.libhdfs + // Opt-in S3-compliant alias schemes (e.g. `blob`) from `fs.comet.s3Compliant.schemes`. Read + // from the Hadoop config, like the libhdfs list above, so `core-site.xml` is honored. The + // native object_store gate no longer claims aliases, so admit them here where the Hadoop + // config is available. + val s3CompliantSchemes = NativeConfig.resolveS3CompliantSchemes(hadoopConf) + + // Classify each root path once; see RootPathInfo. + val roots = classifyRootPaths(rootUris, libhdfsSchemes, s3CompliantSchemes) + + val scanOnlyLibhdfsSchemes = roots.flatMap(_.scheme).filter(libhdfs.scanOnly.contains).toSet + if (scanOnlyLibhdfsSchemes.nonEmpty) { + return Some( + s"$scanName does not support filesystem scheme(s) " + + s"${scanOnlyLibhdfsSchemes.toSeq.sorted.mkString(", ")} listed in " + + s"$COMET_LIBHDFS_SCHEMES_KEY only through reader options") + } + val unsupportedFsSchemes = roots.iterator.flatMap { root => + root.scheme.filter(_ => + !root.isLibhdfs && !isNativelyReadableScheme(root.uri, s3CompliantSchemes)) + }.toSet + if (unsupportedFsSchemes.nonEmpty) { + return Some(s"Unsupported filesystem schemes: ${unsupportedFsSchemes.mkString(", ")}") + } + // An early check on the root paths, with both scheme lists from the Hadoop conf as native + // reads them (not libhdfsSchemes above). The serde checks the listed files. + val multiStoreReason = CometScanUtils.multiStoreFallbackReason( + scanName, + roots.map(_.uri), + s3CompliantSchemes, + NativeConfig.resolveLibhdfsSchemes(hadoopConf), + isBucketedScan) + if (multiStoreReason.nonEmpty) { + return multiStoreReason + } + // A scheme object_store recognizes can still carry a path it rejects (e.g. a directory whose + // name contains a newline -> `%0A` in the URI), which native execution would hard-fail on. + // Only object_store-native schemes are probed: aliases (normalized to s3:// natively) and + // libhdfs schemes route elsewhere. Root paths only -- a rejected character deeper in the tree + // still fails at execution. + roots + .find(root => + root.scheme.isDefined && !root.isLibhdfs && !root.isAlias && + !objectStoreAcceptsPath(root.uri)) + .map(root => + s"$scanName cannot open path '${root.uri}': object_store " + + "rejects it (e.g. an unsupported character in the path)") + } + /** * Ask the native object_store parser whether it can build `url`: both the scheme and, via * `Path::from_url_path`, the key. On any native error (e.g. the library isn't loaded), assume diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometCsvNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometCsvNativeScanExec.scala index ce6fc589093..ade8484fb2d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometCsvNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometCsvNativeScanExec.scala @@ -23,9 +23,9 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.sql.catalyst.csv.CSVOptions 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.v2.BatchScanExec import org.apache.spark.sql.execution.datasources.v2.csv.CSVScan @@ -33,7 +33,6 @@ 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} @@ -60,8 +59,15 @@ case class CometCsvNativeScanExec( // The native operator holds the object store options, so the plan's text leaves it out. override def stringArgs: Iterator[Any] = Iterator(output) + // 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 = { - CometCsvNativeScanExec(nativeOp, output, originalPlan, serializedPlanOpt) + CometCsvNativeScanExec( + nativeOp, + output.map(QueryPlan.normalizeExpressions(_, output)), + if (originalPlan == null) null else originalPlan.doCanonicalize(), + SerializedPlan(None)) } override def equals(obj: Any): Boolean = { @@ -97,65 +103,31 @@ object CometCsvNativeScanExec extends CometOperatorSerde[CometBatchScanExec] { val timeZone = sessionState.conf.sessionLocalTimeZone new CSVOptions(csvScan.options.asScala.toMap, columnPruning, timeZone) } - val hadoopConf = - sessionState.newHadoopConfWithOptions(op.session.sparkContext.conf.getAll.toMap) - val s3CompliantSchemes = NativeConfig.resolveS3CompliantSchemes(hadoopConf) - val libhdfsSchemes = NativeConfig.resolveLibhdfsSchemes(hadoopConf) - val inputPartitions = op.inputPartitions.map(_.asInstanceOf[FilePartition]) - val multiStoreReason = CometScanUtils.multiStoreFallbackReason( + // The Hadoop conf comes from the scan's read options with their case preserved, matching + // Spark's CSVScan.createReaderFactory, so per-read settings like fs.s3a.endpoint reach native. + val input = CometScanUtils.prepareNativeFileScan( "Native CSV 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) => - // CometExecRule falls back to the wrapped Spark scan, so it carries the reason too. - 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 CSV 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 + op, + csvScan.options.asCaseSensitiveMap.asScala.toMap) match { + case Right(input) => input + case Left(reason) => + // CometExecRule falls back to the wrapped scan, so tag it too for the explain output. + withFallbackReason(op, reason) + withFallbackReason(op.wrapped, reason) + return None } val csvOptionsProto = csvOptions2Proto(options) val dataSchemaProto = schema2Proto(csvScan.dataSchema) - val readSchemaFieldNames = csvScan.readDataSchema.fieldNames - val projectionVector = csvScan.dataSchema.zipWithIndex - .filter { case (field, _) => - readSchemaFieldNames.contains(field.name) - } - .map(_._2.asInstanceOf[Integer]) val partitionSchemaProto = schema2Proto(csvScan.readPartitionSchema) - val partitionsProto = filePartitions.map(partition2Proto(_, csvScan.readPartitionSchema)) + val partitionsProto = + input.filePartitions.map(partition2Proto(_, csvScan.readPartitionSchema)) - csvScanBuilder.putAllObjectStoreOptions(objectStoreOptions.asJava) + csvScanBuilder.putAllObjectStoreOptions(input.objectStoreOptions.asJava) csvScanBuilder.setCsvOptions(csvOptionsProto) csvScanBuilder.addAllFilePartitions(partitionsProto.asJava) csvScanBuilder.addAllDataSchema(dataSchemaProto.asJava) - csvScanBuilder.addAllProjectionVector(projectionVector.asJava) + csvScanBuilder.addAllProjectionVector( + CometScanUtils.projectionVector(csvScan.dataSchema, csvScan.readDataSchema).asJava) csvScanBuilder.addAllPartitionSchema(partitionSchemaProto.asJava) Some(builder.setCsvScan(csvScanBuilder).build()) } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometScanUtils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometScanUtils.scala index 5c88b66e2e4..5ae88ae7e67 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometScanUtils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometScanUtils.scala @@ -25,8 +25,10 @@ import java.util.Locale import scala.collection.mutable import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, Literal} +import org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning import org.apache.spark.sql.execution.FileSourceScanExec import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionedFile} +import org.apache.spark.sql.types.StructType import org.apache.comet.objectstore.NativeConfig @@ -112,6 +114,88 @@ object CometScanUtils { groups.valuesIterator.map(_.toVector).toVector } + /** + * What the native CSV and Text scans serialize besides their format options: the object store + * options of every scheme the scan reads, and its file partitions, split so that none mixes + * object stores. + */ + case class NativeFileScanInput( + objectStoreOptions: Map[String, String], + filePartitions: Seq[FilePartition]) + + /** + * Prepares the part of a native file scan's serialization that does not depend on the format: + * the Hadoop conf from the read options, the object store options, and the per-store partition + * split. `Left` is the reason the scan must stay on Spark. `readOptions` is the scan's option + * map with the case of each key preserved, as Spark's reader factory builds the Hadoop conf. + * + * A read option can list a scheme for libhdfs that the session does not. Native reads through + * libhdfs with the executor classpath's configuration, not the read options, so a scheme that + * only the options route there cannot be read natively, and the scan declines it, as the native + * Parquet scan does. + */ + def prepareNativeFileScan( + scanName: String, + op: CometBatchScanExec, + readOptions: Map[String, String]): Either[String, NativeFileScanInput] = { + val sessionState = op.session.sessionState + val hadoopConf = sessionState.newHadoopConfWithOptions(readOptions) + val s3CompliantSchemes = NativeConfig.resolveS3CompliantSchemes(hadoopConf) + val scanLibhdfs = NativeConfig.resolveScanLibhdfsSchemes( + NativeConfig.resolveLibhdfsSchemes(sessionState.newHadoopConf()), + hadoopConf) + val libhdfsSchemes = scanLibhdfs.libhdfs + val inputPartitions = op.inputPartitions.map(_.asInstanceOf[FilePartition]) + val scanOnlyLibhdfsSchemes = inputPartitions.view + .flatMap(_.files.view.flatMap(file => NativeConfig.lowerScheme(file.pathUri))) + .filter(scanLibhdfs.scanOnly.contains) + .toSet + if (scanOnlyLibhdfsSchemes.nonEmpty) { + return Left( + s"$scanName does not support filesystem scheme(s) " + + s"${scanOnlyLibhdfsSchemes.toSeq.sorted.mkString(", ")} listed for libhdfs only " + + "through reader options") + } + val multiStoreReason = multiStoreFallbackReason( + scanName, + inputPartitions.view.flatMap(_.files.view.map(_.pathUri)), + s3CompliantSchemes, + libhdfsSchemes, + isBucketedScan = false) + if (multiStoreReason.nonEmpty) { + return Left(multiStoreReason.get) + } + // 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) => return Left(reason) + } + // Native planning reads each partition through one object store. + val filePartitions = 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]) { + return Left( + s"$scanName would split a partition that mixes object stores, but Spark planned the " + + "operators above it for the scan's own partitioning") + } + Right(NativeFileScanInput(objectStoreOptions, filePartitions)) + } + + /** Indexes into `dataSchema` of the columns in `readDataSchema`, in `dataSchema` order. */ + def projectionVector(dataSchema: StructType, readDataSchema: StructType): Seq[Integer] = { + val readFieldNames = readDataSchema.fieldNames + dataSchema.zipWithIndex + .filter { case (field, _) => readFieldNames.contains(field.name) } + .map(_._2.asInstanceOf[Integer]) + } + /** * Why a native scan of `uris` must fall back because of where they live, or None. The scan * forwards the settings of every scheme it reads, and translates alias settings to the diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometTextNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometTextNativeScanExec.scala new file mode 100644 index 00000000000..7ee18b56703 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometTextNativeScanExec.scala @@ -0,0 +1,150 @@ +/* + * 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.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.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 options = new TextOptions(textScan.options.asScala.toMap) + // The Hadoop conf comes from the scan's read options with their case preserved, matching + // Spark's TextScan.createReaderFactory, so per-read settings like fs.s3a.endpoint reach native. + val input = CometScanUtils.prepareNativeFileScan( + "Native Text scan", + op, + textScan.options.asCaseSensitiveMap.asScala.toMap) match { + case Right(input) => input + case Left(reason) => + // CometExecRule falls back to the wrapped scan, so tag it too for the explain output. + withFallbackReason(op, reason) + withFallbackReason(op.wrapped, reason) + return None + } + // Partition columns are declined by the scan gate, so the file partitions carry no values. + val partitionsProto = + input.filePartitions.map(partition2Proto(_, textScan.readPartitionSchema)) + + textScanBuilder.putAllObjectStoreOptions(input.objectStoreOptions.asJava) + textScanBuilder.setTextOptions(textOptions2Proto(options)) + textScanBuilder.addAllFilePartitions(partitionsProto.asJava) + textScanBuilder.addAllDataSchema(schema2Proto(textScan.dataSchema).asJava) + textScanBuilder.addAllProjectionVector( + CometScanUtils.projectionVector(textScan.dataSchema, textScan.readDataSchema).asJava) + Some(builder.setTextScan(textScanBuilder).build()) + } + + override def createExec(nativeOp: Operator, op: CometBatchScanExec): CometNativeExec = { + CometTextNativeScanExec(nativeOp, op.output, op.wrapped, SerializedPlan(None)) + } + + private def textOptions2Proto(options: TextOptions): OperatorOuterClass.TextOptions = { + val textOptionsBuilder = OperatorOuterClass.TextOptions.newBuilder() + textOptionsBuilder.setWholeText(options.wholeText) + // Send the exact separator bytes Spark computed (already encoded with the file's charset via + // the `encoding` option); the native reader splits on these bytes directly. + options.lineSeparatorInRead.foreach(bytes => + textOptionsBuilder.setLineSep(com.google.protobuf.ByteString.copyFrom(bytes))) + textOptionsBuilder.build() + } +} diff --git a/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala index 6906a39fbfc..40796d09633 100644 --- a/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala +++ b/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala @@ -19,7 +19,9 @@ package org.apache.comet.shims +import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path +import org.apache.hadoop.io.compress.CompressionCodecFactory import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile, RowIndexUtil} import org.apache.spark.sql.types.{DataType, StructType} @@ -30,6 +32,16 @@ object ShimFileFormat { // until they can be placed in the _metadata struct. val ROW_INDEX_TEMPORARY_COLUMN_NAME: String = FileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME + // Whether Spark's text/line readers would decompress this file. Before Spark 4.1 the codec is + // resolved purely through CompressionCodecFactory (the non-standard `.gzip`/`.zstd` extensions + // are read raw on these versions, so they are intentionally not treated as compressed here). + // Build the checker once per scan: the codec factory reads the codec registry from the + // configuration, so a scan over many files should not rebuild it for each path. + def compressedFileCheck(hadoopConf: Configuration): Path => Boolean = { + val factory = new CompressionCodecFactory(hadoopConf) + path => factory.getCodec(path) != null + } + def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = RowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema) diff --git a/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala index 9506a851c08..c26df6877d4 100644 --- a/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala +++ b/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala @@ -19,6 +19,9 @@ package org.apache.comet.shims +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.hadoop.io.compress.CompressionCodecFactory import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat @@ -30,6 +33,16 @@ object ShimFileFormat { // until they can be placed in the _metadata struct. val ROW_INDEX_TEMPORARY_COLUMN_NAME = ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME + // Whether Spark's text/line readers would decompress this file. Before Spark 4.1 the codec is + // resolved purely through CompressionCodecFactory (the non-standard `.gzip`/`.zstd` extensions + // are read raw on these versions, so they are intentionally not treated as compressed here). + // Build the checker once per scan: the codec factory reads the codec registry from the + // configuration, so a scan over many files should not rebuild it for each path. + def compressedFileCheck(hadoopConf: Configuration): Path => Boolean = { + val factory = new CompressionCodecFactory(hadoopConf) + path => factory.getCodec(path) != null + } + def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = ParquetRowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema) diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-4.0/org/apache/comet/shims/ShimFileFormat.scala index 9506a851c08..c26df6877d4 100644 --- a/spark/src/main/spark-4.0/org/apache/comet/shims/ShimFileFormat.scala +++ b/spark/src/main/spark-4.0/org/apache/comet/shims/ShimFileFormat.scala @@ -19,6 +19,9 @@ package org.apache.comet.shims +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.hadoop.io.compress.CompressionCodecFactory import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat @@ -30,6 +33,16 @@ object ShimFileFormat { // until they can be placed in the _metadata struct. val ROW_INDEX_TEMPORARY_COLUMN_NAME = ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME + // Whether Spark's text/line readers would decompress this file. Before Spark 4.1 the codec is + // resolved purely through CompressionCodecFactory (the non-standard `.gzip`/`.zstd` extensions + // are read raw on these versions, so they are intentionally not treated as compressed here). + // Build the checker once per scan: the codec factory reads the codec registry from the + // configuration, so a scan over many files should not rebuild it for each path. + def compressedFileCheck(hadoopConf: Configuration): Path => Boolean = { + val factory = new CompressionCodecFactory(hadoopConf) + path => factory.getCodec(path) != null + } + def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = ParquetRowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema) diff --git a/spark/src/main/spark-4.1/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-4.1/org/apache/comet/shims/ShimFileFormat.scala index 9506a851c08..046560f1185 100644 --- a/spark/src/main/spark-4.1/org/apache/comet/shims/ShimFileFormat.scala +++ b/spark/src/main/spark-4.1/org/apache/comet/shims/ShimFileFormat.scala @@ -19,6 +19,9 @@ package org.apache.comet.shims +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.spark.io.HadoopCodecStreams import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat @@ -30,6 +33,13 @@ object ShimFileFormat { // until they can be placed in the _metadata struct. val ROW_INDEX_TEMPORARY_COLUMN_NAME = ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME + // Whether Spark's text/line readers would decompress this file. Spark 4.1+ resolves the codec + // through HadoopCodecStreams, which also treats the non-standard `.gzip` and `.zstd` extensions + // as compressed (CompressionCodecFactory alone misses those). Keep this in lockstep with the + // Spark version so the native text scan declines exactly the files Spark would decompress. + def compressedFileCheck(hadoopConf: Configuration): Path => Boolean = + path => HadoopCodecStreams.getDecompressionCodec(hadoopConf, path).isDefined + def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = ParquetRowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema) diff --git a/spark/src/main/spark-4.2/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-4.2/org/apache/comet/shims/ShimFileFormat.scala index 4aee5671de8..8e96e59697b 100644 --- a/spark/src/main/spark-4.2/org/apache/comet/shims/ShimFileFormat.scala +++ b/spark/src/main/spark-4.2/org/apache/comet/shims/ShimFileFormat.scala @@ -19,6 +19,9 @@ package org.apache.comet.shims +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.spark.io.HadoopCodecStreams import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat @@ -30,6 +33,13 @@ object ShimFileFormat { // until they can be placed in the _metadata struct. val ROW_INDEX_TEMPORARY_COLUMN_NAME = ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME + // Whether Spark's text/line readers would decompress this file. Spark 4.1+ resolves the codec + // through HadoopCodecStreams, which also treats the non-standard `.gzip` and `.zstd` extensions + // as compressed (CompressionCodecFactory alone misses those). Keep this in lockstep with the + // Spark version so the native text scan declines exactly the files Spark would decompress. + def compressedFileCheck(hadoopConf: Configuration): Path => Boolean = + path => HadoopCodecStreams.getDecompressionCodec(hadoopConf, path).isDefined + def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = ParquetRowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema) diff --git a/spark/src/test/resources/test-data/text-test-1.txt b/spark/src/test/resources/test-data/text-test-1.txt new file mode 100644 index 00000000000..f10289c3616 --- /dev/null +++ b/spark/src/test/resources/test-data/text-test-1.txt @@ -0,0 +1,6 @@ +hello world +foo,bar,baz +"quoted stays literal" +café münchen naïve +日本語テスト +1234567890 diff --git a/spark/src/test/scala/org/apache/comet/csv/CometCsvNativeReadSuite.scala b/spark/src/test/scala/org/apache/comet/csv/CometCsvNativeReadSuite.scala index eeb809764be..7cdfcb99187 100644 --- a/spark/src/test/scala/org/apache/comet/csv/CometCsvNativeReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/csv/CometCsvNativeReadSuite.scala @@ -176,4 +176,46 @@ class CometCsvNativeReadSuite extends CometTestBase { } } } + + test("Native csv read - a libhdfs scheme set only by a read option falls back") { + withTempPath { dir => + spark.range(5).write.csv(dir.toString) + withSQLConf( + CometConf.COMET_CSV_V2_NATIVE_ENABLED.key -> "true", + SQLConf.USE_V1_SOURCE_LIST.key -> "") { + // The option routes `file` paths through libhdfs for this read only. Native cannot read + // them that way, so the scan must stay on Spark, which ignores the option. + val df = spark.read + .schema("id LONG") + .option("fs.comet.libhdfs.schemes", "file,hdfs") + .csv(dir.toString) + val (_, plan) = + checkSparkAnswerAndFallbackReason(df, "listed for libhdfs only through reader options") + assert( + collect(plan) { case s: CometCsvNativeScanExec => s }.isEmpty, + s"Expected Spark's CSV reader but found CometCsvNativeScanExec. Plan:\n$plan") + } + } + } + + test("Native csv read - scans of the same file compare equal") { + withTempPath { dir => + spark.range(5).write.csv(dir.toString) + withSQLConf( + CometConf.COMET_CSV_V2_NATIVE_ENABLED.key -> "true", + SQLConf.USE_V1_SOURCE_LIST.key -> "", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + def scan: CometCsvNativeScanExec = + spark.read + .schema("id LONG") + .csv(dir.toString) + .queryExecution + .executedPlan + .collect { case s: CometCsvNativeScanExec => s } + .head + // Equal canonical forms let Spark reuse the exchanges and stages of a repeated scan. + assert(scan.sameResult(scan)) + } + } + } } diff --git a/spark/src/test/scala/org/apache/comet/objectstore/AbfsScanFallbackSuite.scala b/spark/src/test/scala/org/apache/comet/objectstore/AbfsScanFallbackSuite.scala index 5386697de1b..6751d5aaa77 100644 --- a/spark/src/test/scala/org/apache/comet/objectstore/AbfsScanFallbackSuite.scala +++ b/spark/src/test/scala/org/apache/comet/objectstore/AbfsScanFallbackSuite.scala @@ -80,6 +80,18 @@ class AbfsScanFallbackSuite extends CometTestBase { .write .csv(localDir(dir)) + private def writeText(dir: String, from: Int): Unit = + spark + .range(from.toLong, from + 10L) + .selectExpr("cast(id as string) as value") + .write + .text(localDir(dir)) + + // Text is left off the V1 source list so the scan plans as a V2 TextScan. + private val textConfs = Seq( + CometConf.COMET_TEXT_V2_NATIVE_ENABLED.key -> "true", + SQLConf.USE_V1_SOURCE_LIST.key -> "avro,csv,json,kafka,orc,parquet") + private val csvConfs = Seq(CometConf.COMET_CSV_V2_NATIVE_ENABLED.key -> "true", SQLConf.USE_V1_SOURCE_LIST.key -> "") @@ -181,4 +193,24 @@ class AbfsScanFallbackSuite extends CometTestBase { checkSparkAnswerAndFallbackReason(df, multipleAuthoritiesReason) } } + + // The serializer declines these layouts, so the scan must fall back to Spark's row-producing + // scan rather than stay a columnar CometBatchScanExec (which fails with a ClassCastException). + test("native Text scan over two accounts with different keys falls back to Spark") { + writeText("text-a", 0) + writeText("text-b", 100) + withSQLConf(textConfs ++ accountKeys("accta" -> "a2V5LWE=", "acctb" -> "a2V5LWI="): _*) { + val df = spark.read.text(abfss("data", "accta", "text-a"), abfss("data", "acctb", "text-b")) + checkSparkAnswerAndFallbackReason(df, multipleAuthoritiesReason) + } + } + + test("native Text scan over two containers of one account falls back to Spark") { + writeText("text-c", 0) + writeText("text-d", 100) + withSQLConf(textConfs ++ accountKeys("accta" -> "a2V5LWE="): _*) { + val df = spark.read.text(abfss("data", "accta", "text-c"), abfss("logs", "accta", "text-d")) + checkSparkAnswerAndFallbackReason(df, multipleAuthoritiesReason) + } + } } diff --git a/spark/src/test/scala/org/apache/comet/text/CometTextNativeReadSuite.scala b/spark/src/test/scala/org/apache/comet/text/CometTextNativeReadSuite.scala new file mode 100644 index 00000000000..ab21c8cc659 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/text/CometTextNativeReadSuite.scala @@ -0,0 +1,710 @@ +/* + * 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 java.util.concurrent.atomic.AtomicLong +import java.util.zip.GZIPOutputStream + +import org.scalatest.time.{Millis, Seconds, Span} + +import org.apache.hadoop.fs.{FileSystem, Path} +import org.apache.hadoop.io.compress.GzipCodec +import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} +import org.apache.spark.sql.{CometTestBase, DataFrame} +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, ExtendedExplainInfo} +import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus + +/** A gzip codec whose file extension holds a space, to test codec lookup on decoded paths. */ +class SpacedGzipCodec extends GzipCodec { + override def getDefaultExtension: String = ".packed gz" +} + +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()) + // The blanket V2 input_file_name gate (shared by all native V2 scans) emits this. + checkSparkAnswerAndFallbackReason(df, "is not compatible with input_file_name") + } + } + + test("native text read - fallback when native exec is disabled") { + withSQLConf( + CometConf.COMET_TEXT_V2_NATIVE_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "false", + SQLConf.USE_V1_SOURCE_LIST.key -> "") { + checkSparkAnswerAndFallbackReason( + spark.read.text(TEST_TEXT_PATH), + s"Native Text scan requires ${CometConf.COMET_EXEC_ENABLED.key} to be enabled") + } + } + + test("native text read - wholetext") { + withTextConf { + val df = spark.read.option("wholetext", "true").text(TEST_TEXT_PATH) + checkSparkAnswerAndOperator(df) + } + } + + test("native text read - custom lineSep") { + withTempDir { dir => + val file = new File(dir, "custom-linesep.txt") + Files.write(file.toPath, "a|b|café|日本".getBytes(StandardCharsets.UTF_8)) + withTextConf { + val df = spark.read.option("lineSep", "|").text(file.getAbsolutePath) + checkSparkAnswerAndOperator(df) + } + } + } + + test("native text read - custom lineSep with trailing separator") { + withTempDir { dir => + val file = new File(dir, "trailing-linesep.txt") + // A trailing separator must not produce an extra empty row (matches Spark, SPARK-23577). + Files.write(file.toPath, "a|b|c|".getBytes(StandardCharsets.UTF_8)) + withTextConf { + val df = spark.read.option("lineSep", "|").text(file.getAbsolutePath) + checkSparkAnswerAndOperator(df) + } + } + } + + test("native text read - non-ASCII multibyte custom lineSep") { + withTempDir { dir => + val file = new File(dir, "jp-linesep.txt") + // Japanese ideographic comma (U+3002, 3 UTF-8 bytes) as the record separator. + Files.write(file.toPath, "a。café。日本".getBytes(StandardCharsets.UTF_8)) + withTextConf { + val df = spark.read.option("lineSep", "。").text(file.getAbsolutePath) + checkSparkAnswerAndOperator(df) + } + } + } + + test("native text read - CRLF and CR line endings") { + withTempDir { dir => + val file = new File(dir, "crlf.txt") + // Mix of \r\n, \r, and \n; the default reader splits on all three. + Files.write(file.toPath, "a\r\nb\rc\nd".getBytes(StandardCharsets.UTF_8)) + withTextConf { + checkSparkAnswerAndOperator(spark.read.text(file.getAbsolutePath)) + } + } + } + + test("native text read - invalid UTF-8 bytes match Spark on collect") { + withTempDir { dir => + val file = new File(dir, "invalid-utf8.txt") + // 0xFF is not valid UTF-8; both readers render it as the replacement char on collect. + Files.write(file.toPath, Array[Byte]('a', 0xff.toByte, 'b', '\n', 'c')) + withTextConf { + checkSparkAnswerAndOperator(spark.read.text(file.getAbsolutePath)) + } + } + } + + test("native text read - wholetext count") { + withTextConf { + val df = spark.read + .option("wholetext", "true") + .text(TEST_TEXT_PATH) + .selectExpr("count(*) as c") + checkSparkAnswerAndOperator(df) + } + } + + test("native text read - falls back for files split into byte ranges") { + withTempDir { dir => + val file = new File(dir, "many-lines.txt") + val content = (1 to 2000).map(i => s"line-$i-日本テスト").mkString("\n") + Files.write(file.toPath, content.getBytes(StandardCharsets.UTF_8)) + withSQLConf( + CometConf.COMET_TEXT_V2_NATIVE_ENABLED.key -> "true", + SQLConf.USE_V1_SOURCE_LIST.key -> "", + // Force the file to be split into several byte-range partitions. + SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1024") { + val df = spark.read.text(file.getAbsolutePath) + assert(df.rdd.getNumPartitions > 1, "test requires the file to actually split") + checkSparkAnswerAndFallbackReason( + df, + "Comet native Text scan does not support files split into byte ranges") + } + } + } + + test("native text read - empty file") { + withTempDir { dir => + val file = new File(dir, "empty.txt") + Files.write(file.toPath, Array.emptyByteArray) + withTextConf { + checkSparkAnswerAndOperator(spark.read.text(file.getAbsolutePath)) + } + } + } + + test("native text read - directory of multiple small files") { + withTempDir { dir => + Files.write(new File(dir, "a.txt").toPath, "alpha\nbeta".getBytes(StandardCharsets.UTF_8)) + Files.write(new File(dir, "b.txt").toPath, "日本\ngamma".getBytes(StandardCharsets.UTF_8)) + Files.write(new File(dir, "c.txt").toPath, Array.emptyByteArray) + withTextConf { + checkSparkAnswerAndOperator(spark.read.text(dir.getAbsolutePath)) + } + } + } + + test("native text read - wholetext over a directory of files") { + withTempDir { dir => + Files.write( + new File(dir, "a.txt").toPath, + "line1\nline2\n".getBytes(StandardCharsets.UTF_8)) + Files.write(new File(dir, "b.txt").toPath, "".getBytes(StandardCharsets.UTF_8)) + withTextConf { + // One row per file, including one empty row for the empty file. + checkSparkAnswerAndOperator( + spark.read.option("wholetext", "true").text(dir.getAbsolutePath)) + } + } + } + + test("native text read - wholetext falls back above the file size ceiling") { + withTempDir { dir => + val file = new File(dir, "big-wholetext.txt") + Files.write(file.toPath, ("x" * 100).getBytes(StandardCharsets.UTF_8)) + withTextConf { + // The same ceiling that bounds line mode applies to a wholetext file. + withSQLConf(CometConf.COMET_SCAN_TEXT_MAX_FILE_SIZE.key -> "16") { + val df = spark.read.option("wholetext", "true").text(file.getAbsolutePath) + checkSparkAnswerAndFallbackReason(df, "does not read files larger than") + } + } + } + } + + test("native text read - maxPartitionBytes does not limit wholetext") { + withTempDir { dir => + val file = new File(dir, "wholetext-small-partitions.txt") + Files.write(file.toPath, ("x" * 100).getBytes(StandardCharsets.UTF_8)) + withTextConf { + // maxPartitionBytes only splits line-mode files. Lowering it for parallelism must not + // take a wholetext file under the size ceiling off the native path. + withSQLConf(SQLConf.FILES_MAX_PARTITION_BYTES.key -> "16") { + val df = spark.read.option("wholetext", "true").text(file.getAbsolutePath) + checkSparkAnswerAndOperator(df, Seq(classOf[CometTextNativeScanExec])) + } + } + } + } + + test("native text read - strips a leading UTF-8 BOM in line mode") { + withTempDir { dir => + val file = new File(dir, "bom.txt") + val bom = Array[Byte](0xef.toByte, 0xbb.toByte, 0xbf.toByte) + Files.write(file.toPath, bom ++ "hello\nworld".getBytes(StandardCharsets.UTF_8)) + withTextConf { + // Spark's line reader strips the BOM; native must match. + checkSparkAnswerAndOperator(spark.read.text(file.getAbsolutePath)) + } + } + } + + test("native text read - fallback for ignoreCorruptFiles / ignoreMissingFiles") { + withTempDir { dir => + val file = new File(dir, "a.txt") + Files.write(file.toPath, "a\nb".getBytes(StandardCharsets.UTF_8)) + withTextConf { + checkSparkAnswerAndFallbackReason( + spark.read.option("ignoreCorruptFiles", "true").text(file.getAbsolutePath), + "Comet native Text scan does not support ignoreCorruptFiles") + checkSparkAnswerAndFallbackReason( + spark.read.option("ignoreMissingFiles", "true").text(file.getAbsolutePath), + "Comet native Text scan does not support ignoreMissingFiles") + } + } + } + + test("native text read - fallback for compressed (.gz) files") { + withTempDir { dir => + val file = new File(dir, "data.txt.gz") + val out = new java.util.zip.GZIPOutputStream(new java.io.FileOutputStream(file)) + out.write("a\nb\nc".getBytes(StandardCharsets.UTF_8)) + out.close() + withTextConf { + val df = spark.read.text(file.getAbsolutePath) + checkSparkAnswerAndFallbackReason(df, "Comet does not support compressed text files") + } + } + } + + test("native text read - fallback for .gzip/.zstd compressed files on Spark 4.1+") { + // Spark 4.1+ decompresses the non-standard .gzip/.zstd extensions (via HadoopCodecStreams); + // CompressionCodecFactory alone misses them, so the version-aware shim must catch them. + assume(isSpark41Plus, "Spark 4.1+ treats .gzip/.zstd as compressed") + withTempDir { dir => + val gzip = new File(dir, "data.txt.gzip") + val gzipOut = new java.util.zip.GZIPOutputStream(new java.io.FileOutputStream(gzip)) + gzipOut.write("a\nb\nc".getBytes(StandardCharsets.UTF_8)) + gzipOut.close() + val zstd = new File(dir, "data.txt.zstd") + val zstdOut = new com.github.luben.zstd.ZstdOutputStream(new java.io.FileOutputStream(zstd)) + zstdOut.write("a\nb\nc".getBytes(StandardCharsets.UTF_8)) + zstdOut.close() + withTextConf { + checkSparkAnswerAndFallbackReason( + spark.read.text(gzip.getAbsolutePath), + "Comet does not support compressed text files") + checkSparkAnswerAndFallbackReason( + spark.read.text(zstd.getAbsolutePath), + "Comet does not support compressed text files") + } + } + } + + test("native text read - custom lineSep equal to a BOM is not stripped") { + withTempDir { dir => + val file = new File(dir, "bom-linesep.txt") + val bom = "\uFEFF" + // With lineSep=BOM, Spark leaves the BOM as the separator: `BOM a BOM b` -> ["", "a", "b"]. + Files.write(file.toPath, s"${bom}a${bom}b".getBytes(StandardCharsets.UTF_8)) + withTextConf { + checkSparkAnswerAndOperator(spark.read.option("lineSep", bom).text(file.getAbsolutePath)) + } + } + } + + test("native text read - fallback for a self-overlapping custom lineSep") { + withTempDir { dir => + val file = new File(dir, "overlap.txt") + // "aab" overlaps itself (its "aa" prefix can restart a match). Hadoop's readCustomLine, which + // Spark uses, can split such input differently from the native reader's exact substring scan, + // so Comet must fall back. Spark's own result is the reference. + Files.write(file.toPath, "aaabz".getBytes(StandardCharsets.UTF_8)) + withTextConf { + checkSparkAnswerAndFallbackReason( + spark.read.option("lineSep", "aab").text(file.getAbsolutePath), + "self-overlapping line separator") + } + } + } + + test("native text read - self-overlapping lineSep that straddles the input buffer falls back") { + withTempDir { dir => + val file = new File(dir, "straddle.txt") + // `aab` is found by Hadoop's reader and by an exact search at different places when it + // falls on a read-buffer boundary, so the native scan must stay off. With a 4096-byte buffer + // Spark reads this file as one row; an exact search would split it into two. + val content = "x" * 4092 + "aaaaabTAIL" + Files.write(file.toPath, content.getBytes(StandardCharsets.UTF_8)) + withTextConf { + val df = spark.read + .option("lineSep", "aab") + .option("io.file.buffer.size", "4096") + .text(file.getAbsolutePath) + checkSparkAnswerAndFallbackReason(df, "self-overlapping line separator") + } + } + } + + test("native text read - BOM-only file matches Spark for default and BOM separators") { + withTempDir { dir => + val f = new File(dir, "bom-only.txt") + Files.write(f.toPath, "\uFEFF".getBytes(StandardCharsets.UTF_8)) + withTextConf { + // Default separator: the BOM is stripped, leaving no rows. + checkSparkAnswerAndOperator(spark.read.text(f.getAbsolutePath)) + // lineSep=BOM: the BOM is the (trailing) separator, leaving one empty row. + checkSparkAnswerAndOperator( + spark.read.option("lineSep", "\uFEFF").text(f.getAbsolutePath)) + } + } + } + + // Hadoop's LineRecordReader strips a leading BOM from the first record whatever the separator. + test("native text read - custom lineSep still strips a leading BOM from the first record") { + withTempDir { dir => + val bom = "\uFEFF" + withTextConf { + // `BOM a|b` -> ["a", "b"] + val ab = new File(dir, "bom-a-b.txt") + Files.write(ab.toPath, s"${bom}a|b".getBytes(StandardCharsets.UTF_8)) + checkSparkAnswerAndOperator(spark.read.option("lineSep", "|").text(ab.getAbsolutePath)) + // `BOM|x` -> ["", "x"] + val bx = new File(dir, "bom-sep-x.txt") + Files.write(bx.toPath, s"$bom|x".getBytes(StandardCharsets.UTF_8)) + checkSparkAnswerAndOperator(spark.read.option("lineSep", "|").text(bx.getAbsolutePath)) + } + } + } + + test("native text read - BOM-only file with a custom lineSep has no rows") { + withTempDir { dir => + val f = new File(dir, "bom-only-pipe.txt") + Files.write(f.toPath, "\uFEFF".getBytes(StandardCharsets.UTF_8)) + withTextConf { + checkSparkAnswerAndOperator(spark.read.option("lineSep", "|").text(f.getAbsolutePath)) + } + } + } + + test("native text read - fallback for line.maxlength") { + withTempDir { dir => + val file = new File(dir, "lines.txt") + Files.write(file.toPath, "ab\nabcdefghij\nxy".getBytes(StandardCharsets.UTF_8)) + withTextConf { + // As a per-read option the key reaches the scan's Hadoop conf (stripped), so Spark's line + // reader skips lines >= this length; the native reader does not, so Comet must fall back. + checkSparkAnswerAndFallbackReason( + spark.read + .option("mapreduce.input.linerecordreader.line.maxlength", "5") + .text(file.getAbsolutePath), + "mapreduce.input.linerecordreader.line.maxlength") + } + } + } + + test( + "native text read - line.maxlength from the Hadoop conf is not overridden by a differently " + + "cased read option") { + withTempDir { dir => + val file = new File(dir, "lines-case.txt") + Files.write(file.toPath, "ab\nabcdefghij\nxy".getBytes(StandardCharsets.UTF_8)) + val key = "mapreduce.input.linerecordreader.line.maxlength" + val hadoopConf = spark.sparkContext.hadoopConfiguration + hadoopConf.set(key, "5") + try { + withTextConf { + // Hadoop keys are case sensitive, so this option does not override the limit of 5 and + // Spark skips the long line. The gate must see the option with its case preserved, or it + // reads an unlimited length and the native scan returns the long line too. + checkSparkAnswerAndFallbackReason( + spark.read + .option("mapreduce.input.linerecordreader.line.maxLength", "2147483647") + .text(file.getAbsolutePath), + key) + } + } finally { + hadoopConf.unset(key) + } + } + } + + test("native text read - fallback for files larger than the size ceiling") { + withTempDir { dir => + val file = new File(dir, "big.txt") + Files.write(file.toPath, ("x\n" * 1000).getBytes(StandardCharsets.UTF_8)) + withSQLConf( + CometConf.COMET_TEXT_V2_NATIVE_ENABLED.key -> "true", + SQLConf.USE_V1_SOURCE_LIST.key -> "avro,csv,json,kafka,orc,parquet", + // Tiny ceiling so the ~2KB file exceeds it. + CometConf.COMET_SCAN_TEXT_MAX_FILE_SIZE.key -> "100") { + checkSparkAnswerAndFallbackReason( + spark.read.text(file.getAbsolutePath), + "does not read files larger than") + } + } + } + + test("native text read - partition column fallback: value stays native, select * falls back") { + withTempDir { dir => + // Partitioned layout: /p=1/a.txt + val part = new File(dir, "p=1") + part.mkdirs() + Files.write(new File(part, "a.txt").toPath, "x\ny".getBytes(StandardCharsets.UTF_8)) + withTextConf { + // Selecting only `value` has no partition column -> stays native. + checkSparkAnswerAndOperator(spark.read.text(dir.getAbsolutePath).select("value")) + // Selecting the partition column `p` -> falls back to Spark. + checkSparkAnswerAndFallbackReason( + spark.read.text(dir.getAbsolutePath).selectExpr("value", "p"), + "Comet does not support partition columns in native Text scans") + } + } + } + + test("native text read - multiple unsplit files across separate partitions") { + withTempDir { dir => + (1 to 6).foreach { i => + Files.write( + new File(dir, s"f$i.txt").toPath, + s"a$i\nb$i".getBytes(StandardCharsets.UTF_8)) + } + withSQLConf( + CometConf.COMET_TEXT_V2_NATIVE_ENABLED.key -> "true", + SQLConf.USE_V1_SOURCE_LIST.key -> "avro,csv,json,kafka,orc,parquet", + // A large open cost keeps each small file in its own partition, exercising + // file_partitions[self.partition] beyond index 0. + SQLConf.FILES_MAX_PARTITION_BYTES.key -> "8", + SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "1073741824") { + val df = spark.read.text(dir.getAbsolutePath) + assert(df.rdd.getNumPartitions > 1, "test requires files in separate partitions") + checkSparkAnswerAndOperator(df) + } + } + } + + test("native text read - broadcast join with a text build side runs natively") { + withTempDir { dir => + val lookup = new File(dir, "allow.txt") + Files.write(lookup.toPath, "k1\nk2\nk3".getBytes(StandardCharsets.UTF_8)) + withSQLConf( + CometConf.COMET_TEXT_V2_NATIVE_ENABLED.key -> "true", + SQLConf.USE_V1_SOURCE_LIST.key -> "avro,csv,json,kafka,orc,parquet") { + withParquetTable((1 to 100).map(i => (i, s"k${i % 5}")), "facts") { + val lookupDf = spark.read.text(lookup.getAbsolutePath) + val facts = spark.table("facts") + // Small text build side broadcast-joined to a native Parquet probe side. + val joined = facts.join(lookupDf.hint("broadcast"), facts("_2") === lookupDf("value")) + checkSparkAnswerAndOperator(joined) + } + } + } + } + + // Plan only: asserts the scan stays on Spark and records why, without running a query that + // Spark itself would fail. + private def assertFallsBack(df: DataFrame, reason: String): Unit = { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val plan = df.queryExecution.executedPlan + assert(plan.collect { case s: CometTextNativeScanExec => s }.isEmpty, plan) + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert(reasons.exists(_.contains(reason)), reasons) + } + } + + test("native text read - falls back for a local file with a Hadoop checksum file") { + withTempDir { dir => + val hadoopConf = spark.sparkContext.hadoopConfiguration + val path = new Path(new File(dir, "checked.txt").getAbsolutePath) + // Hadoop's local filesystem writes a `.checked.txt.crc` next to the data. + val out = FileSystem.getLocal(hadoopConf).create(path) + out.write("one\n".getBytes(StandardCharsets.UTF_8)) + out.close() + assert(new File(dir, ".checked.txt.crc").exists()) + // Replace the data without updating the checksum: Spark raises a ChecksumException on read. + Files.write(new File(dir, "checked.txt").toPath, "two\n".getBytes(StandardCharsets.UTF_8)) + withTextConf { + assertFallsBack(spark.read.text(path.toString), "Hadoop checksum") + } + } + } + + test("native text read - codec lookup uses the decoded file name") { + withTempDir { dir => + val file = new File(dir, "small.packed gz") + val out = new GZIPOutputStream(Files.newOutputStream(file.toPath)) + out.write("one\ntwo\n".getBytes(StandardCharsets.UTF_8)) + out.close() + val hadoopConf = spark.sparkContext.hadoopConfiguration + val key = "io.compression.codecs" + val previous = hadoopConf.get(key) + hadoopConf.set(key, classOf[SpacedGzipCodec].getName) + try { + withTextConf { + // The URL-encoded name `small.packed%20gz` hides the registered `.packed gz` suffix. + checkSparkAnswerAndFallbackReason( + spark.read.text(file.getAbsolutePath), + "compressed text files") + } + } finally { + if (previous == null) hadoopConf.unset(key) else hadoopConf.set(key, previous) + } + } + } + + test("native text read - a compressed file in a pruned partition does not block native") { + withTempDir { dir => + val hot = new File(dir, "p=hot") + val cold = new File(dir, "p=cold") + hot.mkdirs() + cold.mkdirs() + Files.write(new File(hot, "a.txt").toPath, "alpha\n".getBytes(StandardCharsets.UTF_8)) + val out = new GZIPOutputStream(Files.newOutputStream(new File(cold, "b.txt.gz").toPath)) + out.write("beta\n".getBytes(StandardCharsets.UTF_8)) + out.close() + withTextConf { + val df = spark.read.text(dir.getAbsolutePath).where("p = 'hot'").select("value") + checkSparkAnswerAndOperator(df, Seq(classOf[CometTextNativeScanExec])) + } + } + } + + test("native text read - a listed file whose name object_store rejects falls back") { + withTempDir { dir => + // Legal for Hadoop, rejected by object_store. The directory root alone looks fine. + Files.write(new File(dir, "a\tb.txt").toPath, "expected\n".getBytes(StandardCharsets.UTF_8)) + withTextConf { + checkSparkAnswerAndFallbackReason( + spark.read.text(dir.getAbsolutePath), + "cannot open path") + } + } + } + + test("native text read - falls back for mandatory file-open options in line mode") { + withTempDir { dir => + val file = new File(dir, "lines.txt") + Files.write(file.toPath, "alpha\nbeta\n".getBytes(StandardCharsets.UTF_8)) + withTextConf { + assertFallsBack( + spark.read + .option("mapreduce.job.input.file.must.fs.s3a.readahead.range", "64K") + .text(file.getAbsolutePath), + "mapreduce.job.input.file.must") + } + } + } + + test("native text read - wholetext ignores line.maxlength") { + withTempDir { dir => + val file = new File(dir, "whole.txt") + Files.write(file.toPath, "alpha\nbeta".getBytes(StandardCharsets.UTF_8)) + val key = "mapreduce.input.linerecordreader.line.maxlength" + val hadoopConf = spark.sparkContext.hadoopConfiguration + try { + withTextConf { + // Spark's whole-file reader never reads this setting, so neither a finite limit nor a + // value that is not an integer may fail or fall back a wholetext query. + Seq("5", "64k").foreach { value => + hadoopConf.set(key, value) + checkSparkAnswerAndOperator( + spark.read.option("wholetext", "true").text(file.getAbsolutePath)) + } + } + } finally { + hadoopConf.unset(key) + } + } + } + + test("native text read - scans of the same file compare equal") { + withTempDir { dir => + val file = new File(dir, "lookup.txt") + Files.write(file.toPath, "a\nb\nc\n".getBytes(StandardCharsets.UTF_8)) + withTextConf { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + def scan: CometTextNativeScanExec = + spark.read + .text(file.getAbsolutePath) + .queryExecution + .executedPlan + .collect { case s: CometTextNativeScanExec => s } + .head + // Equal canonical forms let Spark reuse the exchanges and stages of a repeated lookup. + assert(scan.sameResult(scan)) + } + } + } + } + + test("native text read - reports input bytes and records to the task metrics") { + withTempDir { dir => + (0 until 4).foreach { i => + Files.write( + new File(dir, s"part-$i.txt").toPath, + "a\nb\nc\n".getBytes(StandardCharsets.UTF_8)) + } + val bytes = new AtomicLong(0) + val records = new AtomicLong(0) + val listener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { + if (taskEnd.taskMetrics != null) { + bytes.addAndGet(taskEnd.taskMetrics.inputMetrics.bytesRead) + records.addAndGet(taskEnd.taskMetrics.inputMetrics.recordsRead) + } + } + } + spark.sparkContext.addSparkListener(listener) + try { + withTextConf { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = spark.read.text(dir.getAbsolutePath) + // Run only the native query: a Spark pass would add the same bytes and hide a + // missing metric. + assert(df.queryExecution.executedPlan.collect { case s: CometTextNativeScanExec => + s + }.nonEmpty) + assert(df.collect().length == 12) + } + } + // The listener bus is asynchronous, so wait for the four tasks' totals. + eventually(timeout(Span(10, Seconds)), interval(Span(100, Millis))) { + assert(bytes.get() == 24L, s"bytesRead=${bytes.get()}") + assert(records.get() == 12L, s"recordsRead=${records.get()}") + } + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + } +}