content: """Basic IO operators for reading and writing data from Uptrain."""
  
  from __future__ import annotations
  import typing as t
  
  import polars as pl
  import pandas as pd
  
  if t.TYPE_CHECKING:
      from uptrain.framework import Settings
  from uptrain.operators.base import (
      TransformOp,
      register_op,
      TYPE_TABLE_OUTPUT,
      OpBaseModel,
  )
  from uptrain.utilities import lazy_load_dep
  
  # -----------------------------------------------------------
  # Read from text file formats - csv, json
  # -----------------------------------------------------------
  
  
  @register_op
  class CsvReader(TransformOp):
      """Reads data from a csv file.
  
      Attributes:
          fpath (str): Path to the csv file.
          batch_size (Optional[int]): Number of rows to read at a time. Defaults to None, which reads the entire file.
  
      """
  
      fpath: str
      batch_size: t.Optional[int] = None
  
      def setup(self, settings: Settings):
          self._executor = TextReaderExecutor(self)
          return self
  
      def run(self) -> TYPE_TABLE_OUTPUT:
          return {"output": self._executor.run()}
  
  
  @register_op
  class JsonReader(TransformOp):
      """Reads data from a json file.
  
      Attributes:
          fpath (str): Path to the json file.
          batch_size (Optional[int]): Number of rows to read at a time. Defaults to None, which reads the entire file.
  
      """
  
      fpath: str
      batch_size: t.Optional[int] = None
  
      def setup(self, settings: Settings):
          self._executor = TextReaderExecutor(self)
          return self
  
      def run(self) -> TYPE_TABLE_OUTPUT:
          return {"output": self._executor.run()}
  
  
  class TextReaderExecutor:
      op: t.Union[CsvReader, JsonReader]
      dataset: pl.DataFrame
      rows_read: int
  
      def __init__(self, op: t.Union[CsvReader, JsonReader]):
          self.op = op
          self.rows_read = 0
          if isinstance(self.op, CsvReader):
              self.dataset = pl.read_csv(self.op.fpath)
          elif isinstance(self.op, JsonReader):
              self.dataset = pl.read_ndjson(self.op.fpath)
              null_count = self.dataset.null_count().to_dicts()[0]
              for _, value in null_count.items():
                  if value == self.dataset.shape[0]:
                      ## read from pandas
                      pd_df = pd.read_json(self.op.fpath, lines=True)
                      self.dataset = pl.DataFrame(pd_df)
                      break
  
      @property
      def is_incremental(self) -> bool:
          return self.op.batch_size is not None
  
      def run(self) -> pl.DataFrame | None:
          if not self.is_incremental:
              data = self.dataset
          else:
              if self.rows_read >= len(self.dataset):
                  data = None
              else:
                  assert self.op.batch_size is not None
                  self.rows_read += self.op.batch_size
                  data = self.dataset.slice(self.rows_read, self.op.batch_size)
          return data
  
  
  # -----------------------------------------------------------
  # Read from a Delta Table, which uses parquet files under the hood
  # -----------------------------------------------------------
  
  
  @register_op
  class DeltaReader(TransformOp):
      """Reads data from a Delta Lake table.
  
      Attributes:
          fpath (str): File path to the Delta Lake table.
          batch_split (bool): Whether to read the table in batches. Defaults to False.
  
      """
  
      fpath: str
      batch_split: bool = False
      _dataset: t.Any  # pyarrow dataset
      _batch_generator: t.Optional[t.Iterator[t.Any]]  # record batch generator
  
      def setup(self, settings: Settings):
          lazy_load_dep("pyarrow", "pyarrow>=10.0.0")
          dl = lazy_load_dep("deltatable", "deltalake>=0.9")
  
          self._dataset = dl.DeltaTable(self.fpath).to_pyarrow_dataset()
          if self.is_incremental:
              self._batch_generator = iter(self._dataset.to_batches())
          return self
  
      @property
      def is_incremental(self) -> bool:
          return self.batch_split is True
  
      def run(self) -> TYPE_TABLE_OUTPUT:
          if not self.is_incremental:
              data = pl.from_arrow(self._dataset.to_table())
          else:
              try:
                  data = pl.from_arrow(next(self._batch_generator))  # type: ignore
              except StopIteration:
                  data = None
  
          if data is not None:
              assert isinstance(data, pl.DataFrame)
          return {"output": data}
  
  
  # -----------------------------------------------------------
  # Writer objects
  # -----------------------------------------------------------
  
  
  @register_op
  class DeltaWriter(OpBaseModel):
      fpath: str
      columns: t.Optional[list[str]] = None
  
      def setup(self, settings: Settings):
          return self
  
      def run(self, data: pl.DataFrame) -> TYPE_TABLE_OUTPUT:
          if self.columns is None:
              self.columns = list(data.columns)
          assert set(self.columns) == set(data.columns)
          data.write_delta(self.fpath, mode="append")
          return {"output": data}
  
      def to_reader(self):
          return DeltaReader(fpath=self.fpath)  # type: ignore
  
  
  @register_op
  class JsonWriter(OpBaseModel):
      fpath: str
      columns: t.Optional[list[str]] = None
  
      def setup(self, settings: Settings):
          return self
  
      def to_reader(self):
          return JsonReader(fpath=self.fpath)  # type: ignore
  
      def run(self, data: pl.DataFrame) -> TYPE_TABLE_OUTPUT:
          if self.columns is None:
              self.columns = list(data.columns)
          assert set(self.columns) == set(data.columns)
          with open(self.fpath, "a") as f:
              f.write(data.write_ndjson())
          return {"output": data}
  