diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b703c8b4de4..15331d6778e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -59,11 +59,17 @@ jobs: sccache: s3 - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6 - uses: ./.github/actions/setup-prebuild - # Use uvx for ruff to avoid building the Rust extension (saves ~4.5 min) + - name: Install uv + if: github.repository != 'vortex-data/vortex' + uses: spiraldb/actions/.github/actions/setup-uv@0.18.5 + with: + sync: false + prune-cache: false + # Use uv tool run for ruff to avoid building the Rust extension (saves ~4.5 min) - name: Python Lint - Format - run: uvx ruff format --check . + run: uv tool run ruff format --check . - name: Python Lint - Ruff - run: uvx ruff check . + run: uv tool run ruff check . # PyRight needs the project for type information, so use uv run - name: Python Lint - PyRight env: @@ -289,10 +295,14 @@ jobs: sccache: s3 - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6 - uses: ./.github/actions/setup-prebuild + - name: Install cargo-hack + if: github.repository != 'vortex-data/vortex' + shell: bash + run: command -v cargo-hack >/dev/null || cargo install --locked cargo-hack - name: Rust Lint - Clippy No Default Features shell: bash run: | - cargo hack --no-dev-deps --ignore-private clippy --profile ci --no-default-features -- -D warnings + cargo hack --ignore-private clippy --profile ci --no-default-features -- -D warnings rust-test-other: name: "Rust tests (${{ matrix.os }})" diff --git a/Cargo.lock b/Cargo.lock index 045c72176fd..8e3e27eb61a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9525,6 +9525,7 @@ dependencies = [ "tokio-stream", "tracing", "url", + "uuid", "vortex", "vortex-utils", ] @@ -9760,6 +9761,7 @@ dependencies = [ "glob", "itertools 0.14.0", "kanal", + "log", "object_store 0.13.2", "oneshot", "parking_lot", diff --git a/benchmarks-website/package-lock.json b/benchmarks-website/package-lock.json index f2d4ff4a323..02234d186cf 100644 --- a/benchmarks-website/package-lock.json +++ b/benchmarks-website/package-lock.json @@ -190,9 +190,6 @@ "arm64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -210,9 +207,6 @@ "arm64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -230,9 +224,6 @@ "ppc64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -250,9 +241,6 @@ "s390x" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -270,9 +258,6 @@ "x64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -290,9 +275,6 @@ "x64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -838,9 +820,6 @@ "arm64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -862,9 +841,6 @@ "arm64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -886,9 +862,6 @@ "x64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -910,9 +883,6 @@ "x64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MPL-2.0", "optional": true, "os": [ diff --git a/java/vortex-jni/bin/main/dev/vortex/DateTimeUtil.class b/java/vortex-jni/bin/main/dev/vortex/DateTimeUtil.class new file mode 100644 index 00000000000..61331568fe0 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/DateTimeUtil.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/VortexCleaner.class b/java/vortex-jni/bin/main/dev/vortex/VortexCleaner.class new file mode 100644 index 00000000000..162348d3d7b Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/VortexCleaner.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount$Estimate.class b/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount$Estimate.class new file mode 100644 index 00000000000..f024edfd951 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount$Estimate.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount$Exact.class b/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount$Exact.class new file mode 100644 index 00000000000..9aea31e8dbc Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount$Exact.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount$Unknown.class b/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount$Unknown.class new file mode 100644 index 00000000000..a5f7888ffc7 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount$Unknown.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount.class b/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount.class new file mode 100644 index 00000000000..b06b83518c5 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/DataSource$RowCount.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/DataSource.class b/java/vortex-jni/bin/main/dev/vortex/api/DataSource.class new file mode 100644 index 00000000000..4914cc746d0 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/DataSource.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/Expression$BinaryOp.class b/java/vortex-jni/bin/main/dev/vortex/api/Expression$BinaryOp.class new file mode 100644 index 00000000000..82a2a5dcbab Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/Expression$BinaryOp.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/Expression.class b/java/vortex-jni/bin/main/dev/vortex/api/Expression.class new file mode 100644 index 00000000000..3753c2f6bea Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/Expression.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/Partition.class b/java/vortex-jni/bin/main/dev/vortex/api/Partition.class new file mode 100644 index 00000000000..4db31fe7c83 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/Partition.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/Scan.class b/java/vortex-jni/bin/main/dev/vortex/api/Scan.class new file mode 100644 index 00000000000..a8e2c5b5a16 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/Scan.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/ScanOptions$SelectionMode.class b/java/vortex-jni/bin/main/dev/vortex/api/ScanOptions$SelectionMode.class new file mode 100644 index 00000000000..5fbcaf75f76 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/ScanOptions$SelectionMode.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/ScanOptions.class b/java/vortex-jni/bin/main/dev/vortex/api/ScanOptions.class new file mode 100644 index 00000000000..f95697e450f Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/ScanOptions.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/Session.class b/java/vortex-jni/bin/main/dev/vortex/api/Session.class new file mode 100644 index 00000000000..0cfcb0e9380 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/Session.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/api/VortexWriter.class b/java/vortex-jni/bin/main/dev/vortex/api/VortexWriter.class new file mode 100644 index 00000000000..6cb3670637e Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/api/VortexWriter.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/arrow/ArrowAllocation.class b/java/vortex-jni/bin/main/dev/vortex/arrow/ArrowAllocation.class new file mode 100644 index 00000000000..985a3b61f66 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/arrow/ArrowAllocation.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/jni/NativeDataSource.class b/java/vortex-jni/bin/main/dev/vortex/jni/NativeDataSource.class new file mode 100644 index 00000000000..978d7217cb2 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/jni/NativeDataSource.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/jni/NativeExpression.class b/java/vortex-jni/bin/main/dev/vortex/jni/NativeExpression.class new file mode 100644 index 00000000000..fc501ccd23e Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/jni/NativeExpression.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/jni/NativeFiles.class b/java/vortex-jni/bin/main/dev/vortex/jni/NativeFiles.class new file mode 100644 index 00000000000..4526cf0674c Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/jni/NativeFiles.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/jni/NativeLoader.class b/java/vortex-jni/bin/main/dev/vortex/jni/NativeLoader.class new file mode 100644 index 00000000000..d332a38d540 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/jni/NativeLoader.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/jni/NativeLogging.class b/java/vortex-jni/bin/main/dev/vortex/jni/NativeLogging.class new file mode 100644 index 00000000000..4e23ae8e227 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/jni/NativeLogging.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/jni/NativePartition.class b/java/vortex-jni/bin/main/dev/vortex/jni/NativePartition.class new file mode 100644 index 00000000000..5a5f8e046d0 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/jni/NativePartition.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/jni/NativeRuntime.class b/java/vortex-jni/bin/main/dev/vortex/jni/NativeRuntime.class new file mode 100644 index 00000000000..cf22dd34d15 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/jni/NativeRuntime.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/jni/NativeScan.class b/java/vortex-jni/bin/main/dev/vortex/jni/NativeScan.class new file mode 100644 index 00000000000..b276fecf36e Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/jni/NativeScan.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/jni/NativeSession.class b/java/vortex-jni/bin/main/dev/vortex/jni/NativeSession.class new file mode 100644 index 00000000000..1bf0add4635 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/jni/NativeSession.class differ diff --git a/java/vortex-jni/bin/main/dev/vortex/jni/NativeWriter.class b/java/vortex-jni/bin/main/dev/vortex/jni/NativeWriter.class new file mode 100644 index 00000000000..6ec08d70fa4 Binary files /dev/null and b/java/vortex-jni/bin/main/dev/vortex/jni/NativeWriter.class differ diff --git a/java/vortex-jni/bin/test/dev/vortex/api/TestMinimal$BatchReader.class b/java/vortex-jni/bin/test/dev/vortex/api/TestMinimal$BatchReader.class new file mode 100644 index 00000000000..77dc713ce10 Binary files /dev/null and b/java/vortex-jni/bin/test/dev/vortex/api/TestMinimal$BatchReader.class differ diff --git a/java/vortex-jni/bin/test/dev/vortex/api/TestMinimal$Person.class b/java/vortex-jni/bin/test/dev/vortex/api/TestMinimal$Person.class new file mode 100644 index 00000000000..5d892e3552a Binary files /dev/null and b/java/vortex-jni/bin/test/dev/vortex/api/TestMinimal$Person.class differ diff --git a/java/vortex-jni/bin/test/dev/vortex/api/TestMinimal.class b/java/vortex-jni/bin/test/dev/vortex/api/TestMinimal.class new file mode 100644 index 00000000000..aac8261d63b Binary files /dev/null and b/java/vortex-jni/bin/test/dev/vortex/api/TestMinimal.class differ diff --git a/java/vortex-jni/bin/test/dev/vortex/jni/JNIWriterTest.class b/java/vortex-jni/bin/test/dev/vortex/jni/JNIWriterTest.class new file mode 100644 index 00000000000..cc770bf4a2e Binary files /dev/null and b/java/vortex-jni/bin/test/dev/vortex/jni/JNIWriterTest.class differ diff --git a/vortex-array/src/arrays/dict/take.rs b/vortex-array/src/arrays/dict/take.rs index b77bef19b39..34274e7a77e 100644 --- a/vortex-array/src/arrays/dict/take.rs +++ b/vortex-array/src/arrays/dict/take.rs @@ -156,15 +156,20 @@ pub(crate) fn propagate_take_stats( st.set(Stat::IsConstant, Precision::exact(true)); } } - let inexact_min_max = [Stat::Min, Stat::Max] - .into_iter() - .filter_map(|stat| match source.statistics().get(stat).into_inexact() { - Precision::Exact(scalar) | Precision::Inexact(scalar) => { - scalar.into_value().map(|sv| (stat, Precision::Inexact(sv))) - } - Precision::Absent => None, - }) - .collect::>(); + let inexact_min_max = [ + Stat::Min, + Stat::Max, + Stat::UncompressedSizeInBytes, + Stat::IsConstant, + ] + .into_iter() + .filter_map(|stat| match source.statistics().get(stat).into_inexact() { + Precision::Exact(scalar) | Precision::Inexact(scalar) => { + scalar.into_value().map(|sv| (stat, Precision::Inexact(sv))) + } + Precision::Absent => None, + }) + .collect::>(); st.combine_sets( &(unsafe { StatsSet::new_unchecked(inexact_min_max) }).as_typed_ref(source.dtype()), ) diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index c8e7fa17c36..0f4746fd4d6 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -1,47 +1,197 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use crate::ArrayRef; use crate::IntoArray; +use crate::ToCanonical; use crate::array::ArrayView; use crate::arrays::Extension; use crate::arrays::ExtensionArray; +use crate::arrays::PrimitiveArray; use crate::arrays::extension::ExtensionArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::PType; +use crate::extension::datetime::AnyTemporal; +use crate::extension::datetime::TemporalMetadata; +use crate::extension::datetime::TimeUnit; use crate::scalar_fn::fns::cast::CastReduce; +use crate::validity::Validity; impl CastReduce for Extension { fn cast(array: ArrayView<'_, Extension>, dtype: &DType) -> VortexResult> { - if !array.dtype().eq_ignore_nullability(dtype) { - // Target is not the same extension type. + let DType::Extension(ext_dtype) = dtype else { + // Target is not an extension type. // Delegate to the storage array's cast. return Ok(Some(array.storage_array().cast(dtype.clone())?)); + }; + + if array.ext_dtype().eq_ignore_nullability(ext_dtype) { + let new_storage = match array + .storage_array() + .cast(ext_dtype.storage_dtype().clone()) + .and_then(|a| a.to_canonical().map(|c| c.into_array())) + { + Ok(arr) => arr, + Err(e) => { + tracing::warn!("Failed to cast storage array: {e}"); + return Ok(None); + } + }; + + return Ok(Some( + ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), + )); } - let DType::Extension(ext_dtype) = dtype else { - unreachable!("Already verified we have an extension dtype"); - }; + if let Some(new_storage) = cast_temporal_date_to_timestamp(&array, dtype)? { + return Ok(Some( + ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), + )); + } + + Ok(None) + } +} + +fn cast_temporal_date_to_timestamp( + array: &ArrayView<'_, Extension>, + target_dtype: &DType, +) -> VortexResult> { + let DType::Extension(target_ext_dtype) = target_dtype else { + return Ok(None); + }; + + let Some(source_temporal) = array.ext_dtype().metadata_opt::() else { + return Ok(None); + }; + let Some(target_temporal) = target_ext_dtype.metadata_opt::() else { + return Ok(None); + }; + + let (TemporalMetadata::Date(source_unit), TemporalMetadata::Timestamp(target_unit, _)) = + (source_temporal, target_temporal) + else { + return Ok(None); + }; + + let source_i64 = array + .storage_array() + .cast(DType::Primitive(PType::I64, array.dtype().nullability()))?; + let source_i64 = source_i64.to_primitive(); + + let converted = cast_date_values_to_timestamp(&source_i64, *source_unit, *target_unit)?; + + converted + .into_array() + .cast(target_ext_dtype.storage_dtype().clone()) + .map(Some) +} - let new_storage = match array - .storage_array() - .cast(ext_dtype.storage_dtype().clone()) - { - Ok(arr) => arr, - Err(e) => { - tracing::warn!("Failed to cast storage array: {e}"); - return Ok(None); +fn cast_date_values_to_timestamp( + values: &PrimitiveArray, + source_unit: TimeUnit, + target_unit: TimeUnit, +) -> VortexResult { + let (multiply, divide) = date_to_timestamp_scale(source_unit, target_unit)?; + + let input = values.as_slice::(); + let validity = values.validity()?; + let mut output = BufferMut::with_capacity(input.len()); + + match &validity { + Validity::NonNullable | Validity::AllValid => { + for &value in input { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(convert_temporal_value(value, multiply, divide)?) }; } - }; + } + Validity::AllInvalid => { + for _ in 0..input.len() { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(0i64) }; + } + } + Validity::Array(validity_arr) => { + for (i, &value) in input.iter().enumerate() { + let is_valid = validity_arr + .scalar_at(i)? + .as_bool() + .value() + .unwrap_or(false); + if is_valid { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { + output.push_unchecked(convert_temporal_value(value, multiply, divide)?) + }; + } else { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(0i64) }; + } + } + } + } + + Ok(PrimitiveArray::new(output.freeze(), validity)) +} + +fn date_to_timestamp_scale( + source_unit: TimeUnit, + target_unit: TimeUnit, +) -> VortexResult<(i64, i64)> { + let source_ns = to_nanoseconds(source_unit)?; + let target_ns = to_nanoseconds(target_unit)?; + + if source_ns >= target_ns { + let multiply = source_ns / target_ns; + return Ok((multiply, 1)); + } + + let divide = target_ns / source_ns; + Ok((1, divide)) +} - Ok(Some( - ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), - )) +fn to_nanoseconds(unit: TimeUnit) -> VortexResult { + match unit { + TimeUnit::Nanoseconds => Ok(1), + TimeUnit::Microseconds => Ok(1_000), + TimeUnit::Milliseconds => Ok(1_000_000), + TimeUnit::Seconds => Ok(1_000_000_000), + TimeUnit::Days => Ok(86_400_000_000_000), } } +fn convert_temporal_value(value: i64, multiply: i64, divide: i64) -> VortexResult { + let mut scaled = i128::from(value) + .checked_mul(i128::from(multiply)) + .ok_or_else(|| { + vortex_error::vortex_err!( + Compute: "Date value {value} overflows while scaling to timestamp" + ) + })?; + + if divide != 1 { + let divisor = i128::from(divide); + if scaled % divisor != 0 { + vortex_bail!( + Compute: "Date value {value} cannot be represented exactly in target timestamp unit" + ); + } + scaled /= divisor; + } + + if scaled < i128::from(i64::MIN) || scaled > i128::from(i64::MAX) { + vortex_bail!(Compute: "Date value {value} overflows target timestamp range"); + } + + i64::try_from(scaled) + .map_err(|_| vortex_error::vortex_err!(Compute: "Date value {value} overflows target timestamp range")) +} + #[cfg(test)] mod tests { use std::sync::LazyLock; @@ -61,6 +211,7 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::executor::VortexSessionExecute; + use crate::extension::datetime::Date; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; use crate::session::ArraySession; @@ -101,6 +252,57 @@ mod tests { assert_eq!(output.dtype(), &new_dtype); } + #[test] + fn cast_date_days_to_timestamp_nanoseconds() { + let source_dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased(); + let target_dtype = Timestamp::new(TimeUnit::Nanoseconds, Nullability::NonNullable).erased(); + + let arr = ExtensionArray::new(source_dtype, buffer![0i32, 1, -1].into_array()); + let output = arr + .into_array() + .cast(DType::Extension(target_dtype.clone())) + .unwrap() + .to_extension(); + + assert_eq!(output.dtype(), &DType::Extension(target_dtype)); + + let storage = output.storage_array().to_primitive(); + assert_eq!( + storage.as_slice::(), + &[0, 86_400_000_000_000, -86_400_000_000_000] + ); + } + + #[test] + fn cast_date_days_to_timestamp_seconds_nullable() { + let source_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); + let target_dtype = Timestamp::new(TimeUnit::Seconds, Nullability::Nullable).erased(); + + let arr = ExtensionArray::new( + source_dtype, + PrimitiveArray::from_option_iter([Some(0i32), None, Some(2)]).into_array(), + ); + + let output = arr + .into_array() + .cast(DType::Extension(target_dtype.clone())) + .unwrap() + .to_extension(); + + assert_eq!(output.dtype(), &DType::Extension(target_dtype)); + + let storage = output.storage_array().to_primitive(); + assert_eq!( + storage.scalar_at(0).unwrap().as_primitive().as_::(), + Some(0) + ); + assert!(storage.scalar_at(1).unwrap().is_null()); + assert_eq!( + storage.scalar_at(2).unwrap().as_primitive().as_::(), + Some(172_800) + ); + } + #[test] fn cast_different_ext_dtype() { let original_dtype = diff --git a/vortex-array/src/arrow/convert.rs b/vortex-array/src/arrow/convert.rs index 7c64d937704..419caacd82d 100644 --- a/vortex-array/src/arrow/convert.rs +++ b/vortex-array/src/arrow/convert.rs @@ -13,6 +13,7 @@ use arrow_array::GenericByteArray; use arrow_array::GenericByteViewArray; use arrow_array::GenericListArray; use arrow_array::GenericListViewArray; +use arrow_array::MapArray as ArrowMapArray; use arrow_array::NullArray as ArrowNullArray; use arrow_array::OffsetSizeTrait; use arrow_array::PrimitiveArray as ArrowPrimitiveArray; @@ -459,6 +460,18 @@ impl FromArrowArray<&ArrowFixedSizeListArray> for ArrayRef { } } +impl FromArrowArray<&ArrowMapArray> for ArrayRef { + fn from_arrow(value: &ArrowMapArray, nullable: bool) -> VortexResult { + // Arrow Map is logically List> with i32 offsets. + // We convert it to a ListArray of structs. + let entries = Self::from_arrow(value.entries() as &dyn ArrowArray, false)?; + let offsets = value.offsets().clone().into_array(); + let nulls = nulls(value.nulls(), nullable); + + Ok(ListArray::try_new(entries, offsets, nulls)?.into_array()) + } +} + impl FromArrowArray<&ArrowNullArray> for ArrayRef { fn from_arrow(value: &ArrowNullArray, nullable: bool) -> VortexResult { assert!(nullable); @@ -520,6 +533,7 @@ impl FromArrowArray<&dyn ArrowArray> for ArrayRef { DataType::ListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::LargeListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::FixedSizeList(..) => Self::from_arrow(array.as_fixed_size_list(), nullable), + DataType::Map(..) => Self::from_arrow(array.as_map(), nullable), DataType::Null => Self::from_arrow(as_null_array(array), nullable), DataType::Timestamp(u, _) => match u { ArrowTimeUnit::Second => { @@ -684,6 +698,7 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::Decimal; use crate::arrays::FixedSizeList; use crate::arrays::List; @@ -696,6 +711,7 @@ mod tests { use crate::arrays::listview::ListViewArrayExt; use crate::arrays::struct_::StructArrayExt; use crate::arrow::FromArrowArray as _; + use crate::arrow::executor::ArrowArrayExecutor as _; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -1501,4 +1517,57 @@ mod tests { ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).unwrap(); } + + #[test] + fn test_map_array_conversion() { + use arrow_array::MapArray; + use arrow_array::builder::MapBuilder; + use arrow_array::builder::StringBuilder; + + // Build a MapArray: map + let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + // First map entry: {"a": 1, "b": 2} + builder.keys().append_value("a"); + builder.values().append_value(1); + builder.keys().append_value("b"); + builder.values().append_value(2); + builder.append(true).unwrap(); + + // Second map entry: null + builder.append(false).unwrap(); + + // Third map entry: {"c": 3} + builder.keys().append_value("c"); + builder.values().append_value(3); + builder.append(true).unwrap(); + + let arrow_map = builder.finish(); + assert_eq!(arrow_map.len(), 3); + + // Convert Arrow MapArray → Vortex ListArray + let vortex_array = ArrayRef::from_arrow(&arrow_map, true).unwrap(); + assert_eq!(vortex_array.len(), 3); + + // Verify it's stored as List> + let list_array = vortex_array.as_::(); + assert_eq!(list_array.elements().len(), 3); // 3 total key-value pairs + let struct_elements = list_array.elements().as_::(); + assert_eq!(struct_elements.names().len(), 2); // key and value fields + + // Convert back to Arrow as a MapArray + let map_dtype = arrow_map.data_type().clone(); + let arrow_back = vortex_array + .execute_arrow( + Some(&map_dtype), + &mut crate::LEGACY_SESSION.create_execution_ctx(), + ) + .unwrap(); + let map_back = arrow_back + .as_any() + .downcast_ref::() + .expect("Should be a MapArray"); + assert_eq!(map_back.len(), 3); + assert_eq!(map_back.entries().len(), 3); + assert!(map_back.is_null(1)); // Second entry was null + } } diff --git a/vortex-array/src/arrow/executor/map.rs b/vortex-array/src/arrow/executor/map.rs new file mode 100644 index 00000000000..9a84281e60d --- /dev/null +++ b/vortex-array/src/arrow/executor/map.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::MapArray as ArrowMapArray; +use arrow_array::StructArray as ArrowStructArray; +use arrow_schema::FieldRef; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrow::executor::list::to_arrow_list; + +/// Convert a Vortex List> array into an Arrow MapArray. +pub(super) fn to_arrow_map( + array: ArrayRef, + entries_field: &FieldRef, + ordered: bool, + ctx: &mut ExecutionCtx, +) -> VortexResult { + // First, convert to Arrow ListArray since Map uses i32 offsets. + let list_array = to_arrow_list::(array, entries_field, ctx)?; + + // Downcast to GenericListArray to extract its components. + let Some(list_array) = list_array + .as_any() + .downcast_ref::>() + else { + vortex_bail!("to_arrow_list returned a non-ListArray when building a MapArray"); + }; + + // Extract components from the ListArray. + let (_list_field, offsets, entries, nulls) = list_array.clone().into_parts(); + + // The entries should be a StructArray. Downcast it. + let Some(entries_struct) = entries.as_any().downcast_ref::() else { + vortex_bail!("Map entries must be a StructArray"); + }; + let entries_struct = entries_struct.clone(); + + // Build the MapArray from the components. + let map_array = ArrowMapArray::try_new( + Arc::clone(entries_field), + offsets, + entries_struct, + nulls, + ordered, + )?; + + Ok(Arc::new(map_array)) +} diff --git a/vortex-array/src/arrow/executor/mod.rs b/vortex-array/src/arrow/executor/mod.rs index 49bd87bcbdf..504511b3ae6 100644 --- a/vortex-array/src/arrow/executor/mod.rs +++ b/vortex-array/src/arrow/executor/mod.rs @@ -14,6 +14,7 @@ mod dictionary; mod fixed_size_list; mod list; mod list_view; +mod map; pub mod null; pub mod primitive; mod run_end; @@ -48,6 +49,7 @@ use crate::arrow::executor::dictionary::to_arrow_dictionary; use crate::arrow::executor::fixed_size_list::to_arrow_fixed_list; use crate::arrow::executor::list::to_arrow_list; use crate::arrow::executor::list_view::to_arrow_list_view; +use crate::arrow::executor::map::to_arrow_map; use crate::arrow::executor::null::to_arrow_null; use crate::arrow::executor::primitive::to_arrow_primitive; use crate::arrow::executor::run_end::to_arrow_run_end; @@ -188,11 +190,11 @@ pub(crate) fn execute_arrow_naive( DataType::RunEndEncoded(ends_type, values_type) => { to_arrow_run_end(array, ends_type.data_type(), values_type, ctx) } + DataType::Map(entries_field, ordered) => to_arrow_map(array, entries_field, *ordered, ctx), dt @ (DataType::Date32 | DataType::Date64) => to_arrow_date(array, dt, ctx), dt @ (DataType::Time32(_) | DataType::Time64(_)) => to_arrow_time(array, dt, ctx), dt @ DataType::Timestamp(..) => to_arrow_timestamp(array, dt, ctx), DataType::FixedSizeBinary(_) - | DataType::Map(..) | DataType::Duration(_) | DataType::Interval(_) | DataType::Union(..) => { diff --git a/vortex-array/src/dtype/arrow.rs b/vortex-array/src/dtype/arrow.rs index 8dd80082533..b107746e3e9 100644 --- a/vortex-array/src/dtype/arrow.rs +++ b/vortex-array/src/dtype/arrow.rs @@ -191,6 +191,14 @@ impl FromArrowType<(&DataType, Nullability)> for DType { | DataType::LargeListView(e) => { DType::List(Arc::new(Self::from_arrow(e.as_ref())), nullability) } + DataType::Map(entries_field, _ordered) => { + // Map is logically List>. + // The entries_field contains a Struct type with the key and value fields. + DType::List( + Arc::new(Self::from_arrow(entries_field.as_ref())), + nullability, + ) + } DataType::FixedSizeList(e, size) => DType::FixedSizeList( Arc::new(Self::from_arrow(e.as_ref())), *size as u32, diff --git a/vortex-array/src/stats/mod.rs b/vortex-array/src/stats/mod.rs index ceb085e0815..8c001ddd2cb 100644 --- a/vortex-array/src/stats/mod.rs +++ b/vortex-array/src/stats/mod.rs @@ -39,6 +39,7 @@ pub const PRUNING_STATS: &[Stat] = &[ Stat::Sum, Stat::NullCount, Stat::NaNCount, + Stat::UncompressedSizeInBytes, ]; pub fn as_stat_bitset_bytes(stats: &[Stat]) -> Vec { diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index 0be774bfbc5..b0a8ea67423 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -52,6 +52,14 @@ impl Deref for BtrBlocksCompressor { } } +impl std::fmt::Debug for BtrBlocksCompressor { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BtrBlocksCompressor") + .field("compressor", &self.0) + .finish() + } +} + impl Default for BtrBlocksCompressor { fn default() -> Self { BtrBlocksCompressorBuilder::default().build() diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 2215bd4db8e..21a93b19cf4 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -34,6 +34,7 @@ object_store = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "fs"] } tokio-stream = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } +uuid = { workspace = true, features = ["v7"] } vortex = { workspace = true, features = ["object_store", "tokio", "files"] } vortex-utils = { workspace = true, features = ["dashmap"] } diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 6a9ac0fdc52..ce6c9eb0aa8 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -16,7 +16,6 @@ use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::projection::ProjectionExpr; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::utils::collect_columns; -use datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr; use datafusion_physical_plan::expressions as df_expr; use itertools::Itertools; use vortex::dtype::DType; @@ -54,10 +53,13 @@ pub(crate) fn make_vortex_predicate( expr_convertor: &dyn ExpressionConvertor, predicate: &[Arc], ) -> DFResult> { - let exprs = predicate + let exprs: Vec<_> = predicate .iter() - .map(|e| expr_convertor.convert(e.as_ref())) - .collect::>>()?; + .filter_map(|e| { + // If conversion fails, skip this expression (equivalent to lit(true) in AND conjunction) + expr_convertor.convert(e.as_ref()).ok() + }) + .collect(); Ok(and_collect(exprs)) } @@ -149,18 +151,6 @@ impl DefaultExpressionConvertor { /// Attempts to convert a DataFusion CaseExpr to a Vortex expression. fn try_convert_case_expr(&self, case_expr: &df_expr::CaseExpr) -> DFResult { - // DataFusion CaseExpr has: - // - expr(): Optional base expression (for "CASE expr WHEN ..." form) - // - when_then_expr(): Vec of (when, then) pairs - // - else_expr(): Optional else expression - - // We don't support the "CASE expr WHEN value1 THEN result1" form yet - if case_expr.expr().is_some() { - return Err(exec_datafusion_err!( - "CASE expr WHEN form is not yet supported, only searched CASE is supported" - )); - } - let when_then_pairs = case_expr.when_then_expr(); if when_then_pairs.is_empty() { return Err(exec_datafusion_err!( @@ -168,21 +158,28 @@ impl DefaultExpressionConvertor { )); } - // Convert all when/then pairs to (condition, value) tuples + let base_expr = case_expr + .expr() + .map(|expr| self.convert(expr.as_ref())) + .transpose()?; + let mut pairs = Vec::with_capacity(when_then_pairs.len()); for (when_expr, then_expr) in when_then_pairs { - let condition = self.convert(when_expr.as_ref())?; + let when_expr = self.convert(when_expr.as_ref())?; + let condition = if let Some(base_expr) = &base_expr { + Binary.new_expr(Operator::Eq, [base_expr.clone(), when_expr]) + } else { + when_expr + }; let value = self.convert(then_expr.as_ref())?; pairs.push((condition, value)); } - // Convert optional else expression let else_value = case_expr .else_expr() .map(|e| self.convert(e.as_ref())) .transpose()?; - // Build a single n-ary CASE WHEN expression from DataFusion WHEN/THEN pairs Ok(nested_case_when(pairs, else_value)) } } @@ -298,7 +295,7 @@ impl ExpressionConvertor for DefaultExpressionConvertor { let r = projection_expr.expr.apply(|node| { // We only pull column children of scalar functions that we can't push into the scan. if let Some(scalar_fn_expr) = node.as_any().downcast_ref::() - && !can_scalar_fn_be_pushed_down(scalar_fn_expr) + && !can_scalar_fn_be_pushed_down(scalar_fn_expr, input_schema) { scan_projection.extend( collect_columns(node) @@ -407,12 +404,6 @@ fn try_operator_from_df(value: &DFOperator) -> DFResult { } fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> bool { - // We currently do not support pushdown of dynamic expressions in DF. - // See issue: https://github.com/vortex-data/vortex/issues/4034 - if is_dynamic_physical_expr(df_expr) { - return false; - } - let expr = df_expr.as_any(); if let Some(binary) = expr.downcast_ref::() { can_binary_be_pushed_down(binary, schema) @@ -443,9 +434,15 @@ fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> .iter() .all(|e| can_be_pushed_down_impl(e, schema)) } else if let Some(scalar_fn) = expr.downcast_ref::() { - can_scalar_fn_be_pushed_down(scalar_fn) + can_scalar_fn_be_pushed_down(scalar_fn, schema) } else if let Some(case_expr) = expr.downcast_ref::() { can_case_be_pushed_down(case_expr, schema) + } else if expr + .downcast_ref::() + .is_some() + { + // Assume dynamic filters can be pushed down - the child won't be specified until execution time + true } else { tracing::debug!(%df_expr, "DataFusion expression can't be pushed down"); false @@ -485,13 +482,16 @@ fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> b } fn can_case_be_pushed_down(case_expr: &df_expr::CaseExpr, schema: &Schema) -> bool { - // We only support the "searched CASE" form (CASE WHEN cond THEN result ...) - // not the "simple CASE" form (CASE expr WHEN value THEN result ...) - if case_expr.expr().is_some() { + if case_expr.when_then_expr().is_empty() { + return false; + } + + if let Some(base_expr) = case_expr.expr() + && !can_be_pushed_down_impl(base_expr, schema) + { return false; } - // Check all when/then pairs for (when_expr, then_expr) in case_expr.when_then_expr() { if !can_be_pushed_down_impl(when_expr, schema) || !can_be_pushed_down_impl(then_expr, schema) @@ -500,7 +500,6 @@ fn can_case_be_pushed_down(case_expr: &df_expr::CaseExpr, schema: &Schema) -> bo } } - // Check the optional else clause if let Some(else_expr) = case_expr.else_expr() && !can_be_pushed_down_impl(else_expr, schema) { @@ -544,9 +543,13 @@ fn supported_data_types(dt: &DataType) -> bool { } /// Checks if a scalar function can be pushed down. -/// Currently only GetFieldFunc is supported. -fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr) -> bool { +/// Currently only GetFieldFunc is supported, and its arguments must also be pushable. +fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool { ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() + && scalar_fn + .args() + .iter() + .all(|arg| can_be_pushed_down_impl(arg, schema)) } // TODO(adam): Replace with `DataType::is_decimal` once its released. @@ -727,6 +730,51 @@ mod tests { ); } + #[test] + fn test_expr_from_df_case_when_with_else() { + let when_then_expr = vec![( + Arc::new(df_expr::Column::new("active", 0)) as Arc, + Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "yes".to_string(), + )))) as Arc, + )]; + let case_expr = df_expr::CaseExpr::try_new( + None, + when_then_expr, + Some(Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "no".to_string(), + )))) as Arc), + ) + .unwrap(); + + let result = DefaultExpressionConvertor::default() + .convert(&case_expr) + .unwrap(); + + assert_snapshot!(result.display_tree().to_string(), @r#" + vortex.case_when(case_when(pairs=1, else=true)) + ├── when_0: vortex.get_item(active) + │ └── input: vortex.root() + ├── then_0: vortex.literal("yes") + └── else: vortex.literal("no") + "#); + } + + #[test] + fn test_expr_from_df_case_when_without_else_pushable() { + let when_then_expr = vec![( + Arc::new(df_expr::Column::new("active", 0)) as Arc, + Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "yes".to_string(), + )))) as Arc, + )]; + let case_expr = Arc::new(df_expr::CaseExpr::try_new(None, when_then_expr, None).unwrap()) + as Arc; + + let schema = Schema::new(vec![Field::new("active", DataType::Boolean, false)]); + assert!(can_be_pushed_down_impl(&case_expr, &schema)); + } + #[rstest] // Supported types #[case::null(DataType::Null, true)] diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs index d3115c56b09..07098b28b49 100644 --- a/vortex-datafusion/src/lib.rs +++ b/vortex-datafusion/src/lib.rs @@ -95,6 +95,7 @@ pub mod v2; #[cfg(test)] mod tests; +pub use convert::exprs::DefaultExpressionConvertor; pub use convert::exprs::ExpressionConvertor; pub use persistent::*; diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 2ccfbf8ca7b..043f1d90342 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -37,6 +37,8 @@ use datafusion_execution::cache::cache_manager::CachedFileMetadataEntry; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr::LexRequirement; use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::ExecutionPlanProperties; +use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use futures::FutureExt; use futures::StreamExt as _; use futures::TryStreamExt as _; @@ -72,6 +74,7 @@ use crate::convert::TryToDataFusion; use crate::convert::stats::is_constant_to_distinct_count; const DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES: usize = MAX_POSTSCRIPT_SIZE as usize + EOF_SIZE; +const DEFAULT_TARGET_FILE_SIZE_MB: usize = 128; /// DataFusion [`FileFormat`] implementation for `.vortex` files. /// @@ -160,6 +163,12 @@ config_namespace! { /// Values smaller than `MAX_POSTSCRIPT_SIZE + EOF_SIZE` will be clamped to that minimum /// during footer parsing. pub footer_initial_read_size_bytes: usize, default = DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES + /// Target file size in megabytes for written Vortex files. + /// + /// When greater than 0 for non-partitioned writes, Vortex bypasses + /// DataFusion's file demuxer and splits output files based on + /// approximate byte size rather than row count. + pub target_file_size_mb: usize, default = DEFAULT_TARGET_FILE_SIZE_MB /// Whether to enable projection pushdown into the underlying Vortex scan. /// /// When enabled, projection expressions may be partially evaluated during @@ -620,8 +629,42 @@ impl FileFormat for VortexFormat { return not_impl_err!("Overwrites are not implemented yet for Vortex"); } + let target_file_size = (self.opts.target_file_size_mb > 0) + .then(|| { + u64::try_from(self.opts.target_file_size_mb) + .map_err(|e| { + internal_datafusion_err!( + "target_file_size_mb cannot be represented as u64: {e}" + ) + }) + .map(|v| v.saturating_mul(1024 * 1024).max(1)) + }) + .transpose()?; + + // For non-partitioned writes, force a single input stream so VortexSink + // performs one coordinated write per statement instead of one + // independent write per CPU/input partition. + // + // Use coalescing rather than repartitioning to avoid introducing a + // shuffle/dispatcher step that can interleave batches from different + // input partitions. + // + // For partitioned writes, keep DataFusion's demuxer behavior. + let input: Arc = if conf.table_partition_cols.is_empty() + && input.output_partitioning().partition_count() > 1 + { + Arc::new(CoalescePartitionsExec::new(input)) + } else { + input + }; + let schema = Arc::clone(conf.output_schema()); - let sink = Arc::new(VortexSink::new(conf, schema, self.session.clone())); + let sink = Arc::new(VortexSink::new( + conf, + schema, + self.session.clone(), + target_file_size, + )); Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) } @@ -709,4 +752,19 @@ mod tests { let format = VortexFormat::new_with_options(VortexSession::default(), opts); assert_eq!(format.options().footer_initial_read_size_bytes, 12345); } + + #[test] + fn format_plumbs_target_file_size_mb() { + let mut opts = VortexTableOptions::default(); + opts.set("target_file_size_mb", "123").unwrap(); + + let format = VortexFormat::new_with_options(VortexSession::default(), opts); + assert_eq!(format.options().target_file_size_mb, 123); + } + + #[test] + fn format_target_file_size_default_is_128mb() { + let opts = VortexTableOptions::default(); + assert_eq!(opts.target_file_size_mb, 128); + } } diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 8a579c8557d..2f17235ed19 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -80,7 +80,7 @@ pub(crate) struct VortexOpener { pub file_pruning_predicate: Option, pub expr_adapter_factory: Arc, /// This is the table's schema without partition columns. It may contain fields which do - /// not exist in the file, and are supplied by the `schema_adapter_factory`. + /// not exist in the file. Missing columns are null-filled during batch remapping. pub table_schema: TableSchema, /// A hint for the desired row count of record batches returned from the scan. pub batch_size: usize, diff --git a/vortex-datafusion/src/persistent/sink.rs b/vortex-datafusion/src/persistent/sink.rs index b05cbe84215..b9954e45e2f 100644 --- a/vortex-datafusion/src/persistent/sink.rs +++ b/vortex-datafusion/src/persistent/sink.rs @@ -4,13 +4,17 @@ use std::any::Any; use std::sync::Arc; +use arrow_schema::Schema; use arrow_schema::SchemaRef; use async_trait::async_trait; use datafusion_common::DataFusionError; use datafusion_common::Result as DFResult; +use datafusion_common::arrow::array::RecordBatch; +use datafusion_common::arrow::array::RecordBatchOptions; use datafusion_common::exec_datafusion_err; use datafusion_common_runtime::JoinSet; use datafusion_common_runtime::SpawnedTask; +use datafusion_datasource::ListingTableUrl; use datafusion_datasource::file_sink_config::FileSink; use datafusion_datasource::file_sink_config::FileSinkConfig; use datafusion_datasource::sink::DataSink; @@ -21,32 +25,129 @@ use datafusion_execution::TaskContext; use datafusion_physical_plan::DisplayAs; use datafusion_physical_plan::DisplayFormatType; use datafusion_physical_plan::metrics::MetricsSet; +use futures::SinkExt; +use futures::Stream; use futures::StreamExt; use object_store::ObjectStore; +use object_store::ObjectStoreExt; use object_store::path::Path; +use tokio::task::JoinHandle; use tokio_stream::wrappers::ReceiverStream; +use uuid::Uuid; use vortex::array::arrow::ArrowSessionExt; use vortex::array::stream::ArrayStreamAdapter; +use vortex::dtype::DType; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteSummary; use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; + +struct WriteOutputOptions<'a> { + base_output_path: &'a ListingTableUrl, + target_file_size: Option, + extension: &'a str, + write_id: &'a str, + partition_column_names: &'a [String], + keep_partition_by_columns: bool, + single_file_output: bool, +} + +#[derive(Clone, Copy)] +struct CompressionEstimate { + prev_compressed_bytes: u64, + prev_uncompressed_bytes: u64, +} + +impl CompressionEstimate { + fn identity() -> Self { + Self { + prev_compressed_bytes: 1, + prev_uncompressed_bytes: 1, + } + } + + fn from_file_sizes(compressed_bytes: u64, uncompressed_bytes: u64) -> DFResult { + if uncompressed_bytes == 0 { + return Err(exec_datafusion_err!( + "Cannot derive compression estimate from zero uncompressed bytes" + )); + } + + Ok(Self { + prev_compressed_bytes: compressed_bytes, + prev_uncompressed_bytes: uncompressed_bytes, + }) + } + + fn estimate_compressed_size(self, uncompressed_bytes: u64) -> DFResult { + if self.prev_uncompressed_bytes == 0 { + return Err(exec_datafusion_err!( + "Compression estimate denominator must be non-zero" + )); + } + + let estimated = u128::from(uncompressed_bytes) + .checked_mul(u128::from(self.prev_compressed_bytes)) + .ok_or_else(|| { + exec_datafusion_err!( + "Compressed size estimate overflow for {} * {}", + uncompressed_bytes, + self.prev_compressed_bytes + ) + })? + / u128::from(self.prev_uncompressed_bytes); + + u64::try_from(estimated).map_err(|_| { + exec_datafusion_err!("Compressed size estimate does not fit in u64: {estimated}") + }) + } +} + +struct ActiveFileWriter { + path: Path, + sender: futures::channel::mpsc::Sender, + task: JoinHandle>, +} pub struct VortexSink { config: FileSinkConfig, schema: SchemaRef, session: VortexSession, + target_file_size: Option, } impl VortexSink { - pub fn new(config: FileSinkConfig, schema: SchemaRef, session: VortexSession) -> Self { + pub fn new( + config: FileSinkConfig, + schema: SchemaRef, + session: VortexSession, + target_file_size: Option, + ) -> Self { Self { config, schema, session, + target_file_size, } } + + fn base_output_path(&self) -> DFResult<&ListingTableUrl> { + self.config + .table_paths + .first() + .ok_or_else(|| exec_datafusion_err!("Vortex sink requires at least one table path")) + } + + fn writer_dtype(&self, writer_schema: &SchemaRef) -> DFResult { + self.session + .arrow() + .from_arrow_schema(writer_schema) + .map_err(|e| { + exec_datafusion_err!("Failed to derive Vortex DType from writer schema: {e}") + }) + } } impl std::fmt::Debug for VortexSink { @@ -87,7 +188,46 @@ impl DataSink for VortexSink { data: SendableRecordBatchStream, context: &Arc, ) -> DFResult { - FileSink::write_all(self, data, context).await + if !self.config.table_partition_cols.is_empty() { + return FileSink::write_all(self, data, context).await; + } + + let object_store = context + .runtime_env() + .object_store(&self.config.object_store_url)?; + let writer_schema = get_writer_schema(&self.config); + let dtype = self.writer_dtype(&writer_schema)?; + let write_id = Uuid::now_v7().simple().to_string(); + let base_output_path = self.base_output_path()?; + let partition_column_names = self + .config + .table_partition_cols + .iter() + .map(|(name, _)| name.clone()) + .collect::>(); + + let summaries = write_record_batch_stream_to_files( + self.session.clone(), + object_store, + dtype, + data, + Arc::clone(&writer_schema), + &WriteOutputOptions { + base_output_path, + target_file_size: self.target_file_size, + extension: &self.config.file_extension, + write_id: &write_id, + partition_column_names: &partition_column_names, + keep_partition_by_columns: self.config.keep_partition_by_columns, + single_file_output: self + .config + .file_output_mode + .single_file_output(base_output_path), + }, + ) + .await?; + + aggregate_write_summaries(summaries) } } @@ -104,63 +244,50 @@ impl FileSink for VortexSink { mut file_stream_rx: DemuxedStreamReceiver, object_store: Arc, ) -> DFResult { - let mut file_write_tasks: JoinSet> = JoinSet::new(); + let writer_schema = get_writer_schema(&self.config); + let dtype = self.writer_dtype(&writer_schema)?; + let mut file_write_tasks: JoinSet>> = JoinSet::new(); - // TODO(adamg): - // 1. We can probably be better at signaling how much memory we're consuming (potentially when reading too), see ParquetSink::spawn_writer_tasks_and_join. while let Some((path, rx)) = file_stream_rx.recv().await { let session = self.session.clone(); let object_store = Arc::clone(&object_store); - let writer_schema = get_writer_schema(&self.config); - let dtype = session - .arrow() - .from_arrow_schema(&writer_schema) - .map_err(|e| { - exec_datafusion_err!("Failed to derive Vortex DType from writer schema: {e}") - })?; - - // We need to spawn work because there's a dependency between the different files. If one file has too many batches buffered, - // the demux task might deadlock itself. - let arrow_session = session.clone(); + let dtype = dtype.clone(); let import_schema = Arc::clone(&writer_schema); + let target_file_size = self.target_file_size; + let extension = self.config.file_extension.clone(); + file_write_tasks.spawn(async move { - let stream = ReceiverStream::new(rx).map(move |rb| { - arrow_session - .arrow() - .from_arrow_record_batch(rb, &import_schema) - }); - - let stream_adapter = ArrayStreamAdapter::new(dtype, stream); - - let mut object_writer = ObjectStoreWrite::new(object_store, &path) - .await - .map_err(|e| exec_datafusion_err!("Failed to create ObjectStoreWrite: {e}"))?; - - let summary = session - .write_options() - .write(&mut object_writer, stream_adapter) - .await - .map_err(|e| exec_datafusion_err!("Failed to write Vortex file: {e}"))?; - - object_writer - .shutdown() - .await - .map_err(|e| exec_datafusion_err!("Failed to shutdown Vortex writer: {e}"))?; - - Ok((path, summary)) + let stream = ReceiverStream::new(rx).map(Ok); + write_record_batch_stream_to_paths( + session, + object_store, + dtype, + stream, + import_schema, + target_file_size, + |file_index| { + if file_index == 0 { + path.clone() + } else { + numbered_path(&path, file_index, &extension) + } + }, + ) + .await }); } - let mut row_count = 0; - + let mut row_count = 0_u64; while let Some(result) = file_write_tasks.join_next().await { match result { - Ok(r) => { - let (path, summary) = r?; - - row_count += summary.row_count(); - - tracing::info!(path = %path, "Successfully written file"); + Ok(summaries) => { + row_count = row_count + .checked_add(aggregate_write_summaries(summaries?)?) + .ok_or_else(|| { + exec_datafusion_err!( + "Row count overflow while aggregating file writer tasks" + ) + })?; } Err(e) => { if e.is_panic() { @@ -181,6 +308,328 @@ impl FileSink for VortexSink { } } +fn aggregate_write_summaries(summaries: Vec<(Path, WriteSummary)>) -> DFResult { + let mut row_count = 0_u64; + for (path, summary) in summaries { + row_count = row_count.checked_add(summary.row_count()).ok_or_else(|| { + exec_datafusion_err!( + "Row count overflow while aggregating sink summaries (current={}, file={})", + row_count, + summary.row_count() + ) + })?; + tracing::debug!(path = %path, "Successfully written file"); + } + + Ok(row_count) +} + +/// Write batches from a single input stream to one or more output files. +/// +/// For collection paths, files are emitted using the `{write_id}_{file_index:05}.{extension}` +/// naming scheme as produced by the underlying writer implementation. +/// For a single-file path, the original target path is used unless rolling is needed, +/// in which case additional files follow the same naming scheme. +async fn write_record_batch_stream_to_files( + session: VortexSession, + object_store: Arc, + dtype: DType, + data: SendableRecordBatchStream, + import_schema: SchemaRef, + output_options: &WriteOutputOptions<'_>, +) -> DFResult> { + let stream = data.map(|batch| { + let batch = batch?; + if output_options.keep_partition_by_columns + || output_options.partition_column_names.is_empty() + { + Ok(batch) + } else { + remove_partition_columns(&batch, output_options.partition_column_names) + } + }); + + write_record_batch_stream_to_paths( + session, + object_store, + dtype, + stream, + import_schema, + output_options.target_file_size, + |file_index| { + output_file_path( + output_options.base_output_path, + file_index, + output_options.extension, + output_options.single_file_output, + output_options.write_id, + ) + }, + ) + .await +} + +async fn write_record_batch_stream_to_paths( + session: VortexSession, + object_store: Arc, + dtype: DType, + data: S, + import_schema: SchemaRef, + target_file_size: Option, + mut output_path: F, +) -> DFResult> +where + S: Stream>, + F: FnMut(usize) -> Path, +{ + let target = target_file_size.map(|t| t.max(1)); + let mut results: Vec<(Path, WriteSummary)> = Vec::new(); + let mut active_writer: Option = None; + let mut uncompressed_bytes_in_file = 0_u64; + let mut file_index = 0_usize; + let mut compression_estimate = CompressionEstimate::identity(); + + futures::pin_mut!(data); + + let write_result: DFResult<()> = async { + while let Some(batch) = data.next().await.transpose()? { + if active_writer.is_none() { + active_writer = Some(start_file_writer( + &session, + Arc::clone(&object_store), + output_path(file_index), + dtype.clone(), + Arc::clone(&import_schema), + )); + } + + let batch_bytes = batch_uncompressed_bytes(&batch)?; + let writer = active_writer.as_mut().ok_or_else(|| { + exec_datafusion_err!("Missing active file writer for sink output") + })?; + send_batch_to_active_writer(writer, batch).await?; + uncompressed_bytes_in_file = uncompressed_bytes_in_file + .checked_add(batch_bytes) + .ok_or_else(|| { + exec_datafusion_err!( + "Uncompressed byte counter overflow for output file {}", + writer.path + ) + })?; + + if let Some(target) = target { + let estimated_compressed = + compression_estimate.estimate_compressed_size(uncompressed_bytes_in_file)?; + if estimated_compressed >= target { + let writer = active_writer.take().ok_or_else(|| { + exec_datafusion_err!( + "Missing active file writer while finalizing rotated output file" + ) + })?; + let file_path = writer.path.clone(); + let summary = finish_file_writer(writer).await?; + if uncompressed_bytes_in_file > 0 { + compression_estimate = CompressionEstimate::from_file_sizes( + summary.size(), + uncompressed_bytes_in_file, + )?; + } + + results.push((file_path, summary)); + uncompressed_bytes_in_file = 0; + file_index += 1; + } + } + } + + if let Some(writer) = active_writer.take() { + let file_path = writer.path.clone(); + let summary = finish_file_writer(writer).await?; + results.push((file_path, summary)); + } + + Ok(()) + } + .await; + + if let Err(err) = write_result { + cleanup_failed_write( + object_store, + active_writer.into_iter().collect(), + results.iter().map(|(path, _)| path.clone()).collect(), + ) + .await; + return Err(err); + } + + Ok(results) +} + +/// Generate a numbered file path from an existing path for size-based splitting. +/// +/// Given `base/file.vortex`, produces `base/file_00000.vortex`. +/// If the path has no recognized extension, appends `_00000.{extension}`. +fn numbered_path(original: &Path, index: usize, extension: &str) -> Path { + let s = original.to_string(); + let suffix = format!(".{extension}"); + if let Some(stem) = s.strip_suffix(&suffix) { + Path::from(format!("{stem}_{index:05}{suffix}")) + } else { + Path::from(format!("{s}_{index:05}.{extension}")) + } +} + +fn start_file_writer( + session: &VortexSession, + object_store: Arc, + path: Path, + dtype: DType, + import_schema: SchemaRef, +) -> ActiveFileWriter { + let (sender, receiver) = futures::channel::mpsc::channel::(1); + let session = session.clone(); + let path_for_task = path.clone(); + + let task = tokio::spawn(async move { + let mut object_writer = ObjectStoreWrite::new(object_store, &path_for_task) + .await + .map_err(|e| { + exec_datafusion_err!( + "Failed to create ObjectStoreWrite for '{}': {e}", + path_for_task + ) + })?; + + let arrow_session = session.clone(); + let stream = receiver.map(move |rb| { + arrow_session + .arrow() + .from_arrow_record_batch(rb, &import_schema) + }); + let stream_adapter = ArrayStreamAdapter::new(dtype, stream); + + let summary = session + .write_options() + .write(&mut object_writer, stream_adapter) + .await + .map_err(|e| { + exec_datafusion_err!("Failed to write Vortex file '{}': {e}", path_for_task) + })?; + + object_writer.shutdown().await.map_err(|e| { + exec_datafusion_err!("Failed to shutdown Vortex writer '{}': {e}", path_for_task) + })?; + + Ok(summary) + }); + + ActiveFileWriter { path, sender, task } +} + +async fn cleanup_failed_write( + object_store: Arc, + active_writers: Vec, + finished_paths: Vec, +) { + let mut cleanup_paths = HashSet::new(); + + for writer in active_writers { + cleanup_paths.insert(writer.path.clone()); + writer.task.abort(); + drop(writer.task.await); + } + + for path in finished_paths { + cleanup_paths.insert(path); + } + + for path in cleanup_paths { + if let Err(e) = object_store.delete(&path).await { + tracing::warn!(path = %path, error = %e, "Failed to delete sink output during error cleanup"); + } + } +} + +async fn send_batch_to_active_writer( + writer: &mut ActiveFileWriter, + batch: RecordBatch, +) -> DFResult<()> { + writer.sender.send(batch).await.map_err(|e| { + exec_datafusion_err!( + "Failed to send batch to active writer for '{}': {e}", + writer.path + ) + }) +} + +async fn finish_file_writer(mut writer: ActiveFileWriter) -> DFResult { + writer.sender.close_channel(); + match writer.task.await { + Ok(result) => result, + Err(e) => Err(exec_datafusion_err!( + "Vortex writer task for '{}' failed to join: {e}", + writer.path + )), + } +} + +fn batch_uncompressed_bytes(batch: &RecordBatch) -> DFResult { + u64::try_from(batch.get_array_memory_size()).map_err(|_| { + exec_datafusion_err!( + "RecordBatch memory size does not fit in u64: {}", + batch.get_array_memory_size() + ) + }) +} + +fn remove_partition_columns( + batch: &RecordBatch, + partition_column_names: &[String], +) -> DFResult { + let partition_names: Vec<_> = partition_column_names.iter().map(String::as_str).collect(); + let (columns, fields): (Vec<_>, Vec<_>) = batch + .columns() + .iter() + .zip(batch.schema().fields()) + .filter(|(_, field)| !partition_names.contains(&field.name().as_str())) + .map(|(array, field)| (Arc::clone(array), (**field).clone())) + .unzip(); + + let schema = Schema::new(fields); + if columns.is_empty() { + let options = RecordBatchOptions::default().with_row_count(Some(batch.num_rows())); + return Ok(RecordBatch::try_new_with_options( + Arc::new(schema), + columns, + &options, + )?); + } + + Ok(RecordBatch::try_new(Arc::new(schema), columns)?) +} + +/// Build the output path for a rolling write. +fn output_file_path( + base_output_path: &ListingTableUrl, + file_index: usize, + extension: &str, + single_file_output: bool, + write_id: &str, +) -> Path { + if single_file_output { + if file_index == 0 { + return base_output_path.prefix().clone(); + } + return numbered_path(base_output_path.prefix(), file_index, extension); + } + + let mut base = base_output_path.prefix().to_string(); + if !base.ends_with('/') { + base.push('/'); + } + Path::from(format!("{base}{write_id}_{file_index:05}.{extension}")) +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -199,12 +648,32 @@ mod tests { use datafusion::logical_expr::Values; use datafusion::logical_expr::dml::InsertOp; use datafusion_common::ScalarValue; + use datafusion_common::exec_datafusion_err; + use datafusion_datasource::ListingTableUrl; use datafusion_datasource::file_format::format_as_file_type; use futures::TryStreamExt; use rstest::rstest; + use tokio::sync::oneshot; + use tokio::time::Duration; use crate::common_tests::TestSessionContext; use crate::persistent::VortexFormatFactory; + use crate::persistent::VortexTableOptions; + use crate::persistent::sink::ActiveFileWriter; + use crate::persistent::sink::finish_file_writer; + + fn split_path( + base_path: &object_store::path::Path, + file_index: usize, + extension: &str, + ) -> object_store::path::Path { + let mut base = base_path.to_string(); + if !base.ends_with('/') { + base.push('/'); + } + let filename = format!("part-{file_index:05}.{extension}"); + object_store::path::Path::from(format!("{base}{filename}")) + } #[tokio::test] async fn test_insert_into_sql() -> anyhow::Result<()> { @@ -304,10 +773,11 @@ mod tests { } /// Reproduction by . + /// Uses a 1MB target file size to exercise file splitting behavior. #[rstest] - #[case(1000, 1)] - #[case(40_961, 4)] - #[case(1_000_000, 4)] + #[case(1_000, 1)] + #[case(5_000_000, 6)] + #[case(10_000_000, 10)] #[tokio::test] async fn test_write_large_batch( #[case] entries: usize, @@ -315,15 +785,26 @@ mod tests { ) -> anyhow::Result<()> { let ctx = TestSessionContext::default(); + let opts = VortexTableOptions { + target_file_size_mb: 1, + ..Default::default() + }; + + let factory = VortexFormatFactory::new().with_options(opts); + + let values: Vec = (0..entries) + .map(|i| i8::try_from(i % 127)) + .collect::>()?; + let data = ctx.session.read_batch(RecordBatch::try_new( Arc::new(Schema::new(vec![Field::new("a", DataType::Int8, false)])), - vec![Arc::new(Int8Array::from(vec![0i8; entries]))], + vec![Arc::new(Int8Array::from(values))], )?)?; let logical_plan = LogicalPlanBuilder::copy_to( data.logical_plan().clone(), "/table/".to_string(), - format_as_file_type(Arc::new(VortexFormatFactory::new())), + format_as_file_type(Arc::new(factory)), Default::default(), vec![], )? @@ -359,50 +840,57 @@ mod tests { entries, count_value ); - let all_data = ctx - .session - .sql("SELECT a FROM '/table/'") - .await? - .collect() + let file_metas = ctx + .store + .list(Some(&"table".into())) + .try_collect::>() .await?; - let mut total_rows = 0; - for batch in all_data { - let col = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - - for i in 0..batch.num_rows() { - assert_eq!( - col.value(i), - 0i8, - "Expected value 0 at row {}, but found {}", - total_rows + i, - col.value(i) - ); - } - total_rows += batch.num_rows(); - } - - assert_eq!( - total_rows, entries, - "Total rows read ({}) doesn't match expected entries ({})", - total_rows, entries + assert!( + file_metas.len() >= expected_files, + "Expected at least {expected_files} files for {entries} values, got {}", + file_metas.len() ); + Ok(()) + } + + #[tokio::test] + async fn test_write_large_batch_default_target_is_128mb() -> anyhow::Result<()> { + let ctx = TestSessionContext::default(); + + let entries = 1_000_000; + let values: Vec = (0..entries) + .map(|i| i8::try_from(i % 127)) + .collect::>()?; + + let data = ctx.session.read_batch(RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int8, false)])), + vec![Arc::new(Int8Array::from(values))], + )?)?; + + let logical_plan = LogicalPlanBuilder::copy_to( + data.logical_plan().clone(), + "/table/".to_string(), + format_as_file_type(Arc::new(VortexFormatFactory::new())), + Default::default(), + vec![], + )? + .build()?; + + ctx.session + .execute_logical_plan(logical_plan) + .await? + .collect() + .await?; + let file_metas = ctx .store .list(Some(&"/table".into())) .try_collect::>() .await?; - assert_eq!( - file_metas.len(), - expected_files, - "Expected {expected_files} files for {entries} values" - ); + assert_eq!(file_metas.len(), 1); Ok(()) } @@ -448,4 +936,1030 @@ mod tests { Ok(()) } + + #[test] + fn test_split_path_basic() { + let path = object_store::path::Path::from("data/output"); + assert_eq!( + split_path(&path, 0, "vortex").to_string(), + "data/output/part-00000.vortex" + ); + assert_eq!( + split_path(&path, 12, "vortex").to_string(), + "data/output/part-00012.vortex" + ); + } + + #[test] + fn test_split_path_preserves_trailing_slash() { + let path = object_store::path::Path::from("nested/path/"); + assert_eq!( + split_path(&path, 3, "vx").to_string(), + "nested/path/part-00003.vx" + ); + } + + #[test] + fn test_numbered_path() { + use super::numbered_path; + + let path = object_store::path::Path::from("table/c1=alpha/abc123.vortex"); + assert_eq!( + numbered_path(&path, 0, "vortex").to_string(), + "table/c1=alpha/abc123_00000.vortex" + ); + assert_eq!( + numbered_path(&path, 5, "vortex").to_string(), + "table/c1=alpha/abc123_00005.vortex" + ); + } + + #[test] + fn test_numbered_path_no_extension() { + use super::numbered_path; + + let path = object_store::path::Path::from("table/output"); + assert_eq!( + numbered_path(&path, 0, "vortex").to_string(), + "table/output_00000.vortex" + ); + } + + #[test] + fn test_output_file_path_single_file_and_collection() { + use super::output_file_path; + + let single = ListingTableUrl::parse("file:///tmp/output.vortex").unwrap(); + assert_eq!( + output_file_path(&single, 0, "vortex", true, "wid").to_string(), + "tmp/output.vortex" + ); + assert_eq!( + output_file_path(&single, 2, "vortex", true, "wid").to_string(), + "tmp/output_00002.vortex" + ); + + let collection = ListingTableUrl::parse("file:///tmp/table/").unwrap(); + assert_eq!( + output_file_path(&collection, 3, "vortex", false, "wid").to_string(), + "tmp/table/wid_00003.vortex" + ); + } + + #[tokio::test] + async fn test_finish_file_writer_waits_for_task_completion() -> anyhow::Result<()> { + let (sender, receiver) = futures::channel::mpsc::channel::(1); + drop(receiver); + + let (gate_tx, gate_rx) = oneshot::channel::<()>(); + + let writer = ActiveFileWriter { + path: object_store::path::Path::from("table/pending.vortex"), + sender, + task: tokio::spawn(async move { + let _ = gate_rx.await; + Err(exec_datafusion_err!( + "synthetic writer failure after completion gate" + )) + }), + }; + + let mut finish_fut = Box::pin(finish_file_writer(writer)); + assert!( + tokio::time::timeout(Duration::from_millis(50), &mut finish_fut) + .await + .is_err(), + "finish_file_writer returned before writer task completed" + ); + + gate_tx + .send(()) + .map_err(|_| anyhow::anyhow!("failed to release writer completion gate in test"))?; + + let err = match finish_fut.await { + Ok(_) => { + return Err(anyhow::anyhow!( + "finish_file_writer unexpectedly succeeded after gate release" + )); + } + Err(err) => err, + }; + assert!( + err.to_string().contains("synthetic writer failure"), + "unexpected error: {err}" + ); + + Ok(()) + } + + #[test] + fn test_remove_partition_columns() -> anyhow::Result<()> { + use datafusion::arrow::array::StringArray; + + use super::remove_partition_columns; + + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("part", DataType::Utf8, false), + Field::new("val", DataType::Int64, false), + ])), + vec![ + Arc::new(StringArray::from(vec!["x", "y"])), + Arc::new(Int64Array::from(vec![1, 2])), + ], + )?; + + let out = remove_partition_columns(&batch, &["part".to_string()])?; + assert_eq!(out.num_columns(), 1); + assert_eq!(out.schema().field(0).name(), "val"); + assert_eq!(out.num_rows(), 2); + + Ok(()) + } + + #[test] + fn test_remove_partition_columns_all_columns_partitioned() -> anyhow::Result<()> { + use datafusion::arrow::array::StringArray; + + use super::remove_partition_columns; + + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("part", DataType::Utf8, false)])), + vec![Arc::new(StringArray::from(vec!["x", "y", "z"]))], + )?; + + let out = remove_partition_columns(&batch, &["part".to_string()])?; + assert_eq!(out.num_columns(), 0); + assert_eq!(out.num_rows(), 3); + + Ok(()) + } + + /// Generate `count` pseudo-random i64 values using a simple LCG. + /// These values resist compression (unlike sequential or modular data), + /// giving more realistic compressed file sizes. + fn pseudo_random_i64s(count: usize, seed: i64) -> Vec { + let mut v = seed; + (0..count) + .map(|_| { + v = v + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + v + }) + .collect() + } + + /// Tests file splitting through the full DataFusion pipeline. + /// + /// Writes ~62MB of pseudo-random Int64 data (near 1:1 compression ratio) + /// via COPY TO with a 16MB target file size. Verifies that exactly 4 files + /// are produced and each file's compressed size is approximately 16MB. + /// + /// This exercises the complete COPY TO write path through DataFusion and + /// VortexSink, unlike a direct `write_stream_to_files` call. + #[tokio::test] + async fn test_file_splitting_62mb_into_4_files() -> anyhow::Result<()> { + use datafusion::datasource::MemTable; + use datafusion_datasource::file_format::format_as_file_type; + + let ctx = TestSessionContext::default(); + + let target_mb = 16_usize; + let opts = VortexTableOptions { + target_file_size_mb: target_mb, + ..Default::default() + }; + let factory = VortexFormatFactory::new().with_options(opts); + + let batch_rows = 8192_usize; + let total_elements = 62 * 1024 * 1024 / 8; // ~8,126,464 i64 values ≈ 62MB Arrow memory + let num_batches = total_elements / batch_rows; + let expected_total_rows = (num_batches * batch_rows) as i64; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut batches = Vec::new(); + for i in 0..num_batches { + let values = pseudo_random_i64s(batch_rows, (i * batch_rows) as i64); + batches.push(RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + )?); + } + + let table = MemTable::try_new(Arc::clone(&schema), vec![batches])?; + ctx.session.register_table("source", Arc::new(table))?; + + let source = ctx.session.table("source").await?; + let logical_plan = LogicalPlanBuilder::copy_to( + source.logical_plan().clone(), + "/table/".to_string(), + format_as_file_type(Arc::new(factory)), + Default::default(), + vec![], + )? + .build()?; + + ctx.session + .execute_logical_plan(logical_plan) + .await? + .collect() + .await?; + + let file_metas = ctx + .store + .list(Some(&"/table".into())) + .try_collect::>() + .await?; + + assert_eq!( + file_metas.len(), + 4, + "Expected 4 files for ~62MB data with {target_mb}MB target, got {} (sizes: {:?})", + file_metas.len(), + file_metas.iter().map(|m| m.size).collect::>() + ); + + let target_bytes = u64::try_from(target_mb * 1024 * 1024)?; + for meta in &file_metas { + assert!( + meta.size > target_bytes / 2, + "File {} is {}B, expected at least {}B (target/2)", + meta.location, + meta.size, + target_bytes / 2 + ); + } + + // Verify total row count. + let result = ctx + .session + .sql("SELECT COUNT(*) as cnt FROM '/table/'") + .await? + .collect() + .await?; + + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + /// Tests file splitting with compressible data through the full pipeline. + /// + /// Uses low-entropy Int64 values (repeating 0..255) which compress ~8:1 in + /// Vortex. With the current code that compares Arrow memory size against + /// `target_file_size`, files are split far too early, producing many tiny + /// compressed files instead of files that are close to the target. + /// + /// For ~32MB of Arrow data (~4MB compressed at 8:1) with a 1MB target: + /// - **Correct**: 4 files of ~1MB compressed each + /// - **Bug**: 32 files of ~0.125MB compressed each + #[tokio::test] + async fn test_file_splitting_compressible_data() -> anyhow::Result<()> { + use datafusion::datasource::MemTable; + use datafusion_datasource::file_format::format_as_file_type; + + let ctx = TestSessionContext::default(); + + let target_mb = 1_usize; + let opts = VortexTableOptions { + target_file_size_mb: target_mb, + ..Default::default() + }; + let factory = VortexFormatFactory::new().with_options(opts); + + // Generate low-entropy Int64 values: repeating 0..255. + // Arrow memory: 4M × 8 bytes = 32MB. + // Vortex compressed: each value only needs ~1 byte → ~4MB total. + let total_elements = 4_000_000_usize; + let batch_rows = 8192_usize; + let num_batches = total_elements / batch_rows; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut batches = Vec::new(); + for i in 0..num_batches { + let values: Vec = (0..batch_rows) + .map(|j| ((i * batch_rows + j) % 256) as i64) + .collect(); + batches.push(RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + )?); + } + + let table = MemTable::try_new(Arc::clone(&schema), vec![batches])?; + ctx.session.register_table("source", Arc::new(table))?; + + let source = ctx.session.table("source").await?; + let logical_plan = LogicalPlanBuilder::copy_to( + source.logical_plan().clone(), + "/table/".to_string(), + format_as_file_type(Arc::new(factory)), + Default::default(), + vec![], + )? + .build()?; + + ctx.session + .execute_logical_plan(logical_plan) + .await? + .collect() + .await?; + + let file_metas = ctx + .store + .list(Some(&"/table".into())) + .try_collect::>() + .await?; + + // With compressible data, there should be few files (not > 10). + // The buggy code produces many tiny files because it splits on Arrow + // memory (32MB / 1MB = 32 files) instead of compressed size (~4MB / 1MB = 4 files). + let total_compressed: u64 = file_metas.iter().map(|m| m.size).sum(); + let target_bytes = u64::try_from(target_mb * 1024 * 1024)?; + + // We should have at most ~(total_compressed / target) + 1 files, not + // ~(arrow_memory / target) files. + let max_expected = usize::try_from(total_compressed / target_bytes + 2)?; + assert!( + file_metas.len() <= max_expected, + "Too many files: got {} but total compressed is {}B with {}B target \ + (expected at most {max_expected}). Files are being split on Arrow memory \ + instead of compressed size. Sizes: {:?}", + file_metas.len(), + total_compressed, + target_bytes, + file_metas.iter().map(|m| m.size).collect::>() + ); + + // Every file except the first should be reasonably sized. The first + // file may be smaller because the compression ratio is unknown until + // the first write completes. + for meta in file_metas.iter().skip(1) { + assert!( + meta.size > target_bytes / 4, + "File {} is {}B, far below target {}B — splitting on Arrow memory, not compressed size", + meta.location, + meta.size, + target_bytes + ); + } + + Ok(()) + } + + #[tokio::test] + async fn test_write_large_batch_target_file_size_disabled() -> anyhow::Result<()> { + use datafusion::datasource::MemTable; + use datafusion_datasource::file_format::format_as_file_type; + + let ctx = TestSessionContext::default(); + + let opts = VortexTableOptions { + // Disable sink-side rolling/splitting. + target_file_size_mb: 0, + ..Default::default() + }; + let factory = VortexFormatFactory::new().with_options(opts); + + let rows_per_partition = 300_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + partitions.push(vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + )?]); + } + + let table = MemTable::try_new(schema, partitions)?; + ctx.session.register_table("source", Arc::new(table))?; + + let source = ctx.session.table("source").await?; + let logical_plan = LogicalPlanBuilder::copy_to( + source.logical_plan().clone(), + "/table/".to_string(), + format_as_file_type(Arc::new(factory)), + Default::default(), + vec![], + )? + .build()?; + + ctx.session + .execute_logical_plan(logical_plan) + .await? + .collect() + .await?; + + let file_metas = ctx + .store + .list(Some(&"/table".into())) + .try_collect::>() + .await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = file_metas + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id with target size disabled; got {:?} from files: {:?}", + unique_write_ids, + file_metas + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + + assert_eq!( + file_metas.len(), + 1, + "Expected exactly one output file when target size is disabled, got {}", + file_metas.len() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) as cnt FROM '/table/'") + .await? + .collect() + .await?; + + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + #[tokio::test] + async fn test_target_file_size_uses_single_sink_input_partition() -> anyhow::Result<()> { + use datafusion::datasource::MemTable; + use datafusion_datasource::file_format::format_as_file_type; + + let ctx = TestSessionContext::default(); + + let opts = VortexTableOptions { + // Enable sink-side sizing, but make the threshold large enough + // that all input data should fit in a single file. + target_file_size_mb: 512, + ..Default::default() + }; + let factory = VortexFormatFactory::new().with_options(opts); + + let rows_per_partition = 300_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + // Build a MemTable with multiple physical input partitions to mimic + // DataFusion's parallel writer inputs. + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + partitions.push(vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + )?]); + } + + let table = MemTable::try_new(schema, partitions)?; + ctx.session.register_table("source", Arc::new(table))?; + + let source = ctx.session.table("source").await?; + let logical_plan = LogicalPlanBuilder::copy_to( + source.logical_plan().clone(), + "/table/".to_string(), + format_as_file_type(Arc::new(factory)), + Default::default(), + vec![], + )? + .build()?; + + ctx.session + .execute_logical_plan(logical_plan) + .await? + .collect() + .await?; + + let file_metas = ctx + .store + .list(Some(&"/table".into())) + .try_collect::>() + .await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = file_metas + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id (single sink stream), got {:?} from files: {:?}", + unique_write_ids, + file_metas + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + + assert!( + file_metas.len() < num_partitions, + "Expected fewer output files than input partitions after coalescing; got {} files for {num_partitions} input partitions", + file_metas.len() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) AS cnt FROM '/table/'") + .await? + .collect() + .await?; + + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + #[tokio::test] + async fn test_insert_sql_target_size_multi_partition_source_single_write_id() + -> anyhow::Result<()> { + use datafusion::datasource::MemTable; + + let ctx = TestSessionContext::default(); + + ctx.session + .sql( + "CREATE EXTERNAL TABLE my_tbl \ + (a BIGINT NOT NULL) \ + STORED AS vortex \ + LOCATION 'table/' \ + OPTIONS(target_file_size_mb '64');", + ) + .await?; + + let rows_per_partition = 300_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + partitions.push(vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + )?]); + } + + let source = MemTable::try_new(schema, partitions)?; + ctx.session.register_table("source", Arc::new(source))?; + + ctx.session + .sql("INSERT INTO my_tbl SELECT a FROM source") + .await? + .collect() + .await?; + + let all_files = ctx.store.list(None).try_collect::>().await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = all_files + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id, got {:?} from files: {:?}", + unique_write_ids, + all_files + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + assert!( + all_files.len() < num_partitions, + "Expected fewer files than input partitions; got {} files for {num_partitions} input partitions", + all_files.len() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) AS cnt FROM my_tbl") + .await? + .collect() + .await?; + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + #[tokio::test] + async fn test_insert_sql_streaming_source_single_write_id() -> anyhow::Result<()> { + use arrow_schema::SchemaRef; + use datafusion::catalog::streaming::StreamingTable; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use datafusion::physical_plan::streaming::PartitionStream; + use futures::stream; + + #[derive(Debug)] + struct StaticPartitionStream { + schema: SchemaRef, + batch: RecordBatch, + } + + impl PartitionStream for StaticPartitionStream { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute( + &self, + _ctx: Arc, + ) -> datafusion::physical_plan::SendableRecordBatchStream { + let schema = Arc::clone(&self.schema); + let batch = self.batch.clone(); + Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::iter(vec![Ok(batch)]), + )) + } + } + + let ctx = TestSessionContext::default(); + + ctx.session + .sql( + "CREATE EXTERNAL TABLE my_tbl \ + (a BIGINT NOT NULL) \ + STORED AS vortex \ + LOCATION 'table/' \ + OPTIONS(target_file_size_mb '64');", + ) + .await?; + + let rows_per_partition = 300_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + )?; + + partitions.push(Arc::new(StaticPartitionStream { + schema: Arc::clone(&schema), + batch, + })); + } + + let source = StreamingTable::try_new(schema, partitions)?; + ctx.session + .register_table("source_stream", Arc::new(source))?; + + ctx.session + .sql("INSERT INTO my_tbl SELECT a FROM source_stream") + .await? + .collect() + .await?; + + let all_files = ctx.store.list(None).try_collect::>().await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = all_files + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id for streaming source insert, got {:?} from files: {:?}", + unique_write_ids, + all_files + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) AS cnt FROM my_tbl") + .await? + .collect() + .await?; + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + #[tokio::test] + async fn test_listing_table_direct_insert_into_streaming_exec_single_write_id() + -> anyhow::Result<()> { + use arrow_schema::SchemaRef; + use datafusion::physical_plan::ExecutionPlan; + use datafusion::physical_plan::collect; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use datafusion::physical_plan::streaming::PartitionStream; + use datafusion::physical_plan::streaming::StreamingTableExec; + use datafusion_expr::dml::InsertOp; + use futures::stream; + + #[derive(Debug)] + struct StaticPartitionStream { + schema: SchemaRef, + batch: RecordBatch, + } + + impl PartitionStream for StaticPartitionStream { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute( + &self, + _ctx: Arc, + ) -> datafusion::physical_plan::SendableRecordBatchStream { + let schema = Arc::clone(&self.schema); + let batch = self.batch.clone(); + Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::iter(vec![Ok(batch)]), + )) + } + } + + let ctx = TestSessionContext::default(); + + ctx.session + .sql( + "CREATE EXTERNAL TABLE my_tbl \ + (a BIGINT NOT NULL) \ + STORED AS vortex \ + LOCATION 'table/' \ + OPTIONS(target_file_size_mb '64');", + ) + .await?; + + let table_provider = ctx.session.table_provider("my_tbl").await?; + + let rows_per_partition = 300_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + )?; + + partitions.push(Arc::new(StaticPartitionStream { + schema: Arc::clone(&schema), + batch, + })); + } + + let input = Arc::new(StreamingTableExec::try_new( + schema, + partitions, + None, + Vec::new(), + false, + None, + )?) as Arc; + + let plan = table_provider + .insert_into(&ctx.session.state(), input, InsertOp::Append) + .await?; + let _count_batches = collect(plan, ctx.session.task_ctx()).await?; + + let all_files = ctx.store.list(None).try_collect::>().await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = all_files + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id for direct insert_into streaming exec, got {:?} from files: {:?}", + unique_write_ids, + all_files + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) AS cnt FROM my_tbl") + .await? + .collect() + .await?; + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + #[tokio::test] + async fn test_listing_table_direct_insert_into_unbounded_streaming_exec_single_write_id() + -> anyhow::Result<()> { + use arrow_schema::SchemaRef; + use datafusion::physical_plan::ExecutionPlan; + use datafusion::physical_plan::collect; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use datafusion::physical_plan::streaming::PartitionStream; + use datafusion::physical_plan::streaming::StreamingTableExec; + use datafusion_expr::dml::InsertOp; + use futures::stream; + + #[derive(Debug)] + struct StaticPartitionStream { + schema: SchemaRef, + batch: RecordBatch, + } + + impl PartitionStream for StaticPartitionStream { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute( + &self, + _ctx: Arc, + ) -> datafusion::physical_plan::SendableRecordBatchStream { + let schema = Arc::clone(&self.schema); + let batch = self.batch.clone(); + Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::iter(vec![Ok(batch)]), + )) + } + } + + let ctx = TestSessionContext::default(); + + ctx.session + .sql( + "CREATE EXTERNAL TABLE my_tbl \ + (a BIGINT NOT NULL) \ + STORED AS vortex \ + LOCATION 'table/' \ + OPTIONS(target_file_size_mb '64');", + ) + .await?; + + let table_provider = ctx.session.table_provider("my_tbl").await?; + + let rows_per_partition = 100_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + )?; + + partitions.push(Arc::new(StaticPartitionStream { + schema: Arc::clone(&schema), + batch, + })); + } + + let input = Arc::new(StreamingTableExec::try_new( + schema, + partitions, + None, + Vec::new(), + true, + None, + )?) as Arc; + + let plan = table_provider + .insert_into(&ctx.session.state(), input, InsertOp::Append) + .await?; + let _count_batches = collect(plan, ctx.session.task_ctx()).await?; + + let all_files = ctx.store.list(None).try_collect::>().await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = all_files + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id for unbounded streaming insert, got {:?} from files: {:?}", + unique_write_ids, + all_files + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) AS cnt FROM my_tbl") + .await? + .collect() + .await?; + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } } diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs index 5c21da0ba57..e04d1d92025 100644 --- a/vortex-file/src/read/driver.rs +++ b/vortex-file/src/read/driver.rs @@ -8,6 +8,8 @@ use std::task::Context; use std::task::Poll; use futures::Stream; +#[cfg(all(test, not(feature = "tokio")))] +use oneshot; use pin_project_lite::pin_project; use tracing::trace; use vortex_buffer::Alignment; diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 661e72240d1..db22f97a8e1 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -3,6 +3,7 @@ //! This module defines the default layout strategy for a Vortex file. +use std::any::Any; use std::sync::Arc; use std::sync::LazyLock; @@ -56,6 +57,7 @@ use vortex_layout::layouts::zoned::writer::ZonedStrategy; use vortex_pco::Pco; use vortex_runend::RunEnd; use vortex_sequence::Sequence; +use vortex_session::SessionVar; use vortex_sparse::Sparse; use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; @@ -124,6 +126,7 @@ pub static ALLOWED_ENCODINGS: LazyLock> = LazyLock::new(|| { }); /// How the compressor was configured on [`WriteStrategyBuilder`]. +#[derive(Clone)] enum CompressorConfig { /// A [`BtrBlocksCompressorBuilder`] that [`WriteStrategyBuilder::build`] will finalize. /// `IntDictScheme` is automatically excluded from the data compressor to prevent recursive @@ -133,12 +136,22 @@ enum CompressorConfig { Opaque(Arc), } +impl std::fmt::Debug for CompressorConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::BtrBlocks(_) => f.write_str("BtrBlocks"), + Self::Opaque(_) => f.write_str("Opaque"), + } + } +} + /// Build a new [writer strategy](LayoutStrategy) to compress and reorganize chunks of a Vortex /// file. /// /// Vortex provides an out-of-the-box file writer that optimizes the layout of chunks on-disk, /// repartitioning and compressing them to strike a balance between size on-disk, /// bulk decoding performance, and IOPS required to perform an indexed read. +#[derive(Clone)] pub struct WriteStrategyBuilder { compressor: CompressorConfig, row_block_size: usize, @@ -147,6 +160,29 @@ pub struct WriteStrategyBuilder { flat_strategy: Option>, } +impl std::fmt::Debug for WriteStrategyBuilder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WriteStrategyBuilder") + .field("compressor", &self.compressor) + .field("row_block_size", &self.row_block_size) + .field( + "field_writers", + &format!("<{} entries>", self.field_writers.len()), + ) + .finish() + } +} + +impl SessionVar for WriteStrategyBuilder { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + impl Default for WriteStrategyBuilder { /// Create a new empty builder. It can be further configured, /// and then finally built yielding the [`LayoutStrategy`]. diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 25db11e4280..0a92a71fe58 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -72,10 +72,14 @@ pub struct VortexWriteOptions { pub trait WriteOptionsSessionExt: SessionExt { /// Create [`VortexWriteOptions`] for writing to a Vortex file. fn write_options(&self) -> VortexWriteOptions { - let session = self.session(); + let maybe_write_strategy_builder = self.get_opt::(); + let strategy = maybe_write_strategy_builder + .map(|opt| opt.clone().build()) + .unwrap_or_else(|| WriteStrategyBuilder::default().build()); + VortexWriteOptions { - strategy: WriteStrategyBuilder::default().build(), - session, + session: self.session(), + strategy, exclude_dtype: false, file_statistics: PRUNING_STATS.to_vec(), max_variable_length_statistics_size: 64, @@ -149,9 +153,15 @@ impl VortexWriteOptions { // serialised array order is deterministic. The serialisation of arrays are done // parallel and with an empty context they can register their encodings to the context // in different order, changing the written bytes from run to run. + // Avoid calling `self.session.arrays()` multiple times in one expression. + // `Ref` holds a session lock; reacquiring it in the same statement can deadlock. + let registry = { + let arrays = self.session.arrays(); + arrays.registry().clone() + }; let ctx = ArrayContext::new(ALLOWED_ENCODINGS.iter().cloned().sorted().collect()) // Configure a registry just to ensure only known encodings are interned. - .with_registry(self.session.arrays().registry().clone()); + .with_registry(registry); let dtype = stream.dtype().clone(); let (mut ptr, eof) = SequenceId::root().split(); @@ -481,3 +491,20 @@ impl WriteSummary { self.footer.row_count() } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_write_options_from_session_vars() { + let session = VortexSession::empty(); + let fetched_write_strategy = session.get_opt::(); + assert!(fetched_write_strategy.is_none()); + drop(fetched_write_strategy); + + let session = session.set(WriteStrategyBuilder::default()); + let fetched_write_strategy = session.get_opt::(); + assert!(fetched_write_strategy.is_some()); + } +} diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml index b3f6448484a..ee1a5e76ffb 100644 --- a/vortex-io/Cargo.toml +++ b/vortex-io/Cargo.toml @@ -24,6 +24,7 @@ bytes = { workspace = true } futures = { workspace = true, features = ["std", "executor"] } glob = { workspace = true } kanal = { workspace = true } +log = { workspace = true } object_store = { workspace = true, optional = true, features = ["fs"] } oneshot = { workspace = true } parking_lot = { workspace = true } diff --git a/vortex-io/src/runtime/single.rs b/vortex-io/src/runtime/single.rs index 2a113e3715e..743dc587f48 100644 --- a/vortex-io/src/runtime/single.rs +++ b/vortex-io/src/runtime/single.rs @@ -9,8 +9,13 @@ use futures::Stream; use futures::StreamExt; use futures::future::BoxFuture; use futures::stream::LocalBoxStream; +#[cfg(not(feature = "tokio"))] +use oneshot; use parking_lot::Mutex; use smol::LocalExecutor; +// Prefer tokio::sync::oneshot when tokio feature is enabled +#[cfg(feature = "tokio")] +use tokio::sync::oneshot; use vortex_error::vortex_panic; use crate::runtime::AbortHandle; @@ -19,7 +24,6 @@ use crate::runtime::BlockingRuntime; use crate::runtime::Executor; use crate::runtime::Handle; use crate::runtime::smol::SmolAbortHandle; - /// A runtime that drives all work on the current thread. /// /// This is subtly different from using a current-thread runtime to drive a future since it is diff --git a/vortex-session/src/lib.rs b/vortex-session/src/lib.rs index 805fb1fa364..005d3b90e16 100644 --- a/vortex-session/src/lib.rs +++ b/vortex-session/src/lib.rs @@ -64,6 +64,15 @@ impl VortexSession { self } + /// Inserts a new session variable of type `V` with the supplied value. + /// + /// # Panics + /// + /// If a variable of that type already exists. + pub fn set(self, val: V) -> Self { + self.with_some(val) + } + /// Allow deserializing unknown plugin IDs as non-executable foreign placeholders. pub fn allow_unknown(self) -> Self { let mut policy = ::get_mut::(&self); @@ -129,21 +138,16 @@ impl SessionExt for VortexSession { /// Returns the scope variable of type `V`, or inserts a default one if it does not exist. fn get(&self) -> Ref<'_, V> { - // NOTE(ngates): we don't use `entry().or_insert_with_key()` here because the DashMap - // would immediately acquire an exclusive write lock. - if let Some(v) = self.0.get(&TypeId::of::()) { - return Ref(v.map(|v| { - (**v) - .as_any() - .downcast_ref::() - .vortex_expect("Type mismatch - this is a bug") - })); + let type_id = TypeId::of::(); + if self.0.contains_key(&type_id) { + return self + .get_opt::() + .vortex_expect("Session variable disappeared between lookup and get"); } - // If we get here, the value was not present, so we insert the default with a write lock. Ref(self .0 - .entry(TypeId::of::()) + .entry(type_id) .or_insert_with(|| Box::new(V::default())) .downgrade() .map(|v| { diff --git a/vortex-web/package-lock.json b/vortex-web/package-lock.json index b54fce03a73..e97c9d1c96f 100644 --- a/vortex-web/package-lock.json +++ b/vortex-web/package-lock.json @@ -1849,9 +1849,6 @@ "arm64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1869,9 +1866,6 @@ "arm64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -1889,9 +1883,6 @@ "ppc64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1909,9 +1900,6 @@ "s390x" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1929,9 +1917,6 @@ "x64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1949,9 +1934,6 @@ "x64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -2417,9 +2399,6 @@ "arm64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -2437,9 +2416,6 @@ "arm64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -2457,9 +2433,6 @@ "x64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -2477,9 +2450,6 @@ "x64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -4718,9 +4688,6 @@ "arm64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -4742,9 +4709,6 @@ "arm64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -4766,9 +4730,6 @@ "x64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -4790,9 +4751,6 @@ "x64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MPL-2.0", "optional": true, "os": [