parthchandra commented on code in PR #6784:
URL: https://github.com/apache/datafusion-comet/pull/6784#discussion_r4253098327
##########
spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala:
##########
@@ -458,6 +459,147 @@ 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 read a thread-local that
Spark's FileScanRDD
+ // sets; the native DataFusion scan does not, so they would return ""
/ -1. Same guard as
+ // the native Parquet path.
+ if (plan.exists(node =>
+ node.expressions.exists(_.exists {
+ case _: InputFileName | _: InputFileBlockStart | _:
InputFileBlockLength => true
+ case _ => false
+ }))) {
+ return withFallbackReason(
+ scanExec,
+ "Native Text scan is not compatible with input_file_name, " +
+ "input_file_block_start, or input_file_block_length")
+ }
+ val fallbackReasons = new ListBuffer[String]()
+ val schemaSupported =
+ CometBatchScanExec.isSchemaSupported(scan.readDataSchema,
fallbackReasons)
+ if (!schemaSupported) {
+ fallbackReasons += s"Schema ${scan.readDataSchema} is not supported"
+ }
+ // Spark's text source allows only a single data column
(TextScan.verifyReadSchema). A
+ // user-forced multi-column read schema is invalid; fall back so Spark
raises its structured
+ // AnalysisException rather than letting the native reader fail with
an opaque error.
+ if (scan.readDataSchema.size > 1) {
+ fallbackReasons += "Text data source supports only a single column"
+ }
+ // Partition columns are not implemented for native Text scans (unlike
the native CSV path,
+ // which forwards them). Fall back to Spark for partitioned reads
rather than drop them.
+ if (scan.readPartitionSchema.nonEmpty) {
+ fallbackReasons += "Comet does not support partition columns in
native Text scans"
+ }
+ // The native reader loads each file whole. Spark splits a large
uncompressed text file
+ // into byte-range partitions (TextScan.isSplitable), and each split
would then re-read and
+ // buffer the whole file with no memory-pool reservation. Restrict
native execution to
+ // unsplit files (the small build-side lookups this targets); Spark
handles large files.
+ val hasSplitFile = scanExec.inputPartitions.exists {
+ case fp: FilePartition => fp.files.exists(_.start != 0)
+ case _ => false
+ }
+ if (hasSplitFile) {
+ fallbackReasons += "Comet native Text scan does not support files
split into byte " +
+ "ranges (large files)"
+ }
+ // 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
+ // (e.g. .gz, .bz2). Fall back when any input file is compressed. Use
the options-aware
+ // Hadoop conf so a codec registered via a per-read option is honored
(as Spark does).
+ val hadoopConf =
session.sessionState.newHadoopConfWithOptions(scan.options.asScala.toMap)
+ val codecFactory = new
org.apache.hadoop.io.compress.CompressionCodecFactory(hadoopConf)
+ val hasCompressedFile = scan.fileIndex.inputFiles.exists { path =>
+ codecFactory.getCodec(new org.apache.hadoop.fs.Path(path)) != null
+ }
+ if (hasCompressedFile) {
+ fallbackReasons += "Comet does not support compressed text files"
Review Comment:
Done. The test now writes a .zstd file with ZstdOutputStream as well as the
.gzip one, and checks both fall back and match Spark.
(CometTextNativeReadSuite.scala:295)
##########
native/core/src/execution/operators/text_scan.rs:
##########
@@ -0,0 +1,642 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Native Text file scan, mirroring Spark's `text` data source.
+//!
+//! Spark's text reader produces a single `value: string` column. With the
default options each
+//! line of the file is one row (split on `\n`, `\r\n`, or `\r`, with the
terminator stripped).
+//! With `wholetext=true` the whole file is a single row. A custom `lineSep`
splits on exactly that
+//! byte sequence instead of the universal-newline set. See the
`TextOptions`/`HadoopFileLinesReader`
+//! logic in Spark's `org.apache.spark.sql.execution.datasources.text` package.
+
+use crate::execution::operators::ExecutionError;
+use arrow::array::StringBuilder;
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::{RecordBatch, RecordBatchOptions};
+use datafusion::common::Result;
+use datafusion::datasource::object_store::ObjectStoreUrl;
+use datafusion::physical_expr::projection::ProjectionExprs;
+use datafusion::physical_plan::metrics::{BaselineMetrics,
ExecutionPlanMetricsSet};
+use datafusion::physical_plan::DisplayFormatType;
+use datafusion_comet_common::decode_utf8_spark_lossy;
+use datafusion_comet_proto::spark_operator::TextOptions;
+use datafusion_datasource::file::FileSource;
+use datafusion_datasource::file_compression_type::FileCompressionType;
+use datafusion_datasource::file_groups::FileGroup;
+use datafusion_datasource::file_scan_config::{FileScanConfig,
FileScanConfigBuilder};
+use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
+use datafusion_datasource::projection::{ProjectionOpener, SplitProjection};
+use datafusion_datasource::source::DataSourceExec;
+use datafusion_datasource::{as_file_source, PartitionedFile, TableSchema};
+use futures::StreamExt;
+use object_store::{ObjectStore, ObjectStoreExt};
+use std::borrow::Cow;
+use std::fmt;
+use std::io::Read;
+use std::sync::Arc;
+
+/// Batch size used when the execution framework does not set one before
opening the file.
+const DEFAULT_TEXT_BATCH_SIZE: usize = 8192;
+
+/// UTF-8 byte-order mark. Spark's line reader strips a leading one; wholetext
keeps it.
+const UTF8_BOM: &[u8] = &[0xEF, 0xBB, 0xBF];
+
+/// Cap on a batch's decoded `value` buffer bytes, kept well below Arrow's 2GB
32-bit offset limit
+/// so a batch of many rows (or a few long lines) cannot overflow. Measured
against the decoded
+/// string length, since a lossy UTF-8 decode can expand ill-formed bytes up
to 3x.
+const MAX_BATCH_VALUE_BYTES: usize = 1 << 30;
+
+pub fn init_text_datasource_exec(
+ object_store_url: ObjectStoreUrl,
+ file_groups: Vec<Vec<PartitionedFile>>,
+ data_schema: SchemaRef,
+ _partition_schema: SchemaRef,
+ projection_vector: Vec<usize>,
+ text_options: &TextOptions,
+) -> Result<Arc<DataSourceExec>, ExecutionError> {
+ let text_source = TextSource::new(data_schema, text_options);
+
+ let file_groups = file_groups.into_iter().map(FileGroup::new).collect();
+
+ let file_scan_config = FileScanConfigBuilder::new(object_store_url,
Arc::new(text_source))
+ .with_file_groups(file_groups)
+ .with_projection_indices(Some(projection_vector))?
+ .build();
+
+ Ok(DataSourceExec::from_data_source(file_scan_config))
+}
+
+/// A [`FileSource`] for Spark's text format. The file schema is always the
single `value` column;
+/// projection (and partition columns, if any) is handled by
[`ProjectionOpener`], exactly as the
+/// CSV source does.
+#[derive(Debug, Clone)]
+struct TextSource {
+ table_schema: TableSchema,
+ projection: SplitProjection,
+ batch_size: Option<usize>,
+ metrics: ExecutionPlanMetricsSet,
+ whole_text: bool,
+ // The raw separator bytes when a custom `lineSep` is set; None means
universal newline
+ // (`\n`, `\r\n`, `\r`).
+ line_sep: Option<Vec<u8>>,
+}
+
+impl TextSource {
+ fn new(file_schema: SchemaRef, options: &TextOptions) -> Self {
+ let table_schema: TableSchema = file_schema.into();
+ // `line_sep` arrives as the exact separator bytes Spark computed
(`lineSeparatorInRead`,
+ // already encoded with the file's charset), so we split on them
directly.
+ let line_sep = options.line_sep.clone().filter(|sep| !sep.is_empty());
+ Self {
+ projection: SplitProjection::unprojected(&table_schema),
+ table_schema,
+ batch_size: None,
+ metrics: ExecutionPlanMetricsSet::new(),
+ whole_text: options.whole_text,
+ line_sep,
+ }
+ }
+}
+
+impl From<TextSource> for Arc<dyn FileSource> {
+ fn from(source: TextSource) -> Self {
+ as_file_source(source)
+ }
+}
+
+impl FileSource for TextSource {
+ fn create_file_opener(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ base_config: &FileScanConfig,
+ partition_index: usize,
+ ) -> Result<Arc<dyn FileOpener>> {
+ // The inner opener must emit batches whose schema is the file schema
projected by
+ // `file_indices` (that is what ProjectionOpener feeds its projector).
For text that is
+ // either the single `value` column or, for a projection-less count,
an empty schema.
+ let file_schema = self.table_schema.file_schema();
+ let projected_file_schema =
Arc::new(file_schema.project(&self.projection.file_indices)?);
+ let opener = Arc::new(TextOpener {
+ projected_file_schema,
+ file_compression_type: base_config.file_compression_type,
+ object_store,
+ partition_index,
+ batch_size: self.batch_size.unwrap_or(DEFAULT_TEXT_BATCH_SIZE),
+ metrics: self.metrics.clone(),
+ whole_text: self.whole_text,
+ line_sep: self.line_sep.clone(),
+ }) as Arc<dyn FileOpener>;
+ ProjectionOpener::try_new(self.projection.clone(), opener, file_schema)
+ }
+
+ fn table_schema(&self) -> &TableSchema {
+ &self.table_schema
+ }
+
+ fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
+ let mut conf = self.clone();
+ conf.batch_size = Some(batch_size);
+ Arc::new(conf)
+ }
+
+ fn try_pushdown_projection(
+ &self,
+ projection: &ProjectionExprs,
+ ) -> Result<Option<Arc<dyn FileSource>>> {
+ let mut source = self.clone();
+ let new_projection = self.projection.source.try_merge(projection)?;
+ source.projection =
SplitProjection::new(self.table_schema.file_schema(), &new_projection);
+ Ok(Some(Arc::new(source)))
+ }
+
+ fn projection(&self) -> Option<&ProjectionExprs> {
+ Some(&self.projection.source)
+ }
+
+ fn metrics(&self) -> &ExecutionPlanMetricsSet {
+ &self.metrics
+ }
+
+ fn file_type(&self) -> &str {
+ "text"
+ }
+
+ // A custom line separator can straddle a byte-range boundary, and we read
whole files rather
+ // than boundary-aligned ranges, so this source must not let DataFusion
repartition files by
+ // byte range.
+ fn supports_repartitioning(&self) -> bool {
+ false
+ }
+
+ fn fmt_extra(&self, t: DisplayFormatType, f: &mut fmt::Formatter) ->
fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, ", whole_text={}", self.whole_text)
+ }
+ DisplayFormatType::TreeRender => Ok(()),
+ }
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(
+ &Arc<dyn datafusion::physical_plan::PhysicalExpr>,
+ ) -> Result<datafusion::common::tree_node::TreeNodeRecursion>,
+ ) -> Result<datafusion::common::tree_node::TreeNodeRecursion> {
+
datafusion::physical_plan::apply_expression_roots(self.projection.source.iter(),
f)
+ }
+}
+
+/// A [`FileOpener`] that reads a whole text file and yields the projected
file schema in batches.
+struct TextOpener {
+ // The file schema projected by `file_indices`: the single `value` column,
or an empty schema
+ // for a projection-less count.
+ projected_file_schema: SchemaRef,
+ file_compression_type: FileCompressionType,
+ object_store: Arc<dyn ObjectStore>,
+ partition_index: usize,
+ batch_size: usize,
+ metrics: ExecutionPlanMetricsSet,
+ whole_text: bool,
+ line_sep: Option<Vec<u8>>,
+}
+
+impl FileOpener for TextOpener {
+ fn open(&self, partitioned_file: PartitionedFile) ->
Result<FileOpenFuture> {
+ let store = Arc::clone(&self.object_store);
+ let projected_file_schema = Arc::clone(&self.projected_file_schema);
+ let file_compression_type = self.file_compression_type;
+ let batch_size = self.batch_size;
+ let whole_text = self.whole_text;
+ let line_sep = self.line_sep.clone();
+ let baseline_metrics = BaselineMetrics::new(&self.metrics,
self.partition_index);
+
+ Ok(Box::pin(async move {
+ let range = partitioned_file.range.clone();
+ let location = partitioned_file.object_meta.location;
+ let raw = store.get(&location).await?.bytes().await?;
+
+ // Decompress only when the file is actually compressed. On
Comet's plan config the file
+ // is always UNCOMPRESSED (CometScanRule falls back for compressed
files), so we borrow
+ // the fetched bytes directly and avoid a full-file copy; the
decompress arm is a
+ // defensive handler should compression ever be plumbed through. A
byte range never
+ // coexists with decompression (compressed files are
non-splittable).
+ //
+ // NOTE: the whole object is read into memory. CometScanRule
restricts native text to
+ // unsplit files (it falls back when Spark splits a file into byte
ranges), so in
+ // practice this reads each small file once; the range filter
below is a correctness
+ // safety net for any split that still reaches here.
+ let buf: Cow<[u8]> = if file_compression_type.is_compressed() {
+ let mut decoder =
file_compression_type.convert_read(std::io::Cursor::new(raw))?;
+ let mut decoded = Vec::new();
+ decoder.read_to_end(&mut decoded)?;
+ Cow::Owned(decoded)
+ } else {
+ Cow::Borrowed(raw.as_ref())
+ };
+
+ let mut timer = baseline_metrics.elapsed_compute().timer();
+ // Spark's line reader strips a leading UTF-8 BOM at the start of
the file, but only with
+ // the default separator. With a custom lineSep the BOM is
ordinary data (Spark leaves it
+ // for the splitter, so e.g. a file `BOM a BOM b` read with
lineSep=BOM yields
+ // ["", "a", "b"]). Only strip in line mode, at file offset 0
(native reads unsplit files
+ // only), and when no custom lineSep is set. wholetext keeps the
BOM.
+ let at_file_start = range.as_ref().map(|r| r.start ==
0).unwrap_or(true);
+ let content = strip_leading_bom(&buf, whole_text, at_file_start,
line_sep.is_some());
+ let mut lines = split_lines(content, whole_text,
line_sep.as_deref());
+ // Spark may split an uncompressed file into byte ranges
(TextScan.isSplitable). A line
+ // belongs to the split whose [start, end) contains the line's
start offset -- so each
+ // line is emitted by exactly one split, with no duplication or
gaps across splits. See
+ // DataFusion's CsvOpener for the same rule. wholetext files are
never split.
+ if !whole_text {
+ if let Some(range) = range {
+ lines = filter_lines_by_range(content, lines, range.start,
range.end);
+ }
+ }
+ let batches = build_batches(&projected_file_schema, lines,
batch_size)?;
+ timer.stop();
+
+ Ok(futures::stream::iter(batches.into_iter().map(Ok)).boxed())
+ }))
+ }
+}
+
+/// Strip a single leading UTF-8 BOM the way Spark's default line reader does:
only in line mode
+/// (not wholetext), only for the split that starts at file offset 0, and only
with the default
+/// separator. With a custom lineSep the BOM is left in place as ordinary data.
+fn strip_leading_bom(
+ buf: &[u8],
+ whole_text: bool,
+ at_file_start: bool,
+ custom_line_sep: bool,
+) -> &[u8] {
+ if !whole_text && at_file_start && !custom_line_sep &&
buf.starts_with(UTF8_BOM) {
+ &buf[UTF8_BOM.len()..]
+ } else {
+ buf
+ }
+}
+
+/// Split a decompressed file buffer into text lines following Spark's
semantics.
+fn split_lines<'a>(buf: &'a [u8], whole_text: bool, line_sep: Option<&[u8]>)
-> Vec<&'a [u8]> {
+ if whole_text {
+ // One row per file; the whole content (including internal separators)
is a single value.
+ return vec![buf];
+ }
+ match line_sep {
+ Some(sep) => split_on_separator(buf, sep),
+ None => split_universal(buf),
+ }
+}
+
+/// Universal-newline splitting: a line ends at `\n`, `\r\n`, or `\r`, and the
terminator is
+/// stripped. A trailing terminator does not produce a final empty line
(Hadoop LineRecordReader).
+fn split_universal(buf: &[u8]) -> Vec<&[u8]> {
+ let mut lines = Vec::new();
+ let mut start = 0;
+ let mut i = 0;
+ while i < buf.len() {
+ match buf[i] {
+ b'\n' => {
+ lines.push(&buf[start..i]);
+ i += 1;
+ start = i;
+ }
+ b'\r' => {
+ lines.push(&buf[start..i]);
+ i += 1;
+ if i < buf.len() && buf[i] == b'\n' {
+ i += 1;
+ }
+ start = i;
+ }
+ _ => i += 1,
+ }
+ }
+ if start < buf.len() {
+ lines.push(&buf[start..]);
+ }
+ lines
+}
+
+/// Split on an exact separator byte sequence; the trailing separator does not
produce an empty
+/// final line.
+fn split_on_separator<'a>(buf: &'a [u8], sep: &[u8]) -> Vec<&'a [u8]> {
+ let mut lines = Vec::new();
+ let mut start = 0;
+ let mut i = 0;
+ while i + sep.len() <= buf.len() {
+ if &buf[i..i + sep.len()] == sep {
Review Comment:
Fixed. Instead of reproducing Hadoop's buffered readCustomLine, the scan
falls back when the custom lineSep overlaps itself — the only case where
Hadoop's one-byte-backtrack matcher can split differently from an exact scan.
Overlap-free separators (\r\n, CJK, etc.) stay native. Added a test with
lineSep="aab". (CometScanRule.scala:531; new separatorSelfOverlaps helper)
##########
spark/src/test/scala/org/apache/comet/text/CometTextNativeReadSuite.scala:
##########
@@ -0,0 +1,424 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.text
+
+import java.io.File
+import java.nio.charset.StandardCharsets
+import java.nio.file.Files
+
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.comet.CometTextNativeScanExec
+import org.apache.spark.sql.functions.input_file_name
+import org.apache.spark.sql.internal.SQLConf
+
+import org.apache.comet.CometConf
+import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus
+
+class CometTextNativeReadSuite extends CometTestBase {
+ private val TEST_TEXT_PATH = "src/test/resources/test-data/text-test-1.txt"
+
+ private def withTextConf(f: => Unit): Unit = {
+ withSQLConf(
+ CometConf.COMET_TEXT_V2_NATIVE_ENABLED.key -> "true",
+ // Route only `text` through the V2 path; keep every other format on V1
so a Parquet probe
+ // side (used by the broadcast-join test) is not pushed through the
unsupported V2 scan.
+ SQLConf.USE_V1_SOURCE_LIST.key -> "avro,csv,json,kafka,orc,parquet")(f)
+ }
+
+ test("native text read - lines to value column") {
+ withTextConf {
+ val df = spark.read.text(TEST_TEXT_PATH)
+ checkSparkAnswerAndOperator(df)
+ }
+ }
+
+ test("native text read - count") {
+ withTextConf {
+ val df = spark.read.text(TEST_TEXT_PATH).selectExpr("count(*) as c")
+ checkSparkAnswerAndOperator(df)
+ }
+ }
+
+ test("native text read - filter and projection") {
+ withTextConf {
+ val df = spark.read.text(TEST_TEXT_PATH).where("value like '%bar%'")
+ checkSparkAnswerAndOperator(df)
+ }
+ }
+
+ test("native text read - preserves line order within a file") {
+ withTempDir { dir =>
+ val file = new File(dir, "ordered.txt")
+ val lines = (1 to 500).map(i => s"row-$i")
+ Files.write(file.toPath,
lines.mkString("\n").getBytes(StandardCharsets.UTF_8))
+ withTextConf {
+ val df = spark.read.text(file.getAbsolutePath)
+ // Assert the scan actually ran natively (a Spark fallback would also
preserve order).
+ checkSparkAnswerAndOperator(df, Seq(classOf[CometTextNativeScanExec]))
+ // checkSparkAnswer sorts both sides, so assert order explicitly here.
+ val got = df.collect().map(_.getString(0)).toSeq
+ assert(got == lines, "native text scan must preserve within-file line
order")
+ }
+ }
+ }
+
+ test("native text read - fallback for input_file_name") {
+ withTextConf {
+ val df = spark.read.text(TEST_TEXT_PATH).withColumn("f",
input_file_name())
+ checkSparkAnswerAndFallbackReason(
+ df,
+ "Native Text scan is not compatible with input_file_name")
Review Comment:
Fixed. Dropped the dead Text-specific guard; the test now asserts the
shared V2 message, which the blanket gate at the top of transformV2Scan emits.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
