chore: import upstream snapshot with attribution

python/ray/data/__init__.py

@@ -0,0 +1,209 @@
+# Short term workaround for https://github.com/ray-project/ray/issues/32435
+# Dataset has a hard dependency on pandas, so it doesn't need to be delayed.
+import pandas  # noqa
+from packaging.version import parse as parse_version
+
+from ray.data._internal.utils.arrow_utils import get_pyarrow_version
+
+from ray.data._internal.compute import ActorPoolStrategy, TaskPoolStrategy
+from ray.data._internal.execution.interfaces import (
+    ExecutionOptions,
+    ExecutionResources,
+    NodeIdStr,
+)
+from ray.data._internal.logging import configure_logging
+from ray.data._internal.random_config import RandomSeedConfig
+from ray.data.context import DataContext, DatasetContext
+from ray.data.dataset import (
+    Dataset,
+    Schema,
+    SinkMode,
+    ClickHouseTableSettings,
+    SaveMode,
+)
+from ray.data._internal.logical.operators.n_ary_operator import (
+    MixStoppingCondition,
+)
+from ray.data.stats import DatasetSummary
+from ray.data.datasource import (
+    BlockBasedFileDatasink,
+    Datasink,
+    Datasource,
+    FileShuffleConfig,
+    ReadTask,
+    RowBasedFileDatasink,
+)
+from ray.data.iterator import DataIterator, DatasetIterator
+from ray.data.preprocessor import Preprocessor
+from ray.data.read_api import (  # noqa: F401
+    KafkaAuthConfig,  # noqa: F401
+    from_arrow,
+    from_arrow_refs,
+    from_blocks,
+    from_daft,
+    from_dask,
+    from_huggingface,
+    from_items,
+    from_mars,
+    from_modin,
+    from_numpy,
+    from_numpy_refs,
+    from_pandas,
+    from_pandas_refs,
+    from_spark,
+    from_tf,
+    from_torch,
+    range,
+    range_tensor,
+    read_audio,
+    read_avro,
+    read_bigquery,
+    read_binary_files,
+    read_clickhouse,
+    read_csv,
+    read_databricks_tables,
+    read_datasource,
+    read_delta,
+    read_delta_sharing_tables,
+    read_kafka,
+    read_hudi,
+    read_iceberg,
+    read_images,
+    read_json,
+    read_lance,
+    read_mcap,
+    read_mongo,
+    read_numpy,
+    read_parquet,
+    read_snowflake,
+    read_sql,
+    read_text,
+    read_tfrecords,
+    read_unity_catalog,
+    read_videos,
+    read_webdataset,
+    read_zarr,
+)
+from ray.data.catalog import (
+    Catalog,
+    ReaderFormat,
+    ResolvedSource,
+    DatabricksUnityCatalog,
+)
+
+# Module-level cached global functions for callable classes. It needs to be defined here
+# since it has to be process-global across cloudpickled funcs.
+_map_actor_context = None
+
+configure_logging()
+
+try:
+    import pyarrow as pa
+
+    # Import these arrow extension types to ensure that they are registered.
+    from ray.data._internal.tensor_extensions.arrow import (  # noqa
+        ArrowTensorType,
+        ArrowVariableShapedTensorType,
+    )
+
+    # https://github.com/apache/arrow/pull/38608 deprecated `PyExtensionType`, and
+    # disabled it's deserialization by default. To ensure that users can load data
+    # written with earlier version of Ray Data, we enable auto-loading of serialized
+    # tensor extensions.
+    #
+    # NOTE: `PyExtensionType` is deleted from Arrow >= 21.0
+    pyarrow_version = get_pyarrow_version()
+    if pyarrow_version is None or pyarrow_version >= parse_version("21.0.0"):
+        pass
+    else:
+        from ray._common.utils import env_bool
+
+        RAY_DATA_AUTOLOAD_PYEXTENSIONTYPE = env_bool(
+            "RAY_DATA_AUTOLOAD_PYEXTENSIONTYPE", False
+        )
+
+        if (
+            pyarrow_version >= parse_version("14.0.1")
+            and RAY_DATA_AUTOLOAD_PYEXTENSIONTYPE
+        ):
+            pa.PyExtensionType.set_auto_load(True)
+
+except ModuleNotFoundError:
+    pass
+
+
+__all__ = [
+    "ActorPoolStrategy",
+    "BlockBasedFileDatasink",
+    "ClickHouseTableSettings",
+    "Dataset",
+    "DataContext",
+    "DatasetContext",  # Backwards compatibility alias.
+    "DatasetSummary",
+    "DataIterator",
+    "DatasetIterator",  # Backwards compatibility alias.
+    "Datasink",
+    "Datasource",
+    "ExecutionOptions",
+    "ExecutionResources",
+    "FileShuffleConfig",
+    "MixStoppingCondition",
+    "NodeIdStr",
+    "RandomSeedConfig",
+    "ReadTask",
+    "RowBasedFileDatasink",
+    "Schema",
+    "SinkMode",
+    "SaveMode",
+    "TaskPoolStrategy",
+    "from_daft",
+    "from_dask",
+    "from_items",
+    "from_arrow",
+    "from_arrow_refs",
+    "from_blocks",
+    "from_mars",
+    "from_modin",
+    "from_numpy",
+    "from_numpy_refs",
+    "from_pandas",
+    "from_pandas_refs",
+    "from_spark",
+    "from_tf",
+    "from_torch",
+    "from_huggingface",
+    "range",
+    "range_tensor",
+    "read_audio",
+    "read_avro",
+    "read_text",
+    "read_binary_files",
+    "read_clickhouse",
+    "read_csv",
+    "read_datasource",
+    "read_delta",
+    "read_delta_sharing_tables",
+    "read_kafka",
+    "read_hudi",
+    "read_iceberg",
+    "read_images",
+    "read_json",
+    "read_lance",
+    "read_mcap",
+    "read_numpy",
+    "read_mongo",
+    "read_parquet",
+    "read_snowflake",
+    "read_sql",
+    "read_tfrecords",
+    "read_unity_catalog",
+    "read_videos",
+    "read_zarr",
+    "read_webdataset",
+    "Catalog",
+    "ReaderFormat",
+    "ResolvedSource",
+    "DatabricksUnityCatalog",
+    "KafkaAuthConfig",
+    "Preprocessor",
+]