diff --git a/Cargo.lock b/Cargo.lock index b864f54533..52689f19ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -610,19 +610,6 @@ dependencies = [ "pin-project-lite", ] -[[package]] -name = "async-compat" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1ba85bc55464dcbf728b56d97e119d673f4cf9062be330a9a26f3acf504a590" -dependencies = [ - "futures-core", - "futures-io", - "once_cell", - "pin-project-lite", - "tokio", -] - [[package]] name = "async-compression" version = "0.4.43" @@ -1277,9 +1264,9 @@ dependencies = [ [[package]] name = "azure_storage_blob" -version = "1.1.0-beta.1" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ea46a35353e11e78c4feace9656cb81e2532d3e814c6ff512e9fbb629b00878" +checksum = "17b10207ecf7d666df6940b50051f433b3cd5d2b9b1dd190613208d7a84e7eed" dependencies = [ "async-stream", "async-trait", @@ -1296,9 +1283,9 @@ dependencies = [ [[package]] name = "azure_storage_common" -version = "0.1.0" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f8d86465df4adce8ab1246205be0b99165683808701cd6f0f84e382ad7cd504" +checksum = "0af2e6aeb8d76b17fc998f453c320913f73787b944e3cc29509d19411fa0321d" dependencies = [ "azure_core", "serde", @@ -2648,7 +2635,6 @@ name = "daft-csv" version = "0.3.0-dev0" dependencies = [ "arrow-schema 59.0.0", - "async-compat", "async-stream", "common-error", "common-py-serde", @@ -3372,7 +3358,6 @@ dependencies = [ "indexmap", "log", "parquet", - "path_macro", "pyo3", "serde", "snafu", @@ -6593,12 +6578,6 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" -[[package]] -name = "path_macro" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6e819bbd49d5939f682638fa54826bf1650abddcd65d000923de8ad63cc7d15" - [[package]] name = "pem" version = "3.0.6" @@ -7312,9 +7291,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.12.3" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" dependencies = [ "aho-corasick", "memchr", @@ -7324,9 +7303,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.14" +version = "0.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" dependencies = [ "aho-corasick", "memchr", @@ -7341,9 +7320,9 @@ checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" [[package]] name = "regex-syntax" -version = "0.8.10" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "relative-path" @@ -8159,9 +8138,7 @@ version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" dependencies = [ - "bincode", "serde", - "unty", ] [[package]] @@ -8775,6 +8752,7 @@ checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" dependencies = [ "bytes", "futures-core", + "futures-io", "futures-sink", "pin-project-lite", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 84f8daa0de..6ec07db74c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -263,7 +263,6 @@ arrow-schema = "59.0.0" arrow-select = "59.0.0" approx = "0.5.1" arc-swap = "1.7.1" -async-compat = "0.2.5" async-compression = {version = "0.4.43", features = [ "tokio", "brotli", @@ -351,21 +350,20 @@ opentelemetry-otlp = {version = "0.33", features = ["grpc-tonic", "logs"]} opentelemetry_sdk = {version = "0.33", features = ["logs"]} parking_lot = "0.12.5" parquet = "59.0.0" -path_macro = "1.0.0" pretty_assertions = "1.4.1" proptest = "1.9.0" prost = "0.13.5" prost-types = "0.13.5" rand = "^0.9" rayon = "1.12.0" -regex = "1.12.2" +regex = "1.13" rstest = "0.26.1" reqwest = {version = "0.13", default-features = false, features = ["json"]} serde_json = "1.0.145" sha1 = "^0.11" simdutf8 = "0.1.5" sketches-ddsketch = {version = "0.2.2", features = ["use_serde"]} -smallvec = {version = "1.15.1", features = ["serde", "impl_bincode"]} +smallvec = {version = "1.15.1", features = ["serde"]} snafu = {version = "0.9.0", features = ["futures"]} sqlparser = "0.59.0" sysinfo = "0.39" @@ -383,7 +381,7 @@ tokio = {version = "1.48.0", features = [ "rt", "rt-multi-thread" ]} -tokio-stream = {version = "0.1.14", features = ["fs", "io-util", "time"]} +tokio-stream = {version = "0.1.14", default-features = false, features = ["fs", "io-util"]} tokio-util = "0.7.17" tonic = "0.14.2" tonic-build = "0.14.2" @@ -393,7 +391,7 @@ typetag = "0.2.21" rustfft = "6" url = "2.4.0" uuid = {version = "1.19.0", features = ["v4", "v7", "serde", "fast-rng"]} -xxhash-rust = {version = "0.8.15", features = ["const_xxh3", "const_xxh64", "const_xxh32", "xxh64", "xxh3", "xxh32"]} +xxhash-rust = {version = "0.8.15", features = ["const_xxh3", "xxh64", "xxh3", "xxh32"]} libc = {version = "^0.2.150", default-features = false} hex = "0.4.3" @@ -402,7 +400,7 @@ version = "2.0.1" features = ["derive", "serde"] [workspace.dependencies.derive_more] -features = ["display", "from"] +features = ["display"] version = "2.1.0" [workspace.dependencies.log] diff --git a/daft/file/file.py b/daft/file/file.py index 1bcf017661..7cd6fcc50e 100644 --- a/daft/file/file.py +++ b/daft/file/file.py @@ -143,7 +143,7 @@ def position(self) -> int | None: @property def offset(self) -> int | None: - """Deprecated alias for `position`. The byte offset for range reads, or None for full-file reads.""" + """Deprecated alias for `position`. The byte offset for range reads, or None for full-file reads. Will be removed in v0.9.0.""" warnings.warn( "`File.offset` is deprecated; use `File.position` instead.", DeprecationWarning, @@ -153,7 +153,7 @@ def offset(self) -> int | None: @property def length(self) -> int | None: - """Deprecated alias for the byte-range read window size, or None for full-file reads. + """Deprecated alias for the byte-range read window size, or None for full-file reads Will be removed in v0.9.0. Note: this returns the requested range size (caller intent), not the derived file size. Use `File.size()` for the actual file size. diff --git a/daft/io/source.py b/daft/io/source.py index 6b5ae7fce5..b995fa330d 100644 --- a/daft/io/source.py +++ b/daft/io/source.py @@ -1,20 +1,18 @@ from __future__ import annotations -import warnings from abc import ABC, abstractmethod from typing import TYPE_CHECKING from daft.daft import McapSourceConfig, ParquetSourceConfig, PyDataSourceTask, StorageConfig if TYPE_CHECKING: - from collections.abc import AsyncIterator, Iterator + from collections.abc import AsyncIterator from daft.dataframe import DataFrame from daft.io.clustering import ClusteringKeys from daft.io.partitioning import PartitionField from daft.io.pushdowns import Pushdowns from daft.recordbatch import RecordBatch - from daft.recordbatch.micropartition import MicroPartition from daft.schema import Schema @@ -124,31 +122,16 @@ def schema(self) -> Schema: """Returns the schema of the record batches produced by this task.""" ... + @abstractmethod async def read(self) -> AsyncIterator[RecordBatch]: """Yields record batches. Called from an async execution context. - The default implementation delegates to the deprecated - get_micro_partitions for backwards compatibility. - New subclasses should override this method directly. + Subclasses must override this method. Each yielded batch must match :attr:`schema`. """ - try: - parts = self.get_micro_partitions() - except NotImplementedError: - raise NotImplementedError( - f"{type(self).__name__} must implement async def read(self) -> AsyncIterator[RecordBatch]" - ) - warnings.warn( - f"{type(self).__name__}.get_micro_partitions() is deprecated — override 'async def read()' instead.", - DeprecationWarning, - stacklevel=2, - ) - for mp in parts: - for rb in mp.get_record_batches(): - yield rb - - def get_micro_partitions(self) -> Iterator[MicroPartition]: - """Deprecated: override read instead.""" - raise NotImplementedError + ... + # https://mypy.readthedocs.io/en/latest/more_types.html#typing-async-generators + # This tricks the type checker into treating this as an async generator. + yield # type: ignore[misc] @staticmethod def parquet( diff --git a/src/daft-csv/Cargo.toml b/src/daft-csv/Cargo.toml index 11793e3b26..86abdfb4a2 100644 --- a/src/daft-csv/Cargo.toml +++ b/src/daft-csv/Cargo.toml @@ -1,7 +1,6 @@ [dependencies] log = {workspace = true} arrow-schema = {workspace = true} -async-compat = {workspace = true} async-stream = {workspace = true} common-error = {path = "../common/error", default-features = false} common-py-serde = {path = "../common/py-serde", default-features = false} @@ -24,7 +23,7 @@ serde = {workspace = true} smallvec = {workspace = true} snafu = {workspace = true} tokio = {workspace = true} -tokio-util = {workspace = true} +tokio-util = {workspace = true, features = ["compat"]} [dev-dependencies] rstest = {workspace = true} diff --git a/src/daft-csv/src/metadata.rs b/src/daft-csv/src/metadata.rs index a80320a821..a1ec65292b 100644 --- a/src/daft-csv/src/metadata.rs +++ b/src/daft-csv/src/metadata.rs @@ -1,7 +1,6 @@ use std::{collections::HashSet, sync::Arc}; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; -use async_compat::CompatExt; use common_error::DaftResult; use csv_async::{AsyncReaderBuilder, ByteRecord}; use daft_compression::CompressionCodec; @@ -14,7 +13,7 @@ use tokio::{ fs::File, io::{AsyncBufRead, AsyncRead, BufReader}, }; -use tokio_util::io::StreamReader; +use tokio_util::{compat::TokioAsyncReadCompatExt, io::StreamReader}; use crate::{CsvParseOptions, schema::merge_schema}; diff --git a/src/daft-csv/src/read.rs b/src/daft-csv/src/read.rs index 3d6a48134a..b05e0cbf0e 100644 --- a/src/daft-csv/src/read.rs +++ b/src/daft-csv/src/read.rs @@ -2,7 +2,6 @@ use core::str; use std::{collections::HashMap, num::NonZeroUsize, sync::Arc}; use arrow_schema::Field as ArrowField; -use async_compat::{Compat, CompatExt}; use common_error::{DaftError, DaftResult}; use common_runtime::get_io_runtime; use csv_async::{AsyncReader, AsyncReaderBuilder}; @@ -26,7 +25,10 @@ use tokio::{ io::{AsyncBufRead, AsyncRead, BufReader}, task::JoinHandle, }; -use tokio_util::io::StreamReader; +use tokio_util::{ + compat::{Compat, TokioAsyncReadCompatExt}, + io::StreamReader, +}; use crate::{CsvConvertOptions, CsvParseOptions, CsvReadOptions, metadata::read_csv_schema_single}; diff --git a/src/daft-io/Cargo.toml b/src/daft-io/Cargo.toml index 4d526cbc83..2d37e5756e 100644 --- a/src/daft-io/Cargo.toml +++ b/src/daft-io/Cargo.toml @@ -9,7 +9,7 @@ aws-smithy-http-client = "1.1.2" aws-smithy-runtime-api = "1.8.0" azure_core = {version = "1.1.0", features = ["reqwest", "reqwest_rustls", "tokio", "hmac_rust", "xml"], default-features = false} azure_identity = {version = "1.0.0", features = ["tokio"], default-features = false} -azure_storage_blob = {version = "1.1.0-beta.1", default-features = false} +azure_storage_blob = {version = "1.1.0", default-features = false} bytes = {workspace = true} common-error = {path = "../common/error", default-features = false} common-file-formats = {path = "../common/file-formats", default-features = false} @@ -38,7 +38,7 @@ opendal = {version = "0.58", default-features = false, features = [ ]} pyo3 = {workspace = true, optional = true} rand = {workspace = true} -regex = "1.12.2" +regex = {workspace = true} reqwest = {workspace = true, default-features = false, features = ["stream", "json"]} reqwest-middleware = "0.5" reqwest-retry = "0.9" diff --git a/src/daft-io/src/azure_auth.rs b/src/daft-io/src/azure_auth.rs index f6386250df..0e4d4b8dd5 100644 --- a/src/daft-io/src/azure_auth.rs +++ b/src/daft-io/src/azure_auth.rs @@ -11,6 +11,10 @@ //! `azure_storage::authorization_policy`. //! - [`SasTokenPolicy`]: appends a shared access signature's query parameters //! to each request. +//! - [`ListBlobsDelimiterPolicy`]: adds the List Blobs `delimiter` query +//! parameter so the service returns virtual directories. The parameter has +//! to be on the URL before [`SharedKeyAuthorizationPolicy`] signs the +//! canonicalized resource. use std::sync::Arc; @@ -127,6 +131,57 @@ impl Policy for SharedKeyAuthorizationPolicy { } } +/// Value stored on the request [`Context`] so [`ListBlobsDelimiterPolicy`] adds +/// a List Blobs `delimiter`. +/// +/// Only List Blobs calls (`comp=list`) that carry this value are modified. +#[derive(Debug, Clone)] +pub(crate) struct ListBlobsDelimiter(pub String); + +/// A per-call pipeline [`Policy`] that appends `delimiter` to List Blobs URLs. +/// +/// Register this on the client's per-call policies so it runs before per-try +/// auth. SharedKey signatures include every query parameter, and a delimiter +/// added after signing is rejected. +#[derive(Debug, Clone)] +pub(crate) struct ListBlobsDelimiterPolicy; + +impl ListBlobsDelimiterPolicy { + pub(crate) fn new() -> Arc { + Arc::new(Self) + } +} + +#[async_trait] +impl Policy for ListBlobsDelimiterPolicy { + async fn send( + &self, + ctx: &Context, + request: &mut Request, + next: &[Arc], + ) -> PolicyResult { + if let Some(ListBlobsDelimiter(delimiter)) = ctx.value::() { + append_list_blobs_delimiter(request.url_mut(), delimiter); + } + next[0].send(ctx, request, &next[1..]).await + } +} + +/// Add `delimiter` to a List Blobs URL. Other requests, and URLs that already +/// have a delimiter, are left alone so retries do not duplicate the parameter. +fn append_list_blobs_delimiter(url: &mut Url, delimiter: &str) { + let is_list = url + .query_pairs() + .any(|(key, value)| key == "comp" && value == "list"); + if !is_list { + return; + } + if url.query_pairs().any(|(key, _)| key == "delimiter") { + return; + } + url.query_pairs_mut().append_pair("delimiter", delimiter); +} + /// A per-try pipeline [`Policy`] that appends a shared access signature's /// query parameters to each request URL. #[derive(Debug, Clone)] @@ -346,6 +401,61 @@ mod tests { assert_eq!(token.token.secret(), "my-fabric-token"); } + #[test] + fn list_blobs_delimiter_is_added_only_to_list_requests() { + let mut list_url = Url::parse( + "https://myaccount.blob.core.windows.net/mycontainer?comp=list&restype=container&prefix=a%2Fb", + ) + .unwrap(); + append_list_blobs_delimiter(&mut list_url, "/"); + let pairs: Vec<_> = list_url + .query_pairs() + .map(|(key, value)| (key.into_owned(), value.into_owned())) + .collect(); + assert!(pairs.contains(&("comp".to_string(), "list".to_string()))); + assert!(pairs.contains(&("restype".to_string(), "container".to_string()))); + assert!(pairs.contains(&("prefix".to_string(), "a/b".to_string()))); + assert!(pairs.contains(&("delimiter".to_string(), "/".to_string()))); + + // Retries must not append a second delimiter. + append_list_blobs_delimiter(&mut list_url, "/"); + assert_eq!( + list_url + .query_pairs() + .filter(|(key, _)| key == "delimiter") + .count(), + 1 + ); + + let original = "https://myaccount.blob.core.windows.net/mycontainer/blob?comp=metadata"; + let mut other = Url::parse(original).unwrap(); + append_list_blobs_delimiter(&mut other, "/"); + assert_eq!(other.as_str(), original); + } + + #[test] + fn string_to_sign_includes_list_blobs_delimiter() { + // SharedKey canonicalizes every query parameter. The delimiter policy + // has to run first or Azure rejects the signature. + let mut headers = Headers::new(); + headers.insert(MS_DATE, "Wed, 05 Aug 2026 20:00:00 GMT"); + headers.insert( + HeaderName::from_static("x-ms-version"), + "2025-01-05".to_string(), + ); + let mut url = Url::parse( + "https://myaccount.blob.core.windows.net/mycontainer?comp=list&restype=container", + ) + .unwrap(); + append_list_blobs_delimiter(&mut url, "/"); + + let str_to_sign = string_to_sign(&headers, &url, &Method::Get, "myaccount"); + assert!( + str_to_sign.contains("\ndelimiter:/\n"), + "string-to-sign missing delimiter:\n{str_to_sign}" + ); + } + #[test] fn sas_token_policy_parses_with_and_without_question_mark() { for sas in [ diff --git a/src/daft-io/src/azure_blob.rs b/src/daft-io/src/azure_blob.rs index 111e2baae1..42282df15e 100644 --- a/src/daft-io/src/azure_blob.rs +++ b/src/daft-io/src/azure_blob.rs @@ -3,7 +3,7 @@ use std::{any::Any, sync::Arc}; use async_trait::async_trait; use azure_core::{ credentials::{Secret, TokenCredential}, - http::{ClientOptions, StatusCode, Url, policies::Policy}, + http::{ClientOptions, Context, StatusCode, Url, pager::PagerOptions, policies::Policy}, }; use azure_identity::{ ClientSecretCredential, DeveloperToolsCredential, ManagedIdentityCredential, @@ -14,8 +14,7 @@ use azure_storage_blob::{ BlobServiceClient, BlobServiceClientOptions, models::{ BlobClientDownloadOptions, BlobClientGetPropertiesResultHeaders, - BlobContainerClientListBlobsHierarchicalOptions, BlobContainerClientListBlobsOptions, - BlobItem, ContainerItem, HttpRange, + BlobContainerClientListBlobsOptions, BlobItem, ContainerItem, HttpRange, }, }; use common_io_config::AzureConfig; @@ -25,7 +24,10 @@ use snafu::{IntoError, ResultExt, Snafu}; use crate::{ FileFormat, GetResult, InvalidRangeRequestSnafu, - azure_auth::{SasTokenPolicy, SharedKeyAuthorizationPolicy, StaticBearerCredential}, + azure_auth::{ + ListBlobsDelimiter, ListBlobsDelimiterPolicy, SasTokenPolicy, SharedKeyAuthorizationPolicy, + StaticBearerCredential, + }, object_io::{FileMetadata, FileType, LSResult, ObjectSource}, range::GetRange, stats::IOStatsRef, @@ -286,12 +288,12 @@ struct ParsedBlobList { prefixes: Vec, } -/// Parse a List Blobs (hierarchical / delimiter) XML body. +/// Parse a delimiter List Blobs XML body. /// -/// `azure_storage_blob` 1.1.0-beta.1 deserializes `blob_prefixes` as -/// `Option>`, which treats a second `` sibling as a -/// duplicate field. HNS accounts also attach `` to prefixes. Parse -/// the XML ourselves so posix listing keeps working. +/// `ListBlobsResponse` does not include `BlobPrefix`, and a second +/// `` sibling is treated as a duplicate field. HNS accounts also +/// attach `` to prefixes. Read the XML directly so posix listing +/// still returns virtual directories. fn parse_hierarchical_listing_xml(xml: &[u8]) -> Result { let text = std::str::from_utf8(xml).map_err(|e| format!("invalid utf-8: {e}"))?; let text = text.trim_start_matches('\u{feff}'); @@ -476,6 +478,8 @@ impl AzureBlobSource { let endpoint = ensure_trailing_slash(endpoint); let client_options = ClientOptions { + // Per-call so `delimiter` is on the URL before per-try SharedKey signing. + per_call_policies: vec![ListBlobsDelimiterPolicy::new() as _], per_try_policies, ..Default::default() }; @@ -725,23 +729,28 @@ impl AzureBlobSource { }; let request_path = format!("{}://{}{}", protocol, container_name, prefix); if posix { - // Setting a delimiter triggers "directory-mode" which is a posix-like ls for the current directory. - let options = BlobContainerClientListBlobsHierarchicalOptions { + // `delimiter` is directory-mode: one level of blobs and virtual directories. + // The per-call policy adds the query parameter when this context value is set. + let options = BlobContainerClientListBlobsOptions { prefix: request_prefix.clone(), + method_options: PagerOptions { + context: Context::new() + .with_value(ListBlobsDelimiter(AZURE_DELIMITER.to_string())), + ..Default::default() + }, ..Default::default() }; - let mut pager = - match container_client.list_blobs_hierarchical(AZURE_DELIMITER, Some(options)) { - Ok(pager) => pager, - Err(error) => { - yield Err(Error::RequestFailedForPath { - path: request_path, - source: error, - } - .into()); - return; + let mut pager = match container_client.list_blobs(Some(options)) { + Ok(pager) => pager.into_pages(), + Err(error) => { + yield Err(Error::RequestFailedForPath { + path: request_path, + source: error, } - }; + .into()); + return; + } + }; while let Some(page) = pager.next().await { if let Some(is) = io_stats.clone() { is.mark_list_requests(1); @@ -1100,6 +1109,20 @@ mod tests { let listing = parse_hierarchical_listing_xml(xml.as_bytes()).unwrap(); assert_eq!(listing.prefixes, vec!["iceberg/", "test_ls/"]); assert_eq!(listing.blobs, vec![("mvp.parquet".to_string(), Some(1234))]); + + // `list_blobs` pages by deserializing only `NextMarker`. Delimiter responses + // still have to parse, including an empty marker on the last page. + #[derive(serde::Deserialize)] + struct ListPage { + #[serde(rename = "NextMarker")] + next_marker: Option, + } + let page: ListPage = azure_core::xml::from_xml(xml.as_bytes()).unwrap(); + assert_eq!(page.next_marker.as_deref(), Some("")); + + let continued = xml.replace("", "page-2"); + let page: ListPage = azure_core::xml::from_xml(continued.as_bytes()).unwrap(); + assert_eq!(page.next_marker.as_deref(), Some("page-2")); } #[test] diff --git a/src/daft-logical-plan/src/ops/key_filtering.rs b/src/daft-logical-plan/src/ops/key_filtering.rs index eca651afaf..ff3f84e351 100644 --- a/src/daft-logical-plan/src/ops/key_filtering.rs +++ b/src/daft-logical-plan/src/ops/key_filtering.rs @@ -1,13 +1,11 @@ use common_error::{DaftError, DaftResult}; use common_hashable_float_wrapper::FloatWrapper; -use educe::Educe; #[cfg(feature = "python")] use pyo3::prelude::*; use serde::{Deserialize, Serialize}; -#[derive(Educe, Clone, Serialize, Deserialize)] +#[derive(Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] #[cfg_attr(debug_assertions, derive(Debug))] -#[educe(PartialEq, Eq, Hash)] pub struct KeyFilteringConfig { pub left_key_columns: Vec, pub right_key_columns: Vec, diff --git a/src/daft-logical-plan/src/ops/stage_checkpoint_keys.rs b/src/daft-logical-plan/src/ops/stage_checkpoint_keys.rs index fcbaebcd0a..1698c14d1b 100644 --- a/src/daft-logical-plan/src/ops/stage_checkpoint_keys.rs +++ b/src/daft-logical-plan/src/ops/stage_checkpoint_keys.rs @@ -1,7 +1,6 @@ use std::sync::Arc; use common_checkpoint_config::CheckpointConfig; -use educe::Educe; use serde::{Deserialize, Serialize}; use crate::{ @@ -20,9 +19,8 @@ use crate::{ /// must not move user operators below it, otherwise post-filter keys would /// be staged instead of post-anti-join keys, defeating source-oriented /// tracking. -#[derive(Educe, Clone, Serialize, Deserialize)] +#[derive(Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] #[cfg_attr(debug_assertions, derive(Debug))] -#[educe(PartialEq, Eq, Hash)] pub struct StageCheckpointKeys { pub plan_id: Option, pub node_id: Option, diff --git a/src/daft-parquet/Cargo.toml b/src/daft-parquet/Cargo.toml index 4374ee9aeb..d10c14568d 100644 --- a/src/daft-parquet/Cargo.toml +++ b/src/daft-parquet/Cargo.toml @@ -27,7 +27,6 @@ tokio-stream = {workspace = true} bincode = {workspace = true} fastrand = "2.1.0" parquet = {workspace = true} -path_macro = {workspace = true} tango-bench.workspace = true [features] diff --git a/src/daft-parquet/src/read.rs b/src/daft-parquet/src/read.rs index 94bbc630cf..db5a19f229 100644 --- a/src/daft-parquet/src/read.rs +++ b/src/daft-parquet/src/read.rs @@ -712,18 +712,9 @@ mod tests { #[test] fn test_invalid_utf8_parquet_reading() { - let parquet: Arc = path_macro::path!( - env!("CARGO_MANIFEST_DIR") - / ".." - / ".." - / "tests" - / "assets" - / "parquet-data" - / "invalid_utf8.parquet" - ) - .to_str() - .unwrap() - .into(); + let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + path.push("../../tests/assets/parquet-data/invalid_utf8.parquet"); + let parquet: Arc = path.to_str().unwrap().into(); let io_client = Arc::new(IOClient::new(IOConfig::default().into()).unwrap()); let runtime = get_io_runtime(true); let file_metadata = runtime diff --git a/tests/io/test_datasource_join.py b/tests/io/test_datasource_join.py index b9f0de375e..1e1ced7396 100644 --- a/tests/io/test_datasource_join.py +++ b/tests/io/test_datasource_join.py @@ -10,14 +10,14 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator import pyarrow as pa from daft import col from daft.io.pushdowns import Pushdowns from daft.io.source import DataSource, DataSourceTask -from daft.recordbatch import MicroPartition +from daft.recordbatch import RecordBatch from daft.schema import Schema # --------------------------------------------------------------------------- @@ -37,11 +37,11 @@ def schema(self) -> Schema: return Schema.from_pyarrow_schema(projected) return Schema.from_pyarrow_schema(self._table.schema) - def get_micro_partitions(self) -> Iterator[MicroPartition]: + async def read(self) -> AsyncIterator[RecordBatch]: t = self._table if self._columns: t = t.select(self._columns) - yield MicroPartition.from_arrow(t) + yield RecordBatch.from_arrow_table(t) class InMemorySource(DataSource):